opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: workspace quotas (#4184)

Browse Source
This commit is contained in:
Garrett Delfosse
2022-09-30 14:01:20 -04:00
committed by GitHub
parent f9b7588963
commit 69c73b2d28
28 changed files with 712 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/coderd/coderd.go
+46 -17
View File
@@ -15,7 +15,6 @@ import (
"cdr.dev/slog"
"github.com/coder/coder/coderd"
agplaudit "github.com/coder/coder/coderd/audit"
"github.com/coder/coder/coderd/httpapi"
"github.com/coder/coder/coderd/httpmw"
"github.com/coder/coder/codersdk"
@@ -43,9 +42,10 @@ func New(ctx context.Context, options *Options) (*API, error) {
Entitlement: codersdk.EntitlementNotEntitled,
Enabled: false,
},
auditLogs: codersdk.EntitlementNotEntitled,
browserOnly: codersdk.EntitlementNotEntitled,
scim: codersdk.EntitlementNotEntitled,
auditLogs: codersdk.EntitlementNotEntitled,
browserOnly: codersdk.EntitlementNotEntitled,
scim: codersdk.EntitlementNotEntitled,
workspaceQuota: codersdk.EntitlementNotEntitled,
},
cancelEntitlementsLoop: cancelFunc,
}
@@ -67,6 +67,13 @@ func New(ctx context.Context, options *Options) (*API, error) {
r.Get("/", api.licenses)
r.Delete("/{id}", api.deleteLicense)
})
r.Route("/workspace-quota", func(r chi.Router) {
r.Use(apiKeyMiddleware)
r.Route("/{user}", func(r chi.Router) {
r.Use(httpmw.ExtractUserParam(options.Database))
r.Get("/", api.workspaceQuota)
})
})
})
if len(options.SCIMAPIKey) != 0 {
@@ -96,8 +103,10 @@ type Options struct {
AuditLogging bool
// Whether to block non-browser connections.
BrowserOnly bool
SCIMAPIKey []byte
BrowserOnly bool
SCIMAPIKey []byte
UserWorkspaceQuota int
EntitlementsUpdateInterval time.Duration
Keys map[string]ed25519.PublicKey
}
@@ -112,11 +121,12 @@ type API struct {
}
type entitlements struct {
hasLicense bool
activeUsers codersdk.Feature
auditLogs codersdk.Entitlement
browserOnly codersdk.Entitlement
scim codersdk.Entitlement
hasLicense bool
activeUsers codersdk.Feature
auditLogs codersdk.Entitlement
browserOnly codersdk.Entitlement
scim codersdk.Entitlement
workspaceQuota codersdk.Entitlement
}
func (api *API) Close() error {
@@ -140,9 +150,10 @@ func (api *API) updateEntitlements(ctx context.Context) error {
Enabled: false,
Entitlement: codersdk.EntitlementNotEntitled,
},
auditLogs: codersdk.EntitlementNotEntitled,
scim: codersdk.EntitlementNotEntitled,
browserOnly: codersdk.EntitlementNotEntitled,
auditLogs: codersdk.EntitlementNotEntitled,
scim: codersdk.EntitlementNotEntitled,
browserOnly: codersdk.EntitlementNotEntitled,
workspaceQuota: codersdk.EntitlementNotEntitled,
}
// Here we loop through licenses to detect enabled features.
@@ -181,20 +192,22 @@ func (api *API) updateEntitlements(ctx context.Context) error {
if claims.Features.SCIM > 0 {
entitlements.scim = entitlement
}
if claims.Features.WorkspaceQuota > 0 {
entitlements.workspaceQuota = entitlement
}
}
if entitlements.auditLogs != api.entitlements.auditLogs {
auditor := agplaudit.NewNop()
// A flag could be added to the options that would allow disabling
// enhanced audit logging here!
if entitlements.auditLogs != codersdk.EntitlementNotEntitled && api.AuditLogging {
auditor = audit.NewAuditor(
auditor := audit.NewAuditor(
audit.DefaultFilter,
backends.NewPostgres(api.Database, true),
backends.NewSlog(api.Logger),
)
api.AGPL.Auditor.Store(&auditor)
}
api.AGPL.Auditor.Store(&auditor)
}
if entitlements.browserOnly != api.entitlements.browserOnly {
@@ -205,6 +218,13 @@ func (api *API) updateEntitlements(ctx context.Context) error {
api.AGPL.WorkspaceClientCoordinateOverride.Store(&handler)
}
if entitlements.workspaceQuota != api.entitlements.workspaceQuota {
if entitlements.workspaceQuota != codersdk.EntitlementNotEntitled && api.UserWorkspaceQuota > 0 {
enforcer := NewEnforcer(api.Options.UserWorkspaceQuota)
api.AGPL.WorkspaceQuotaEnforcer.Store(&enforcer)
}
}
api.entitlements = entitlements
return nil
@@ -260,6 +280,15 @@ func (api *API) serveEntitlements(rw http.ResponseWriter, r *http.Request) {
"Browser only connections are enabled but your license for this feature is expired.")
}
resp.Features[codersdk.FeatureWorkspaceQuota] = codersdk.Feature{
Entitlement: entitlements.workspaceQuota,
Enabled: api.UserWorkspaceQuota > 0,
}
if entitlements.workspaceQuota == codersdk.EntitlementGracePeriod && api.UserWorkspaceQuota > 0 {
resp.Warnings = append(resp.Warnings,
"Workspace quotas are enabled but your license for this feature is expired.")
}
httpapi.Write(ctx, rw, http.StatusOK, resp)
}
enterprise/coderd/coderdenttest/coderdenttest.go
+20 -12
View File
@@ -39,6 +39,7 @@ type Options struct {
BrowserOnly bool
EntitlementsUpdateInterval time.Duration
SCIMAPIKey []byte
UserWorkspaceQuota int
}
// New constructs a codersdk client connected to an in-memory Enterprise API instance.
@@ -59,6 +60,7 @@ func NewWithAPI(t *testing.T, options *Options) (*codersdk.Client, io.Closer, *c
AuditLogging: true,
BrowserOnly: options.BrowserOnly,
SCIMAPIKey: options.SCIMAPIKey,
UserWorkspaceQuota: options.UserWorkspaceQuota,
Options: oop,
EntitlementsUpdateInterval: options.EntitlementsUpdateInterval,
Keys: map[string]ed25519.PublicKey{
@@ -80,14 +82,15 @@ func NewWithAPI(t *testing.T, options *Options) (*codersdk.Client, io.Closer, *c
}
type LicenseOptions struct {
AccountType string
AccountID string
GraceAt time.Time
ExpiresAt time.Time
UserLimit int64
AuditLog bool
BrowserOnly bool
SCIM bool
AccountType string
AccountID string
GraceAt time.Time
ExpiresAt time.Time
UserLimit int64
AuditLog bool
BrowserOnly bool
SCIM bool
WorkspaceQuota bool
}
// AddLicense generates a new license with the options provided and inserts it.
@@ -119,6 +122,10 @@ func GenerateLicense(t *testing.T, options LicenseOptions) string {
if options.SCIM {
scim = 1
}
workspaceQuota := int64(0)
if options.WorkspaceQuota {
workspaceQuota = 1
}
c := &coderd.Claims{
RegisteredClaims: jwt.RegisteredClaims{
@@ -132,10 +139,11 @@ func GenerateLicense(t *testing.T, options LicenseOptions) string {
AccountID: options.AccountID,
Version: coderd.CurrentVersion,
Features: coderd.Features{
UserLimit: options.UserLimit,
AuditLog: auditLog,
BrowserOnly: browserOnly,
SCIM: scim,
UserLimit: options.UserLimit,
AuditLog: auditLog,
BrowserOnly: browserOnly,
SCIM: scim,
WorkspaceQuota: workspaceQuota,
},
}
tok := jwt.NewWithClaims(jwt.SigningMethodEdDSA, c)
enterprise/coderd/licenses.go
+5 -4
View File
@@ -45,10 +45,11 @@ var key20220812 []byte
var Keys = map[string]ed25519.PublicKey{"2022-08-12": ed25519.PublicKey(key20220812)}
type Features struct {
UserLimit int64 `json:"user_limit"`
AuditLog int64 `json:"audit_log"`
BrowserOnly int64 `json:"browser_only"`
SCIM int64 `json:"scim"`
UserLimit int64 `json:"user_limit"`
AuditLog int64 `json:"audit_log"`
BrowserOnly int64 `json:"browser_only"`
SCIM int64 `json:"scim"`
WorkspaceQuota int64 `json:"workspace_quota"`
}
type Claims struct {
enterprise/coderd/licenses_test.go
+10 -8
View File
@@ -98,18 +98,20 @@ func TestGetLicense(t *testing.T) {
assert.Equal(t, int32(1), licenses[0].ID)
assert.Equal(t, "testing", licenses[0].Claims["account_id"])
assert.Equal(t, map[string]interface{}{
codersdk.FeatureUserLimit: json.Number("0"),
codersdk.FeatureAuditLog: json.Number("1"),
codersdk.FeatureSCIM: json.Number("1"),
codersdk.FeatureBrowserOnly: json.Number("1"),
codersdk.FeatureUserLimit: json.Number("0"),
codersdk.FeatureAuditLog: json.Number("1"),
codersdk.FeatureSCIM: json.Number("1"),
codersdk.FeatureBrowserOnly: json.Number("1"),
codersdk.FeatureWorkspaceQuota: json.Number("0"),
}, licenses[0].Claims["features"])
assert.Equal(t, int32(2), licenses[1].ID)
assert.Equal(t, "testing2", licenses[1].Claims["account_id"])
assert.Equal(t, map[string]interface{}{
codersdk.FeatureUserLimit: json.Number("200"),
codersdk.FeatureAuditLog: json.Number("1"),
codersdk.FeatureSCIM: json.Number("1"),
codersdk.FeatureBrowserOnly: json.Number("1"),
codersdk.FeatureUserLimit: json.Number("200"),
codersdk.FeatureAuditLog: json.Number("1"),
codersdk.FeatureSCIM: json.Number("1"),
codersdk.FeatureBrowserOnly: json.Number("1"),
codersdk.FeatureWorkspaceQuota: json.Number("0"),
}, licenses[1].Claims["features"])
})
}
enterprise/coderd/workspacequota.go
+60
View File
@@ -0,0 +1,60 @@
package coderd
import (
"net/http"
"github.com/coder/coder/coderd/database"
"github.com/coder/coder/coderd/httpapi"
"github.com/coder/coder/coderd/httpmw"
"github.com/coder/coder/coderd/rbac"
"github.com/coder/coder/coderd/workspacequota"
"github.com/coder/coder/codersdk"
)
type enforcer struct {
userWorkspaceLimit int
}
func NewEnforcer(userWorkspaceLimit int) workspacequota.Enforcer {
return &enforcer{
userWorkspaceLimit: userWorkspaceLimit,
}
}
func (e *enforcer) UserWorkspaceLimit() int {
return e.userWorkspaceLimit
}
func (e *enforcer) CanCreateWorkspace(count int) bool {
if e.userWorkspaceLimit == 0 {
return true
}
return count < e.userWorkspaceLimit
}
func (api *API) workspaceQuota(rw http.ResponseWriter, r *http.Request) {
user := httpmw.UserParam(r)
if !api.AGPL.Authorize(r, rbac.ActionRead, rbac.ResourceUser) {
httpapi.ResourceNotFound(rw)
return
}
workspaces, err := api.Database.GetWorkspaces(r.Context(), database.GetWorkspacesParams{
OwnerID: user.ID,
})
if err != nil {
httpapi.Write(r.Context(), rw, http.StatusInternalServerError, codersdk.Response{
Message: "Internal error fetching workspaces.",
Detail: err.Error(),
})
return
}
e := *api.AGPL.WorkspaceQuotaEnforcer.Load()
httpapi.Write(r.Context(), rw, http.StatusOK, codersdk.WorkspaceQuota{
UserWorkspaceCount: len(workspaces),
UserWorkspaceLimit: e.UserWorkspaceLimit(),
})
}
enterprise/coderd/workspacequota_test.go
+122
View File
@@ -0,0 +1,122 @@
package coderd_test
import (
"context"
"testing"
"time"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
"github.com/coder/coder/coderd/coderdtest"
"github.com/coder/coder/coderd/util/ptr"
"github.com/coder/coder/codersdk"
"github.com/coder/coder/enterprise/coderd/coderdenttest"
"github.com/coder/coder/provisioner/echo"
"github.com/coder/coder/provisionersdk/proto"
"github.com/coder/coder/testutil"
)
func TestWorkspaceQuota(t *testing.T) {
t.Parallel()
t.Run("Disabled", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
client := coderdenttest.New(t, &coderdenttest.Options{})
_ = coderdtest.CreateFirstUser(t, client)
coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
WorkspaceQuota: true,
})
q1, err := client.WorkspaceQuota(ctx, codersdk.Me)
require.NoError(t, err)
require.EqualValues(t, q1.UserWorkspaceLimit, 0)
})
t.Run("Enabled", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
max := 3
client := coderdenttest.New(t, &coderdenttest.Options{
UserWorkspaceQuota: max,
})
user := coderdtest.CreateFirstUser(t, client)
coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
WorkspaceQuota: true,
})
q1, err := client.WorkspaceQuota(ctx, codersdk.Me)
require.NoError(t, err)
require.EqualValues(t, q1.UserWorkspaceLimit, max)
// ensure other user IDs work too
u2, err := client.CreateUser(ctx, codersdk.CreateUserRequest{
Email: "whatever@yo.com",
Username: "haha",
Password: "laskjdnvkaj",
OrganizationID: user.OrganizationID,
})
require.NoError(t, err)
q2, err := client.WorkspaceQuota(ctx, u2.ID.String())
require.NoError(t, err)
require.EqualValues(t, q1, q2)
})
t.Run("BlocksBuild", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
max := 1
client := coderdenttest.New(t, &coderdenttest.Options{
UserWorkspaceQuota: max,
Options: &coderdtest.Options{
IncludeProvisionerDaemon: true,
},
})
user := coderdtest.CreateFirstUser(t, client)
coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
WorkspaceQuota: true,
})
q1, err := client.WorkspaceQuota(ctx, codersdk.Me)
require.NoError(t, err)
require.EqualValues(t, q1.UserWorkspaceCount, 0)
require.EqualValues(t, q1.UserWorkspaceLimit, max)
authToken := uuid.NewString()
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
Parse: echo.ParseComplete,
Provision: []*proto.Provision_Response{{
Type: &proto.Provision_Response_Complete{
Complete: &proto.Provision_Complete{
Resources: []*proto.Resource{{
Name: "example",
Type: "aws_instance",
Agents: []*proto.Agent{{
Id: uuid.NewString(),
Name: "example",
Auth: &proto.Agent_Token{
Token: authToken,
},
}},
}},
},
},
}},
})
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
_ = coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
_, err = client.CreateWorkspace(context.Background(), user.OrganizationID, codersdk.Me, codersdk.CreateWorkspaceRequest{
TemplateID: template.ID,
Name: "ajksdnvksjd",
AutostartSchedule: ptr.Ref("CRON_TZ=US/Central 30 9 * * 1-5"),
TTLMillis: ptr.Ref((8 * time.Hour).Milliseconds()),
})
require.Error(t, err)
require.ErrorContains(t, err, "User workspace limit")
// ensure count increments
q1, err = client.WorkspaceQuota(ctx, codersdk.Me)
require.NoError(t, err)
require.EqualValues(t, q1.UserWorkspaceCount, 1)
require.EqualValues(t, q1.UserWorkspaceLimit, max)
})
}