feat: workspace quotas (#4184)

enterprise/coderd/coderd.go

@@ -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)
 }