coder · f0ssel · Sep 30, 2022 · Aug 30, 2022 · Aug 31, 2022 · Aug 31, 2022
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -35,6 +35,7 @@ import (
 	"github.com/coder/coder/coderd/rbac"
 	"github.com/coder/coder/coderd/telemetry"
 	"github.com/coder/coder/coderd/tracing"
+	"github.com/coder/coder/coderd/workspacequota"
 	"github.com/coder/coder/coderd/wsconncache"
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/site"
@@ -55,6 +56,7 @@ type Options struct {
 	CacheDir string
 
 	Auditor                        audit.Auditor
+	WorkspaceQuotaEnforcer         workspacequota.Enforcer
 	AgentConnectionUpdateFrequency time.Duration
 	AgentInactiveDisconnectTimeout time.Duration
 	// APIRateLimit is the minutely throughput rate limit per user or ip.
@@ -120,6 +122,9 @@ func New(options *Options) *API {
 	if options.Auditor == nil {
 		options.Auditor = audit.NewNop()
 	}
+	if options.WorkspaceQuotaEnforcer == nil {
+		options.WorkspaceQuotaEnforcer = workspacequota.NewNop()
+	}
 
 	siteCacheDir := options.CacheDir
 	if siteCacheDir != "" {
@@ -145,10 +150,12 @@ func New(options *Options) *API {
 			Authorizer: options.Authorizer,
 			Logger:     options.Logger,
 		},
-		metricsCache: metricsCache,
-		Auditor:      atomic.Pointer[audit.Auditor]{},
+		metricsCache:           metricsCache,
+		Auditor:                atomic.Pointer[audit.Auditor]{},
+		WorkspaceQuotaEnforcer: atomic.Pointer[workspacequota.Enforcer]{},
 	}
 	api.Auditor.Store(&options.Auditor)
+	api.WorkspaceQuotaEnforcer.Store(&options.WorkspaceQuotaEnforcer)
 	api.workspaceAgentCache = wsconncache.New(api.dialWorkspaceAgentTailnet, 0)
 	api.derpServer = derp.NewServer(key.NewNode(), tailnet.Logger(options.Logger))
 	oauthConfigs := &httpmw.OAuth2Configs{
@@ -516,6 +523,7 @@ type API struct {
 	*Options
 	Auditor                           atomic.Pointer[audit.Auditor]
 	WorkspaceClientCoordinateOverride atomic.Pointer[func(rw http.ResponseWriter) bool]
+	WorkspaceQuotaEnforcer            atomic.Pointer[workspacequota.Enforcer]
 	HTTPAuth                          *HTTPAuthorizer
 
 	// APIHandler serves "/api/v2"

diff --git a/coderd/database/databasefake/databasefake.go b/coderd/database/databasefake/databasefake.go
@@ -698,6 +698,22 @@ func (q *fakeQuerier) GetWorkspaceBuildByID(_ context.Context, id uuid.UUID) (da
 	return database.WorkspaceBuild{}, sql.ErrNoRows
 }
 
+func (q *fakeQuerier) GetWorkspaceCountByUserID(_ context.Context, id uuid.UUID) (int64, error) {
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+	var count int64
+	for _, workspace := range q.workspaces {
+		if workspace.OwnerID.String() == id.String() {
+			if workspace.Deleted {
+				continue
+			}
+
+			count++
+		}
+	}
+	return count, nil
+}
+
 func (q *fakeQuerier) GetWorkspaceBuildByJobID(_ context.Context, jobID uuid.UUID) (database.WorkspaceBuild, error) {
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()

diff --git a/coderd/database/querier.go b/coderd/database/querier.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspaces.sql b/coderd/database/queries/workspaces.sql
@@ -74,6 +74,16 @@ WHERE
 GROUP BY
 	template_id;
 
+-- name: GetWorkspaceCountByUserID :one
+SELECT
+	COUNT(id)
+FROM
+	workspaces
+WHERE
+	owner_id = @owner_id
+	-- Ignore deleted workspaces
+	AND deleted != true;
+
 -- name: InsertWorkspace :one
 INSERT INTO
 	workspaces (

diff --git a/coderd/workspacequota/workspacequota.go b/coderd/workspacequota/workspacequota.go
@@ -0,0 +1,19 @@
+package workspacequota
+
+type Enforcer interface {
+	UserWorkspaceLimit() int
+	CanCreateWorkspace(count int) bool
+}
+
+type nop struct{}
+
+func NewNop() Enforcer {
+	return &nop{}
+}
+
+func (*nop) UserWorkspaceLimit() int {
+	return 0
+}
+func (*nop) CanCreateWorkspace(_ int) bool {
+	return true
+}
diff --git a/coderd/workspaces.go b/coderd/workspaces.go
@@ -317,6 +317,25 @@ func (api *API) postWorkspacesByOrganization(rw http.ResponseWriter, r *http.Req
 		return
 	}
 
+	workspaceCount, err := api.Database.GetWorkspaceCountByUserID(ctx, user.ID)
+	if err != nil {
+		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Internal error fetching workspace count.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+
+	// make sure the user has not hit their quota limit
+	e := *api.WorkspaceQuotaEnforcer.Load()
+	canCreate := e.CanCreateWorkspace(int(workspaceCount))
+	if !canCreate {
+		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
+			Message: fmt.Sprintf("User workspace limit of %d is already reached.", e.UserWorkspaceLimit()),
+		})
+		return
+	}
+
 	templateVersion, err := api.Database.GetTemplateVersionByID(ctx, template.ActiveVersionID)
 	if err != nil {
 		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
@@ -352,8 +371,10 @@ func (api *API) postWorkspacesByOrganization(rw http.ResponseWriter, r *http.Req
 		return
 	}
 
-	var provisionerJob database.ProvisionerJob
-	var workspaceBuild database.WorkspaceBuild
+	var (
+		provisionerJob database.ProvisionerJob
+		workspaceBuild database.WorkspaceBuild
+	)
 	err = api.Database.InTx(func(db database.Store) error {
 		now := database.Now()
 		workspaceBuildID := uuid.New()

diff --git a/codersdk/features.go b/codersdk/features.go
@@ -15,13 +15,20 @@ const (
 )
 
 const (
-	FeatureUserLimit   = "user_limit"
-	FeatureAuditLog    = "audit_log"
-	FeatureBrowserOnly = "browser_only"
-	FeatureSCIM        = "scim"
+	FeatureUserLimit      = "user_limit"
+	FeatureAuditLog       = "audit_log"
+	FeatureBrowserOnly    = "browser_only"
+	FeatureSCIM           = "scim"
+	FeatureWorkspaceQuota = "workspace_quota"
 )
 
-var FeatureNames = []string{FeatureUserLimit, FeatureAuditLog, FeatureBrowserOnly, FeatureSCIM}
+var FeatureNames = []string{
+	FeatureUserLimit,
+	FeatureAuditLog,
+	FeatureBrowserOnly,
+	FeatureSCIM,
+	FeatureWorkspaceQuota,
+}
 
 type Feature struct {
 	Entitlement Entitlement `json:"entitlement"`

diff --git a/codersdk/workspacequota.go b/codersdk/workspacequota.go
@@ -0,0 +1,26 @@
+package codersdk
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+)
+
+type WorkspaceQuota struct {
+	UserWorkspaceCount int `json:"user_workspace_count"`
+	UserWorkspaceLimit int `json:"user_workspace_limit"`
+}
+
+func (c *Client) WorkspaceQuota(ctx context.Context, userID string) (WorkspaceQuota, error) {
+	res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/workspace-quota/%s", userID), nil)
+	if err != nil {
+		return WorkspaceQuota{}, err
+	}
+	defer res.Body.Close()
+	if res.StatusCode != http.StatusOK {
+		return WorkspaceQuota{}, readBodyAsError(res)
+	}
+	var quota WorkspaceQuota
+	return quota, json.NewDecoder(res.Body).Decode(&quota)
+}
diff --git a/enterprise/cli/features_test.go b/enterprise/cli/features_test.go
@@ -57,14 +57,16 @@ func TestFeaturesList(t *testing.T) {
 		var entitlements codersdk.Entitlements
 		err := json.Unmarshal(buf.Bytes(), &entitlements)
 		require.NoError(t, err, "unmarshal JSON output")
-		assert.Len(t, entitlements.Features, 3)
+		assert.Len(t, entitlements.Features, 4)
 		assert.Empty(t, entitlements.Warnings)
 		assert.Equal(t, codersdk.EntitlementNotEntitled,
 			entitlements.Features[codersdk.FeatureUserLimit].Entitlement)
 		assert.Equal(t, codersdk.EntitlementNotEntitled,
 			entitlements.Features[codersdk.FeatureAuditLog].Entitlement)
 		assert.Equal(t, codersdk.EntitlementNotEntitled,
 			entitlements.Features[codersdk.FeatureBrowserOnly].Entitlement)
+		assert.Equal(t, codersdk.EntitlementNotEntitled,
+			entitlements.Features[codersdk.FeatureWorkspaceQuota].Entitlement)
 		assert.False(t, entitlements.HasLicense)
 	})
 }
diff --git a/enterprise/cli/server.go b/enterprise/cli/server.go
@@ -15,16 +15,18 @@ import (
 
 func server() *cobra.Command {
 	var (
-		auditLogging   bool
-		browserOnly    bool
-		scimAuthHeader string
+		auditLogging       bool
+		browserOnly        bool
+		scimAuthHeader     string
+		userWorkspaceQuota int
 	)
 	cmd := agpl.Server(func(ctx context.Context, options *agplcoderd.Options) (*agplcoderd.API, error) {
 		api, err := coderd.New(ctx, &coderd.Options{
-			AuditLogging: auditLogging,
-			BrowserOnly:  browserOnly,
-			SCIMAPIKey:   []byte(scimAuthHeader),
-			Options:      options,
+			AuditLogging:       auditLogging,
+			BrowserOnly:        browserOnly,
+			SCIMAPIKey:         []byte(scimAuthHeader),
+			UserWorkspaceQuota: userWorkspaceQuota,
+			Options:            options,
 		})
 		if err != nil {
 			return nil, err
@@ -39,6 +41,8 @@ func server() *cobra.Command {
 		"Whether Coder only allows connections to workspaces via the browser. "+enterpriseOnly)
 	cliflag.StringVarP(cmd.Flags(), &scimAuthHeader, "scim-auth-header", "", "CODER_SCIM_API_KEY", "",
 		"Enables SCIM and sets the authentication header for the built-in SCIM server. New users are automatically created with OIDC authentication. "+enterpriseOnly)
+	cliflag.IntVarP(cmd.Flags(), &userWorkspaceQuota, "user-workspace-quota", "", "CODER_USER_WORKSPACE_QUOTA", 0,
+		"A positive number applies a limit on how many workspaces each user can create. "+enterpriseOnly)
 
 	return cmd
 }