coder · code-asher · Aug 19, 2024 · Aug 19, 2024 · Aug 19, 2024 · Aug 19, 2024
diff --git a/codersdk/deployment.go b/codersdk/deployment.go
@@ -128,6 +128,17 @@ func (n FeatureName) AlwaysEnable() bool {
 	}[n]
 }
 
+// Enterprise returns true if the feature is an enterprise feature.
+func (n FeatureName) Enterprise() bool {
+	switch n {
+	// Add all features that should be excluded in the Enterprise feature set.
+	case FeatureMultipleOrganizations, FeatureCustomRoles:
+		return false
+	default:
+		return true
+	}
+}
+
 // FeatureSet represents a grouping of features. Rather than manually
 // assigning features al-la-carte when making a license, a set can be specified.
 // Sets are dynamic in the sense a feature can be added to a set, granting the
@@ -152,13 +163,7 @@ func (set FeatureSet) Features() []FeatureName {
 		copy(enterpriseFeatures, FeatureNames)
 		// Remove the selection
 		enterpriseFeatures = slices.DeleteFunc(enterpriseFeatures, func(f FeatureName) bool {
-			switch f {
-			// Add all features that should be excluded in the Enterprise feature set.
-			case FeatureMultipleOrganizations, FeatureCustomRoles:
-				return true
-			default:
-				return false
-			}
+			return !f.Enterprise()
 		})
 
 		return enterpriseFeatures

diff --git a/enterprise/coderd/organizations_test.go b/enterprise/coderd/organizations_test.go
@@ -488,6 +488,46 @@ func TestPatchOrganizationsByUser(t *testing.T) {
 		require.Equal(t, displayName, o.DisplayName) // didn't change
 		require.Equal(t, icon, o.Icon)
 	})
+
+	t.Run("RevokedLicense", func(t *testing.T) {
+		t.Parallel()
+		dv := coderdtest.DeploymentValues(t)
+		dv.Experiments = []string{string(codersdk.ExperimentMultiOrganization)}
+		client, _ := coderdenttest.New(t, &coderdenttest.Options{
+			Options: &coderdtest.Options{
+				DeploymentValues: dv,
+			},
+			LicenseOptions: &coderdenttest.LicenseOptions{
+				Features: license.Features{
+					codersdk.FeatureMultipleOrganizations: 1,
+				},
+			},
+		})
+		ctx := testutil.Context(t, testutil.WaitMedium)
+
+		const displayName = "New Organization"
+		o := coderdenttest.CreateOrganization(t, client, coderdenttest.CreateOrganizationOptions{}, func(request *codersdk.CreateOrganizationRequest) {
+			request.DisplayName = displayName
+			request.Icon = "/emojis/random.png"
+			request.Name = "new-org"
+		})
+
+		// Remove the license to block premium functionality.
+		licenses, err := client.Licenses(ctx)
+		require.NoError(t, err, "get licenses")
+		for _, license := range licenses {
+			// Should be only 1...
+			err := client.DeleteLicense(ctx, license.ID)
+			require.NoError(t, err, "delete license")
+		}
+
+		// Verify functionality is lost.
+		const icon = "/emojis/1f48f-1f3ff.png"
+		o, err = client.UpdateOrganization(ctx, o.Name, codersdk.UpdateOrganizationRequest{
+			Icon: ptr.Ref(icon),
+		})
+		require.ErrorContains(t, err, "Multiple Organizations is a Premium feature")
+	})
 }
 
 func TestPostOrganizationsByUser(t *testing.T) {

diff --git a/enterprise/coderd/roles_test.go b/enterprise/coderd/roles_test.go
@@ -114,7 +114,7 @@ func TestCustomOrganizationRole(t *testing.T) {
 		role, err := owner.CreateOrganizationRole(ctx, templateAdminCustom(first.OrganizationID))
 		require.NoError(t, err, "upsert role")
 
-		// Remove the license to block enterprise functionality
+		// Remove the license to block premium functionality
 		licenses, err := owner.Licenses(ctx)
 		require.NoError(t, err, "get licenses")
 		for _, license := range licenses {
@@ -125,7 +125,7 @@ func TestCustomOrganizationRole(t *testing.T) {
 
 		// Verify functionality is lost
 		_, err = owner.UpdateOrganizationRole(ctx, templateAdminCustom(first.OrganizationID))
-		require.ErrorContains(t, err, "Custom Roles is an Enterprise feature")
+		require.ErrorContains(t, err, "Custom Roles is a Premium feature")
 
 		// Assign the custom template admin role
 		tmplAdmin, _ := coderdtest.CreateAnotherUser(t, owner, first.OrganizationID, rbac.RoleIdentifier{Name: role.Name, OrganizationID: first.OrganizationID})

diff --git a/enterprise/coderd/templates.go b/enterprise/coderd/templates.go
@@ -364,8 +364,12 @@ func (api *API) RequireFeatureMW(feat codersdk.FeatureName) func(http.Handler) h
 			enabled := api.entitlements.Features[feat].Enabled
 			api.entitlementsMu.RUnlock()
 			if !enabled {
+				licenseType := "a Premium"
+				if feat.Enterprise() {
+					licenseType = "an Enterprise"
+				}
 				httpapi.Write(r.Context(), rw, http.StatusForbidden, codersdk.Response{
-					Message: fmt.Sprintf("%s is an Enterprise feature. Contact sales!", feat.Humanize()),
+					Message: fmt.Sprintf("%s is %s feature. Contact sales!", feat.Humanize(), licenseType),
 				})
 				return
 			}