feat: implement premium vs enterprise licenses (#13907)

* feat: implement premium vs enterprise licenses

Implement different sets of licensed features.

codersdk/deployment.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"path/filepath"
 	"reflect"
+	"slices"
 	"strconv"
 	"strings"
 	"time"
@@ -34,6 +35,21 @@ const (
 	EntitlementNotEntitled Entitlement = "not_entitled"
 )
 
+// Weight converts the enum types to a numerical value for easier
+// comparisons. Easier than sets of if statements.
+func (e Entitlement) Weight() int {
+	switch e {
+	case EntitlementEntitled:
+		return 2
+	case EntitlementGracePeriod:
+		return 1
+	case EntitlementNotEntitled:
+		return -1
+	default:
+		return -2
+	}
+}
+
 // FeatureName represents the internal name of a feature.
 // To add a new feature, add it to this set of enums as well as the FeatureNames
 // array below.
@@ -95,8 +111,11 @@ func (n FeatureName) Humanize() string {
 }
 
 // AlwaysEnable returns if the feature is always enabled if entitled.
-// Warning: We don't know if we need this functionality.
-// This method may disappear at any time.
+// This is required because some features are only enabled if they are entitled
+// and not required.
+// E.g: "multiple-organizations" is disabled by default in AGPL and enterprise
+// deployments. This feature should only be enabled for premium deployments
+// when it is entitled.
 func (n FeatureName) AlwaysEnable() bool {
 	return map[FeatureName]bool{
 		FeatureMultipleExternalAuth:       true,
@@ -105,9 +124,54 @@ func (n FeatureName) AlwaysEnable() bool {
 		FeatureWorkspaceBatchActions:      true,
 		FeatureHighAvailability:           true,
 		FeatureCustomRoles:                true,
+		FeatureMultipleOrganizations:      true,
 	}[n]
 }
 
+// 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
+// feature to existing licenses out in the wild.
+// If features were granted al-la-carte, we would need to reissue the existing
+// old licenses to include the new feature.
+type FeatureSet string
+
+const (
+	FeatureSetNone       FeatureSet = ""
+	FeatureSetEnterprise FeatureSet = "enterprise"
+	FeatureSetPremium    FeatureSet = "premium"
+)
+
+func (set FeatureSet) Features() []FeatureName {
+	switch FeatureSet(strings.ToLower(string(set))) {
+	case FeatureSetEnterprise:
+		// Enterprise is the set 'AllFeatures' minus some select features.
+
+		// Copy the list of all features
+		enterpriseFeatures := make([]FeatureName, len(FeatureNames))
+		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:
+				return true
+			default:
+				return false
+			}
+		})
+
+		return enterpriseFeatures
+	case FeatureSetPremium:
+		premiumFeatures := make([]FeatureName, len(FeatureNames))
+		copy(premiumFeatures, FeatureNames)
+		// FeatureSetPremium is just all features.
+		return premiumFeatures
+	}
+	// By default, return an empty set.
+	return []FeatureName{}
+}
+
 type Feature struct {
 	Entitlement Entitlement `json:"entitlement"`
 	Enabled     bool        `json:"enabled"`
@@ -115,6 +179,89 @@ type Feature struct {
 	Actual      *int64      `json:"actual,omitempty"`
 }
 
+// Compare compares two features and returns an integer representing
+// if the first feature (f) is greater than, equal to, or less than the second
+// feature (b). "Greater than" means the first feature has more functionality
+// than the second feature. It is assumed the features are for the same FeatureName.
+//
+// A feature is considered greater than another feature if:
+// 1. Graceful & capable > Entitled & not capable
+// 2. The entitlement is greater
+// 3. The limit is greater
+// 4. Enabled is greater than disabled
+// 5. The actual is greater
+func (f Feature) Compare(b Feature) int {
+	if !f.Capable() || !b.Capable() {
+		// If either is incapable, then it is possible a grace period
+		// feature can be "greater" than an entitled.
+		// If either is "NotEntitled" then we can defer to a strict entitlement
+		// check.
+		if f.Entitlement.Weight() >= 0 && b.Entitlement.Weight() >= 0 {
+			if f.Capable() && !b.Capable() {
+				return 1
+			}
+			if b.Capable() && !f.Capable() {
+				return -1
+			}
+		}
+	}
+
+	// Strict entitlement check. Higher is better
+	entitlementDifference := f.Entitlement.Weight() - b.Entitlement.Weight()
+	if entitlementDifference != 0 {
+		return entitlementDifference
+	}
+
+	// If the entitlement is the same, then we can compare the limits.
+	if f.Limit == nil && b.Limit != nil {
+		return -1
+	}
+	if f.Limit != nil && b.Limit == nil {
+		return 1
+	}
+	if f.Limit != nil && b.Limit != nil {
+		difference := *f.Limit - *b.Limit
+		if difference != 0 {
+			return int(difference)
+		}
+	}
+
+	// Enabled is better than disabled.
+	if f.Enabled && !b.Enabled {
+		return 1
+	}
+	if !f.Enabled && b.Enabled {
+		return -1
+	}
+
+	// Higher actual is better
+	if f.Actual == nil && b.Actual != nil {
+		return -1
+	}
+	if f.Actual != nil && b.Actual == nil {
+		return 1
+	}
+	if f.Actual != nil && b.Actual != nil {
+		difference := *f.Actual - *b.Actual
+		if difference != 0 {
+			return int(difference)
+		}
+	}
+
+	return 0
+}
+
+// Capable is a helper function that returns if a given feature has a limit
+// that is greater than or equal to the actual.
+// If this condition is not true, then the feature is not capable of being used
+// since the limit is not high enough.
+func (f Feature) Capable() bool {
+	if f.Limit != nil && f.Actual != nil {
+		return *f.Limit >= *f.Actual
+	}
+	return true
+}
+
 type Entitlements struct {
 	Features         map[FeatureName]Feature `json:"features"`
 	Warnings         []string                `json:"warnings"`
@@ -125,6 +272,29 @@ type Entitlements struct {
 	RefreshedAt      time.Time               `json:"refreshed_at" format:"date-time"`
 }
 
+// AddFeature will add the feature to the entitlements iff it expands
+// the set of features granted by the entitlements. If it does not, it will
+// be ignored and the existing feature with the same name will remain.
+//
+// All features should be added as atomic items, and not merged in any way.
+// Merging entitlements could lead to unexpected behavior, like a larger user
+// limit in grace period merging with a smaller one in an "entitled" state. This
+// could lead to the larger limit being extended as "entitled", which is not correct.
+func (e *Entitlements) AddFeature(name FeatureName, add Feature) {
+	existing, ok := e.Features[name]
+	if !ok {
+		e.Features[name] = add
+		return
+	}
+
+	// Compare the features, keep the one that is "better"
+	comparison := add.Compare(existing)
+	if comparison > 0 {
+		e.Features[name] = add
+		return
+	}
+}
+
 func (c *Client) Entitlements(ctx context.Context) (Entitlements, error) {
 	res, err := c.Request(ctx, http.MethodGet, "/api/v2/entitlements", nil)
 	if err != nil {

codersdk/deployment_test.go

@@ -3,15 +3,18 @@ package codersdk_test
 import (
 	"bytes"
 	"embed"
+	"encoding/json"
 	"fmt"
 	"runtime"
 	"strings"
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"gopkg.in/yaml.v3"
 
+	"github.com/coder/coder/v2/coderd/util/ptr"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/serpent"
 )
@@ -379,3 +382,182 @@ func TestExternalAuthYAMLConfig(t *testing.T) {
 	output := strings.Replace(out.String(), "value:", "externalAuthProviders:", 1)
 	require.Equal(t, inputYAML, output, "re-marshaled is the same as input")
 }
+
+func TestFeatureComparison(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		Name     string
+		A        codersdk.Feature
+		B        codersdk.Feature
+		Expected int
+	}{
+		{
+			Name:     "Empty",
+			Expected: 0,
+		},
+		// Entitlement check
+		//		Entitled
+		{
+			Name:     "EntitledVsGracePeriod",
+			A:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled},
+			B:        codersdk.Feature{Entitlement: codersdk.EntitlementGracePeriod},
+			Expected: 1,
+		},
+		{
+			Name: "EntitledVsGracePeriodLimits",
+			A:    codersdk.Feature{Entitlement: codersdk.EntitlementEntitled},
+			// Entitled should still win here
+			B:        codersdk.Feature{Entitlement: codersdk.EntitlementGracePeriod, Limit: ptr.Ref[int64](100), Actual: ptr.Ref[int64](50)},
+			Expected: 1,
+		},
+		{
+			Name:     "EntitledVsNotEntitled",
+			A:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled},
+			B:        codersdk.Feature{Entitlement: codersdk.EntitlementNotEntitled},
+			Expected: 3,
+		},
+		{
+			Name:     "EntitledVsUnknown",
+			A:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled},
+			B:        codersdk.Feature{Entitlement: ""},
+			Expected: 4,
+		},
+		//		GracePeriod
+		{
+			Name:     "GracefulVsNotEntitled",
+			A:        codersdk.Feature{Entitlement: codersdk.EntitlementGracePeriod},
+			B:        codersdk.Feature{Entitlement: codersdk.EntitlementNotEntitled},
+			Expected: 2,
+		},
+		{
+			Name:     "GracefulVsUnknown",
+			A:        codersdk.Feature{Entitlement: codersdk.EntitlementGracePeriod},
+			B:        codersdk.Feature{Entitlement: ""},
+			Expected: 3,
+		},
+		//		NotEntitled
+		{
+			Name:     "NotEntitledVsUnknown",
+			A:        codersdk.Feature{Entitlement: codersdk.EntitlementNotEntitled},
+			B:        codersdk.Feature{Entitlement: ""},
+			Expected: 1,
+		},
+		// --
+		{
+			Name:     "EntitledVsGracePeriodCapable",
+			A:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled, Limit: ptr.Ref[int64](100), Actual: ptr.Ref[int64](200)},
+			B:        codersdk.Feature{Entitlement: codersdk.EntitlementGracePeriod, Limit: ptr.Ref[int64](300), Actual: ptr.Ref[int64](200)},
+			Expected: -1,
+		},
+		// UserLimits
+		{
+			// Tests an exceeded limit that is entitled vs a graceful limit that
+			// is not exceeded. This is the edge case that we should use the graceful period
+			// instead of the entitled.
+			Name:     "UserLimitExceeded",
+			A:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled, Limit: ptr.Ref(int64(100)), Actual: ptr.Ref(int64(200))},
+			B:        codersdk.Feature{Entitlement: codersdk.EntitlementGracePeriod, Limit: ptr.Ref(int64(300)), Actual: ptr.Ref(int64(200))},
+			Expected: -1,
+		},
+		{
+			Name:     "UserLimitExceededNoEntitled",
+			A:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled, Limit: ptr.Ref(int64(100)), Actual: ptr.Ref(int64(200))},
+			B:        codersdk.Feature{Entitlement: codersdk.EntitlementNotEntitled, Limit: ptr.Ref(int64(300)), Actual: ptr.Ref(int64(200))},
+			Expected: 3,
+		},
+		{
+			Name:     "HigherLimit",
+			A:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled, Limit: ptr.Ref(int64(110)), Actual: ptr.Ref(int64(200))},
+			B:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled, Limit: ptr.Ref(int64(100)), Actual: ptr.Ref(int64(200))},
+			Expected: 10, // Diff in the limit #
+		},
+		{
+			Name:     "HigherActual",
+			A:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled, Limit: ptr.Ref(int64(100)), Actual: ptr.Ref(int64(300))},
+			B:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled, Limit: ptr.Ref(int64(100)), Actual: ptr.Ref(int64(200))},
+			Expected: 100, // Diff in the actual #
+		},
+		{
+			Name:     "LimitExists",
+			A:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled, Limit: ptr.Ref(int64(100)), Actual: ptr.Ref(int64(50))},
+			B:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled, Limit: nil, Actual: ptr.Ref(int64(200))},
+			Expected: 1,
+		},
+		{
+			Name:     "LimitExistsGrace",
+			A:        codersdk.Feature{Entitlement: codersdk.EntitlementGracePeriod, Limit: ptr.Ref(int64(100)), Actual: ptr.Ref(int64(50))},
+			B:        codersdk.Feature{Entitlement: codersdk.EntitlementGracePeriod, Limit: nil, Actual: ptr.Ref(int64(200))},
+			Expected: 1,
+		},
+		{
+			Name:     "ActualExists",
+			A:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled, Limit: ptr.Ref(int64(100)), Actual: ptr.Ref(int64(50))},
+			B:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled, Limit: ptr.Ref(int64(100)), Actual: nil},
+			Expected: 1,
+		},
+		{
+			Name:     "NotNils",
+			A:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled, Limit: ptr.Ref(int64(100)), Actual: ptr.Ref(int64(50))},
+			B:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled, Limit: nil, Actual: nil},
+			Expected: 1,
+		},
+		{
+			Name:     "EnabledVsDisabled",
+			A:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled, Enabled: true, Limit: ptr.Ref(int64(300)), Actual: ptr.Ref(int64(200))},
+			B:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled, Limit: ptr.Ref(int64(300)), Actual: ptr.Ref(int64(200))},
+			Expected: 1,
+		},
+		{
+			Name:     "NotNils",
+			A:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled, Limit: ptr.Ref(int64(100)), Actual: ptr.Ref(int64(50))},
+			B:        codersdk.Feature{Entitlement: codersdk.EntitlementEntitled, Limit: nil, Actual: nil},
+			Expected: 1,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+
+		t.Run(tc.Name, func(t *testing.T) {
+			t.Parallel()
+
+			r := tc.A.Compare(tc.B)
+			logIt := !assert.Equal(t, tc.Expected, r)
+
+			// Comparisons should be like addition. A - B = -1 * (B - A)
+			r = tc.B.Compare(tc.A)
+			logIt = logIt || !assert.Equalf(t, tc.Expected*-1, r, "the inverse comparison should also be true")
+			if logIt {
+				ad, _ := json.Marshal(tc.A)
+				bd, _ := json.Marshal(tc.B)
+				t.Logf("a = %s\nb = %s", ad, bd)
+			}
+		})
+	}
+}
+
+// TestPremiumSuperSet tests that the "premium" feature set is a superset of the
+// "enterprise" feature set.
+func TestPremiumSuperSet(t *testing.T) {
+	t.Parallel()
+
+	enterprise := codersdk.FeatureSetEnterprise
+	premium := codersdk.FeatureSetPremium
+
+	// Premium > Enterprise
+	require.Greater(t, len(premium.Features()), len(enterprise.Features()), "premium should have more features than enterprise")
+
+	// Premium ⊃ Enterprise
+	require.Subset(t, premium.Features(), enterprise.Features(), "premium should be a superset of enterprise. If this fails, update the premium feature set to include all enterprise features.")
+
+	// Premium = All Features
+	// This is currently true. If this assertion changes, update this test
+	// to reflect the change in feature sets.
+	require.ElementsMatch(t, premium.Features(), codersdk.FeatureNames, "premium should contain all features")
+
+	// This check exists because if you misuse the slices.Delete, you can end up
+	// with zero'd values.
+	require.NotContains(t, enterprise.Features(), "", "enterprise should not contain empty string")
+	require.NotContains(t, premium.Features(), "", "premium should not contain empty string")
+}

enterprise/coderd/coderd.go

@@ -570,7 +570,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 
 	entitlements, err := license.Entitlements(
 		ctx, api.Database,
-		api.Logger, len(agedReplicas), len(api.ExternalAuthConfigs), api.LicenseKeys, map[codersdk.FeatureName]bool{
+		len(agedReplicas), len(api.ExternalAuthConfigs), api.LicenseKeys, map[codersdk.FeatureName]bool{
 			codersdk.FeatureAuditLog:                   api.AuditLogging,
 			codersdk.FeatureBrowserOnly:                api.BrowserOnly,
 			codersdk.FeatureSCIM:                       len(api.SCIMAPIKey) != 0,
@@ -583,7 +583,6 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 			codersdk.FeatureUserRoleManagement:         true,
 			codersdk.FeatureAccessControl:              true,
 			codersdk.FeatureControlSharedPorts:         true,
-			codersdk.FeatureMultipleOrganizations:      true,
 		})
 	if err != nil {
 		return err

enterprise/coderd/coderdenttest/coderdenttest.go

@@ -146,15 +146,55 @@ func NewWithAPI(t *testing.T, options *Options) (
 	return client, provisionerCloser, coderAPI, user
 }
 
+// LicenseOptions is used to generate a license for testing.
+// It supports the builder pattern for easy customization.
 type LicenseOptions struct {
 	AccountType   string
 	AccountID     string
 	DeploymentIDs []string
 	Trial         bool
+	FeatureSet    codersdk.FeatureSet
 	AllFeatures   bool
-	GraceAt       time.Time
-	ExpiresAt     time.Time
-	Features      license.Features
+	// GraceAt is the time at which the license will enter the grace period.
+	GraceAt time.Time
+	// ExpiresAt is the time at which the license will hard expire.
+	// ExpiresAt should always be greater then GraceAt.
+	ExpiresAt time.Time
+	Features  license.Features
+}
+
+func (opts *LicenseOptions) Expired(now time.Time) *LicenseOptions {
+	opts.ExpiresAt = now.Add(time.Hour * 24 * -2)
+	opts.GraceAt = now.Add(time.Hour * 24 * -3)
+	return opts
+}
+
+func (opts *LicenseOptions) GracePeriod(now time.Time) *LicenseOptions {
+	opts.ExpiresAt = now.Add(time.Hour * 24)
+	opts.GraceAt = now.Add(time.Hour * 24 * -1)
+	return opts
+}
+
+func (opts *LicenseOptions) Valid(now time.Time) *LicenseOptions {
+	opts.ExpiresAt = now.Add(time.Hour * 24 * 60)
+	opts.GraceAt = now.Add(time.Hour * 24 * 53)
+	return opts
+}
+
+func (opts *LicenseOptions) UserLimit(limit int64) *LicenseOptions {
+	return opts.Feature(codersdk.FeatureUserLimit, limit)
+}
+
+func (opts *LicenseOptions) Feature(name codersdk.FeatureName, value int64) *LicenseOptions {
+	if opts.Features == nil {
+		opts.Features = license.Features{}
+	}
+	opts.Features[name] = value
+	return opts
+}
+
+func (opts *LicenseOptions) Generate(t *testing.T) string {
+	return GenerateLicense(t, *opts)
 }
 
 // AddFullLicense generates a license with all features enabled.
@@ -195,6 +235,7 @@ func GenerateLicense(t *testing.T, options LicenseOptions) string {
 		Trial:          options.Trial,
 		Version:        license.CurrentVersion,
 		AllFeatures:    options.AllFeatures,
+		FeatureSet:     options.FeatureSet,
 		Features:       options.Features,
 	}
 	tok := jwt.NewWithClaims(jwt.SigningMethodEdDSA, c)

enterprise/coderd/license/doc.go

@@ -0,0 +1,32 @@
+// Package license provides the license parsing and validation logic for Coderd.
+// Licensing in Coderd defines what features are allowed to be used in a
+// given deployment. Without a license, or with a license that grants 0 features,
+// Coderd will refuse to execute some feature code paths. These features are
+// typically gated with a middleware that checks the license before allowing
+// the http request to proceed.
+//
+// Terms:
+// - FeatureName: A specific functionality that Coderd provides, such as
+//                external provisioners.
+//
+// - Feature: Entitlement definition for a FeatureName. A feature can be:
+//				- "entitled": The feature is allowed to be used by the deployment.
+//				- "grace period": The feature is allowed to be used by the deployment,
+//                                but the license is expired. There is a grace period
+//                                before the feature is disabled.
+//				- "not entitled": The deployment is not allowed to use the feature.
+//	                              Either by expiration, or by not being included
+//	                              in the license.
+//            A feature can also be "disabled" that prevents usage of the feature
+//            even if entitled. This is usually a deployment configuration option.
+//
+// - License: A signed JWT that lists the features that are allowed to be used by
+//            a given deployment. A license can have extra properties like,
+//            `IsTrial`, `DeploymentIDs`, etc that can be used to further define
+//            usage of the license.
+/**/
+// - Entitlements: A parsed set of licenses. Yes you can have more than 1 license
+//                 on a deployment! Entitlements will enumerate all features that
+//                 are allowed to be used.
+//
+package license

enterprise/coderd/license/license.go

@@ -10,8 +10,6 @@ import (
 	"github.com/golang-jwt/jwt/v4"
 	"golang.org/x/xerrors"
 
-	"cdr.dev/slog"
-
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbauthz"
 	"github.com/coder/coder/v2/codersdk"
@@ -21,58 +19,103 @@ import (
 func Entitlements(
 	ctx context.Context,
 	db database.Store,
-	logger slog.Logger,
 	replicaCount int,
 	externalAuthCount int,
 	keys map[string]ed25519.PublicKey,
 	enablements map[codersdk.FeatureName]bool,
 ) (codersdk.Entitlements, error) {
 	now := time.Now()
-	// Default all entitlements to be disabled.
-	entitlements := codersdk.Entitlements{
-		Features: map[codersdk.FeatureName]codersdk.Feature{},
-		Warnings: []string{},
-		Errors:   []string{},
-	}
-	for _, featureName := range codersdk.FeatureNames {
-		entitlements.Features[featureName] = codersdk.Feature{
-			Entitlement: codersdk.EntitlementNotEntitled,
-			Enabled:     enablements[featureName],
-		}
-	}
 
 	// nolint:gocritic // Getting unexpired licenses is a system function.
 	licenses, err := db.GetUnexpiredLicenses(dbauthz.AsSystemRestricted(ctx))
 	if err != nil {
-		return entitlements, err
+		return codersdk.Entitlements{}, err
 	}
 
 	// nolint:gocritic // Getting active user count is a system function.
 	activeUserCount, err := db.GetActiveUserCount(dbauthz.AsSystemRestricted(ctx))
 	if err != nil {
-		return entitlements, xerrors.Errorf("query active user count: %w", err)
+		return codersdk.Entitlements{}, xerrors.Errorf("query active user count: %w", err)
 	}
 
 	// always shows active user count regardless of license
-	entitlements.Features[codersdk.FeatureUserLimit] = codersdk.Feature{
-		Entitlement: codersdk.EntitlementNotEntitled,
-		Enabled:     enablements[codersdk.FeatureUserLimit],
-		Actual:      &activeUserCount,
+	entitlements, err := LicensesEntitlements(now, licenses, enablements, keys, FeatureArguments{
+		ActiveUserCount:   activeUserCount,
+		ReplicaCount:      replicaCount,
+		ExternalAuthCount: externalAuthCount,
+	})
+	if err != nil {
+		return entitlements, err
 	}
 
-	allFeatures := false
-	allFeaturesEntitlement := codersdk.EntitlementNotEntitled
+	return entitlements, nil
+}
 
-	// Here we loop through licenses to detect enabled features.
-	for _, l := range licenses {
-		claims, err := ParseClaims(l.JWT, keys)
+type FeatureArguments struct {
+	ActiveUserCount   int64
+	ReplicaCount      int
+	ExternalAuthCount int
+}
+
+// LicensesEntitlements returns the entitlements for licenses. Entitlements are
+// merged from all licenses and the highest entitlement is used for each feature.
+// Arguments:
+//
+//	now: The time to use for checking license expiration.
+//	license: The license to check.
+//	enablements: Features can be explicitly disabled by the deployment even if
+//	             the license has the feature entitled. Features can also have
+//	             the 'feat.AlwaysEnable()' return true to disallow disabling.
+//	featureArguments: Additional arguments required by specific features.
+func LicensesEntitlements(
+	now time.Time,
+	licenses []database.License,
+	enablements map[codersdk.FeatureName]bool,
+	keys map[string]ed25519.PublicKey,
+	featureArguments FeatureArguments,
+) (codersdk.Entitlements, error) {
+	// Default all entitlements to be disabled.
+	entitlements := codersdk.Entitlements{
+		Features: map[codersdk.FeatureName]codersdk.Feature{
+			// always shows active user count regardless of license.
+			codersdk.FeatureUserLimit: {
+				Entitlement: codersdk.EntitlementNotEntitled,
+				Enabled:     enablements[codersdk.FeatureUserLimit],
+				Actual:      &featureArguments.ActiveUserCount,
+			},
+		},
+		Warnings: []string{},
+		Errors:   []string{},
+	}
+
+	// By default, enumerate all features and set them to not entitled.
+	for _, featureName := range codersdk.FeatureNames {
+		entitlements.AddFeature(featureName, codersdk.Feature{
+			Entitlement: codersdk.EntitlementNotEntitled,
+			Enabled:     enablements[featureName],
+		})
+	}
+
+	// TODO: License specific warnings and errors should be tied to the license, not the
+	//   'Entitlements' group as a whole.
+	for _, license := range licenses {
+		claims, err := ParseClaims(license.JWT, keys)
 		if err != nil {
-			logger.Debug(ctx, "skipping invalid license",
-				slog.F("id", l.ID), slog.Error(err))
+			entitlements.Errors = append(entitlements.Errors,
+				fmt.Sprintf("Invalid license (%s) parsing claims: %s", license.UUID.String(), err.Error()))
 			continue
 		}
+
+		// Any valid license should toggle this boolean
 		entitlements.HasLicense = true
+
+		// If any license requires telemetry, the deployment should require telemetry.
+		entitlements.RequireTelemetry = entitlements.RequireTelemetry || claims.RequireTelemetry
+
+		// entitlement is the highest entitlement for any features in this license.
 		entitlement := codersdk.EntitlementEntitled
+		// If any license is a trial license, this should be set to true.
+		// The user should delete the trial license to remove this.
 		entitlements.Trial = claims.Trial
 		if now.After(claims.LicenseExpires.Time) {
 			// if the grace period were over, the validation fails, so if we are after
@@ -80,22 +123,32 @@ func Entitlements(
 			entitlement = codersdk.EntitlementGracePeriod
 		}
 
-		// Add warning if license is expiring soon
-		daysToExpire := int(math.Ceil(claims.LicenseExpires.Sub(now).Hours() / 24))
-		isTrial := entitlements.Trial
-		showWarningDays := 30
-		if isTrial {
-			showWarningDays = 7
-		}
-		isExpiringSoon := daysToExpire > 0 && daysToExpire < showWarningDays
-		if isExpiringSoon {
-			day := "day"
-			if daysToExpire > 1 {
-				day = "days"
-			}
-			entitlements.Warnings = append(entitlements.Warnings, fmt.Sprintf("Your license expires in %d %s.", daysToExpire, day))
+		// Will add a warning if the license is expiring soon.
+		// This warning can be raised multiple times if there is more than 1 license.
+		licenseExpirationWarning(&entitlements, now, claims)
+
+		// 'claims.AllFeature' is the legacy way to set 'claims.FeatureSet = codersdk.FeatureSetEnterprise'
+		// If both are set, ignore the legacy 'claims.AllFeature'
+		if claims.AllFeatures && claims.FeatureSet == "" {
+			claims.FeatureSet = codersdk.FeatureSetEnterprise
 		}
 
+		// Add all features from the feature set defined.
+		for _, featureName := range claims.FeatureSet.Features() {
+			if featureName == codersdk.FeatureUserLimit {
+				// FeatureUserLimit is unique in that it must be specifically defined
+				// in the license. There is no default meaning if no "limit" is set.
+				continue
+			}
+			entitlements.AddFeature(featureName, codersdk.Feature{
+				Entitlement: entitlement,
+				Enabled:     enablements[featureName] || featureName.AlwaysEnable(),
+				Limit:       nil,
+				Actual:      nil,
+			})
+		}
+
+		// Features al-la-carte
 		for featureName, featureValue := range claims.Features {
 			// Can this be negative?
 			if featureValue <= 0 {
@@ -103,86 +156,28 @@ func Entitlements(
 			}
 
 			switch featureName {
-			// User limit has special treatment as our only non-boolean feature.
 			case codersdk.FeatureUserLimit:
+				// User limit has special treatment as our only non-boolean feature.
 				limit := featureValue
-				priorLimit := entitlements.Features[codersdk.FeatureUserLimit]
-				if priorLimit.Limit != nil && *priorLimit.Limit > limit {
-					limit = *priorLimit.Limit
-				}
-				entitlements.Features[codersdk.FeatureUserLimit] = codersdk.Feature{
+				entitlements.AddFeature(codersdk.FeatureUserLimit, codersdk.Feature{
 					Enabled:     true,
 					Entitlement: entitlement,
 					Limit:       &limit,
-					Actual:      &activeUserCount,
-				}
+					Actual:      &featureArguments.ActiveUserCount,
+				})
 			default:
 				entitlements.Features[featureName] = codersdk.Feature{
-					Entitlement: maxEntitlement(entitlements.Features[featureName].Entitlement, entitlement),
+					Entitlement: entitlement,
 					Enabled:     enablements[featureName] || featureName.AlwaysEnable(),
 				}
 			}
 		}
-
-		if claims.AllFeatures {
-			allFeatures = true
-			allFeaturesEntitlement = maxEntitlement(allFeaturesEntitlement, entitlement)
-		}
-		entitlements.RequireTelemetry = entitlements.RequireTelemetry || claims.RequireTelemetry
 	}
 
-	if allFeatures {
-		for _, featureName := range codersdk.FeatureNames {
-			// No user limit!
-			if featureName == codersdk.FeatureUserLimit {
-				continue
-			}
-			feature := entitlements.Features[featureName]
-			feature.Entitlement = maxEntitlement(feature.Entitlement, allFeaturesEntitlement)
-			feature.Enabled = enablements[featureName] || featureName.AlwaysEnable()
-			entitlements.Features[featureName] = feature
-		}
-	}
+	// Now the license specific warnings and errors are added to the entitlements.
 
-	if entitlements.HasLicense {
-		userLimit := entitlements.Features[codersdk.FeatureUserLimit].Limit
-		if userLimit != nil && activeUserCount > *userLimit {
-			entitlements.Warnings = append(entitlements.Warnings, fmt.Sprintf(
-				"Your deployment has %d active users but is only licensed for %d.",
-				activeUserCount, *userLimit))
-		}
-
-		for _, featureName := range codersdk.FeatureNames {
-			// The user limit has it's own warnings!
-			if featureName == codersdk.FeatureUserLimit {
-				continue
-			}
-			// High availability has it's own warnings based on replica count!
-			if featureName == codersdk.FeatureHighAvailability {
-				continue
-			}
-			// External Auth Providers auth has it's own warnings based on the number configured!
-			if featureName == codersdk.FeatureMultipleExternalAuth {
-				continue
-			}
-			feature := entitlements.Features[featureName]
-			if !feature.Enabled {
-				continue
-			}
-			niceName := featureName.Humanize()
-			switch feature.Entitlement {
-			case codersdk.EntitlementNotEntitled:
-				entitlements.Warnings = append(entitlements.Warnings,
-					fmt.Sprintf("%s is enabled but your license is not entitled to this feature.", niceName))
-			case codersdk.EntitlementGracePeriod:
-				entitlements.Warnings = append(entitlements.Warnings,
-					fmt.Sprintf("%s is enabled but your license for this feature is expired.", niceName))
-			default:
-			}
-		}
-	}
-
-	if replicaCount > 1 {
+	// If HA is enabled, ensure the feature is entitled.
+	if featureArguments.ReplicaCount > 1 {
 		feature := entitlements.Features[codersdk.FeatureHighAvailability]
 
 		switch feature.Entitlement {
@@ -200,7 +195,7 @@ func Entitlements(
 		}
 	}
 
-	if externalAuthCount > 1 {
+	if featureArguments.ExternalAuthCount > 1 {
 		feature := entitlements.Features[codersdk.FeatureMultipleExternalAuth]
 
 		switch feature.Entitlement {
@@ -221,6 +216,52 @@ func Entitlements(
 		}
 	}
 
+	if entitlements.HasLicense {
+		userLimit := entitlements.Features[codersdk.FeatureUserLimit]
+		if userLimit.Limit != nil && featureArguments.ActiveUserCount > *userLimit.Limit {
+			entitlements.Warnings = append(entitlements.Warnings, fmt.Sprintf(
+				"Your deployment has %d active users but is only licensed for %d.",
+				featureArguments.ActiveUserCount, *userLimit.Limit))
+		} else if userLimit.Limit != nil && userLimit.Entitlement == codersdk.EntitlementGracePeriod {
+			entitlements.Warnings = append(entitlements.Warnings, fmt.Sprintf(
+				"Your deployment has %d active users but the license with the limit %d is expired.",
+				featureArguments.ActiveUserCount, *userLimit.Limit))
+		}
+
+		// Add a warning for every feature that is enabled but not entitled or
+		// is in a grace period.
+		for _, featureName := range codersdk.FeatureNames {
+			// The user limit has it's own warnings!
+			if featureName == codersdk.FeatureUserLimit {
+				continue
+			}
+			// High availability has it's own warnings based on replica count!
+			if featureName == codersdk.FeatureHighAvailability {
+				continue
+			}
+			// External Auth Providers auth has it's own warnings based on the number configured!
+			if featureName == codersdk.FeatureMultipleExternalAuth {
+				continue
+			}
+
+			feature := entitlements.Features[featureName]
+			if !feature.Enabled {
+				continue
+			}
+			niceName := featureName.Humanize()
+			switch feature.Entitlement {
+			case codersdk.EntitlementNotEntitled:
+				entitlements.Warnings = append(entitlements.Warnings,
+					fmt.Sprintf("%s is enabled but your license is not entitled to this feature.", niceName))
+			case codersdk.EntitlementGracePeriod:
+				entitlements.Warnings = append(entitlements.Warnings,
+					fmt.Sprintf("%s is enabled but your license for this feature is expired.", niceName))
+			default:
+			}
+		}
+	}
+
+	// Wrap up by disabling all features that are not entitled.
 	for _, featureName := range codersdk.FeatureNames {
 		feature := entitlements.Features[featureName]
 		if feature.Entitlement == codersdk.EntitlementNotEntitled {
@@ -261,9 +302,12 @@ type Claims struct {
 	AccountType    string           `json:"account_type,omitempty"`
 	AccountID      string           `json:"account_id,omitempty"`
 	// DeploymentIDs enforces the license can only be used on a set of deployments.
-	DeploymentIDs    []string `json:"deployment_ids,omitempty"`
-	Trial            bool     `json:"trial"`
-	AllFeatures      bool     `json:"all_features"`
+	DeploymentIDs []string            `json:"deployment_ids,omitempty"`
+	Trial         bool                `json:"trial"`
+	FeatureSet    codersdk.FeatureSet `json:"feature_set"`
+	// AllFeatures represents 'FeatureSet = FeatureSetEnterprise'
+	// Deprecated: AllFeatures is deprecated in favor of FeatureSet.
+	AllFeatures      bool     `json:"all_features,omitempty"`
 	Version          uint64   `json:"version"`
 	Features         Features `json:"features"`
 	RequireTelemetry bool     `json:"require_telemetry,omitempty"`
@@ -330,13 +374,21 @@ func keyFunc(keys map[string]ed25519.PublicKey) func(*jwt.Token) (interface{}, e
 	}
 }
 
-// maxEntitlement is the "greater" entitlement between the given values
-func maxEntitlement(e1, e2 codersdk.Entitlement) codersdk.Entitlement {
-	if e1 == codersdk.EntitlementEntitled || e2 == codersdk.EntitlementEntitled {
-		return codersdk.EntitlementEntitled
+// licenseExpirationWarning adds a warning message if the license is expiring soon.
+func licenseExpirationWarning(entitlements *codersdk.Entitlements, now time.Time, claims *Claims) {
+	// Add warning if license is expiring soon
+	daysToExpire := int(math.Ceil(claims.LicenseExpires.Sub(now).Hours() / 24))
+	showWarningDays := 30
+	isTrial := entitlements.Trial
+	if isTrial {
+		showWarningDays = 7
 	}
-	if e1 == codersdk.EntitlementGracePeriod || e2 == codersdk.EntitlementGracePeriod {
-		return codersdk.EntitlementGracePeriod
+	isExpiringSoon := daysToExpire > 0 && daysToExpire < showWarningDays
+	if isExpiringSoon {
+		day := "day"
+		if daysToExpire > 1 {
+			day = "days"
+		}
+		entitlements.Warnings = append(entitlements.Warnings, fmt.Sprintf("Your license expires in %d %s.", daysToExpire, day))
 	}
-	return codersdk.EntitlementNotEntitled
 }

enterprise/coderd/license/license_test.go

@@ -7,9 +7,10 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/exp/slices"
 
-	"cdr.dev/slog"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbmem"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
@@ -30,7 +31,7 @@ func TestEntitlements(t *testing.T) {
 	t.Run("Defaults", func(t *testing.T) {
 		t.Parallel()
 		db := dbmem.New()
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, all)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, all)
 		require.NoError(t, err)
 		require.False(t, entitlements.HasLicense)
 		require.False(t, entitlements.Trial)
@@ -42,7 +43,7 @@ func TestEntitlements(t *testing.T) {
 	t.Run("Always return the current user count", func(t *testing.T) {
 		t.Parallel()
 		db := dbmem.New()
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, all)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, all)
 		require.NoError(t, err)
 		require.False(t, entitlements.HasLicense)
 		require.False(t, entitlements.Trial)
@@ -55,7 +56,7 @@ func TestEntitlements(t *testing.T) {
 			JWT: coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{}),
 			Exp: time.Now().Add(time.Hour),
 		})
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, empty)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, empty)
 		require.NoError(t, err)
 		require.True(t, entitlements.HasLicense)
 		require.False(t, entitlements.Trial)
@@ -79,7 +80,7 @@ func TestEntitlements(t *testing.T) {
 			}),
 			Exp: time.Now().Add(time.Hour),
 		})
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, empty)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, empty)
 		require.NoError(t, err)
 		require.True(t, entitlements.HasLicense)
 		require.False(t, entitlements.Trial)
@@ -102,7 +103,7 @@ func TestEntitlements(t *testing.T) {
 			}),
 			Exp: time.Now().Add(time.Hour),
 		})
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, all)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, all)
 		require.NoError(t, err)
 		require.True(t, entitlements.HasLicense)
 		require.False(t, entitlements.Trial)
@@ -129,7 +130,7 @@ func TestEntitlements(t *testing.T) {
 			Exp: time.Now().AddDate(0, 0, 5),
 		})
 
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, all)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, all)
 
 		require.NoError(t, err)
 		require.True(t, entitlements.HasLicense)
@@ -158,7 +159,7 @@ func TestEntitlements(t *testing.T) {
 			Exp: time.Now().AddDate(0, 0, 5),
 		})
 
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, all)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, all)
 
 		require.NoError(t, err)
 		require.True(t, entitlements.HasLicense)
@@ -188,7 +189,7 @@ func TestEntitlements(t *testing.T) {
 			Exp: time.Now().AddDate(0, 0, 5),
 		})
 
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, all)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, all)
 
 		require.NoError(t, err)
 		require.True(t, entitlements.HasLicense)
@@ -217,7 +218,7 @@ func TestEntitlements(t *testing.T) {
 			Exp: time.Now().AddDate(0, 0, 5),
 		})
 
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, all)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, all)
 
 		require.NoError(t, err)
 		require.True(t, entitlements.HasLicense)
@@ -237,7 +238,7 @@ func TestEntitlements(t *testing.T) {
 			JWT: coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{}),
 			Exp: time.Now().Add(time.Hour),
 		})
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, all)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, all)
 		require.NoError(t, err)
 		require.True(t, entitlements.HasLicense)
 		require.False(t, entitlements.Trial)
@@ -299,7 +300,7 @@ func TestEntitlements(t *testing.T) {
 			}),
 			Exp: time.Now().Add(time.Hour),
 		})
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, empty)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, empty)
 		require.NoError(t, err)
 		require.True(t, entitlements.HasLicense)
 		require.Contains(t, entitlements.Warnings, "Your deployment has 2 active users but is only licensed for 1.")
@@ -327,7 +328,7 @@ func TestEntitlements(t *testing.T) {
 			}),
 			Exp: time.Now().Add(60 * 24 * time.Hour),
 		})
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, empty)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, empty)
 		require.NoError(t, err)
 		require.True(t, entitlements.HasLicense)
 		require.Empty(t, entitlements.Warnings)
@@ -350,12 +351,96 @@ func TestEntitlements(t *testing.T) {
 			}),
 		})
 
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, empty)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, empty)
 		require.NoError(t, err)
 		require.True(t, entitlements.HasLicense)
 		require.False(t, entitlements.Trial)
 	})
 
+	t.Run("Enterprise", func(t *testing.T) {
+		t.Parallel()
+		db := dbmem.New()
+		_, err := db.InsertLicense(context.Background(), database.InsertLicenseParams{
+			Exp: time.Now().Add(time.Hour),
+			JWT: coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
+				FeatureSet: codersdk.FeatureSetEnterprise,
+			}),
+		})
+		require.NoError(t, err)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, all)
+		require.NoError(t, err)
+		require.True(t, entitlements.HasLicense)
+		require.False(t, entitlements.Trial)
+
+		// All enterprise features should be entitled
+		enterpriseFeatures := codersdk.FeatureSetEnterprise.Features()
+		for _, featureName := range codersdk.FeatureNames {
+			if featureName == codersdk.FeatureUserLimit {
+				continue
+			}
+			if slices.Contains(enterpriseFeatures, featureName) {
+				require.True(t, entitlements.Features[featureName].Enabled, featureName)
+				require.Equal(t, codersdk.EntitlementEntitled, entitlements.Features[featureName].Entitlement)
+			} else {
+				require.False(t, entitlements.Features[featureName].Enabled, featureName)
+				require.Equal(t, codersdk.EntitlementNotEntitled, entitlements.Features[featureName].Entitlement)
+			}
+		}
+	})
+
+	t.Run("Premium", func(t *testing.T) {
+		t.Parallel()
+		db := dbmem.New()
+		_, err := db.InsertLicense(context.Background(), database.InsertLicenseParams{
+			Exp: time.Now().Add(time.Hour),
+			JWT: coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
+				FeatureSet: codersdk.FeatureSetPremium,
+			}),
+		})
+		require.NoError(t, err)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, all)
+		require.NoError(t, err)
+		require.True(t, entitlements.HasLicense)
+		require.False(t, entitlements.Trial)
+
+		// All premium features should be entitled
+		enterpriseFeatures := codersdk.FeatureSetPremium.Features()
+		for _, featureName := range codersdk.FeatureNames {
+			if featureName == codersdk.FeatureUserLimit {
+				continue
+			}
+			if slices.Contains(enterpriseFeatures, featureName) {
+				require.True(t, entitlements.Features[featureName].Enabled, featureName)
+				require.Equal(t, codersdk.EntitlementEntitled, entitlements.Features[featureName].Entitlement)
+			} else {
+				require.False(t, entitlements.Features[featureName].Enabled, featureName)
+				require.Equal(t, codersdk.EntitlementNotEntitled, entitlements.Features[featureName].Entitlement)
+			}
+		}
+	})
+
+	t.Run("SetNone", func(t *testing.T) {
+		t.Parallel()
+		db := dbmem.New()
+		_, err := db.InsertLicense(context.Background(), database.InsertLicenseParams{
+			Exp: time.Now().Add(time.Hour),
+			JWT: coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
+				FeatureSet: "",
+			}),
+		})
+		require.NoError(t, err)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, all)
+		require.NoError(t, err)
+		require.True(t, entitlements.HasLicense)
+		require.False(t, entitlements.Trial)
+
+		for _, featureName := range codersdk.FeatureNames {
+			require.False(t, entitlements.Features[featureName].Enabled, featureName)
+			require.Equal(t, codersdk.EntitlementNotEntitled, entitlements.Features[featureName].Entitlement)
+		}
+	})
+
+	// AllFeatures uses the deprecated 'AllFeatures' boolean.
 	t.Run("AllFeatures", func(t *testing.T) {
 		t.Parallel()
 		db := dbmem.New()
@@ -365,16 +450,24 @@ func TestEntitlements(t *testing.T) {
 				AllFeatures: true,
 			}),
 		})
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, all)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, all)
 		require.NoError(t, err)
 		require.True(t, entitlements.HasLicense)
 		require.False(t, entitlements.Trial)
+
+		// All enterprise features should be entitled
+		enterpriseFeatures := codersdk.FeatureSetEnterprise.Features()
 		for _, featureName := range codersdk.FeatureNames {
 			if featureName == codersdk.FeatureUserLimit {
 				continue
 			}
-			require.True(t, entitlements.Features[featureName].Enabled)
-			require.Equal(t, codersdk.EntitlementEntitled, entitlements.Features[featureName].Entitlement)
+			if slices.Contains(enterpriseFeatures, featureName) {
+				require.True(t, entitlements.Features[featureName].Enabled, featureName)
+				require.Equal(t, codersdk.EntitlementEntitled, entitlements.Features[featureName].Entitlement)
+			} else {
+				require.False(t, entitlements.Features[featureName].Enabled, featureName)
+				require.Equal(t, codersdk.EntitlementNotEntitled, entitlements.Features[featureName].Entitlement)
+			}
 		}
 	})
 
@@ -387,17 +480,25 @@ func TestEntitlements(t *testing.T) {
 				AllFeatures: true,
 			}),
 		})
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, empty)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, empty)
 		require.NoError(t, err)
 		require.True(t, entitlements.HasLicense)
 		require.False(t, entitlements.Trial)
+		// All enterprise features should be entitled
+		enterpriseFeatures := codersdk.FeatureSetEnterprise.Features()
 		for _, featureName := range codersdk.FeatureNames {
 			if featureName == codersdk.FeatureUserLimit {
 				continue
 			}
+
 			feature := entitlements.Features[featureName]
-			require.Equal(t, featureName.AlwaysEnable(), feature.Enabled)
-			require.Equal(t, codersdk.EntitlementEntitled, feature.Entitlement)
+			if slices.Contains(enterpriseFeatures, featureName) {
+				require.Equal(t, featureName.AlwaysEnable(), feature.Enabled)
+				require.Equal(t, codersdk.EntitlementEntitled, feature.Entitlement)
+			} else {
+				require.False(t, entitlements.Features[featureName].Enabled, featureName)
+				require.Equal(t, codersdk.EntitlementNotEntitled, entitlements.Features[featureName].Entitlement)
+			}
 		}
 	})
 
@@ -412,23 +513,30 @@ func TestEntitlements(t *testing.T) {
 				ExpiresAt:   dbtime.Now().Add(time.Hour),
 			}),
 		})
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, all)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 1, coderdenttest.Keys, all)
 		require.NoError(t, err)
 		require.True(t, entitlements.HasLicense)
 		require.False(t, entitlements.Trial)
+		// All enterprise features should be entitled
+		enterpriseFeatures := codersdk.FeatureSetEnterprise.Features()
 		for _, featureName := range codersdk.FeatureNames {
 			if featureName == codersdk.FeatureUserLimit {
 				continue
 			}
-			require.True(t, entitlements.Features[featureName].Enabled)
-			require.Equal(t, codersdk.EntitlementGracePeriod, entitlements.Features[featureName].Entitlement)
+			if slices.Contains(enterpriseFeatures, featureName) {
+				require.True(t, entitlements.Features[featureName].Enabled, featureName)
+				require.Equal(t, codersdk.EntitlementGracePeriod, entitlements.Features[featureName].Entitlement)
+			} else {
+				require.False(t, entitlements.Features[featureName].Enabled, featureName)
+				require.Equal(t, codersdk.EntitlementNotEntitled, entitlements.Features[featureName].Entitlement)
+			}
 		}
 	})
 
 	t.Run("MultipleReplicasNoLicense", func(t *testing.T) {
 		t.Parallel()
 		db := dbmem.New()
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 2, 1, coderdenttest.Keys, all)
+		entitlements, err := license.Entitlements(context.Background(), db, 2, 1, coderdenttest.Keys, all)
 		require.NoError(t, err)
 		require.False(t, entitlements.HasLicense)
 		require.Len(t, entitlements.Errors, 1)
@@ -446,7 +554,7 @@ func TestEntitlements(t *testing.T) {
 				},
 			}),
 		})
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 2, 1, coderdenttest.Keys, map[codersdk.FeatureName]bool{
+		entitlements, err := license.Entitlements(context.Background(), db, 2, 1, coderdenttest.Keys, map[codersdk.FeatureName]bool{
 			codersdk.FeatureHighAvailability: true,
 		})
 		require.NoError(t, err)
@@ -468,7 +576,7 @@ func TestEntitlements(t *testing.T) {
 			}),
 			Exp: time.Now().Add(time.Hour),
 		})
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 2, 1, coderdenttest.Keys, map[codersdk.FeatureName]bool{
+		entitlements, err := license.Entitlements(context.Background(), db, 2, 1, coderdenttest.Keys, map[codersdk.FeatureName]bool{
 			codersdk.FeatureHighAvailability: true,
 		})
 		require.NoError(t, err)
@@ -480,7 +588,7 @@ func TestEntitlements(t *testing.T) {
 	t.Run("MultipleGitAuthNoLicense", func(t *testing.T) {
 		t.Parallel()
 		db := dbmem.New()
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 2, coderdenttest.Keys, all)
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 2, coderdenttest.Keys, all)
 		require.NoError(t, err)
 		require.False(t, entitlements.HasLicense)
 		require.Len(t, entitlements.Errors, 1)
@@ -498,7 +606,7 @@ func TestEntitlements(t *testing.T) {
 				},
 			}),
 		})
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 2, coderdenttest.Keys, map[codersdk.FeatureName]bool{
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 2, coderdenttest.Keys, map[codersdk.FeatureName]bool{
 			codersdk.FeatureMultipleExternalAuth: true,
 		})
 		require.NoError(t, err)
@@ -520,7 +628,7 @@ func TestEntitlements(t *testing.T) {
 			}),
 			Exp: time.Now().Add(time.Hour),
 		})
-		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 2, coderdenttest.Keys, map[codersdk.FeatureName]bool{
+		entitlements, err := license.Entitlements(context.Background(), db, 1, 2, coderdenttest.Keys, map[codersdk.FeatureName]bool{
 			codersdk.FeatureMultipleExternalAuth: true,
 		})
 		require.NoError(t, err)
@@ -529,3 +637,236 @@ func TestEntitlements(t *testing.T) {
 		require.Equal(t, "You have multiple External Auth Providers configured but your license is expired. Reduce to one.", entitlements.Warnings[0])
 	})
 }
+
+func TestLicenseEntitlements(t *testing.T) {
+	t.Parallel()
+
+	// We must use actual 'time.Now()' in tests because the jwt library does
+	// not accept a custom time function. The only way to change it is as a
+	// package global, which does not work in t.Parallel().
+
+	// This list comes from coderd.go on launch. This list is a bit arbitrary,
+	// maybe some should be moved to "AlwaysEnabled" instead.
+	defaultEnablements := map[codersdk.FeatureName]bool{
+		codersdk.FeatureAuditLog:                   true,
+		codersdk.FeatureBrowserOnly:                true,
+		codersdk.FeatureSCIM:                       true,
+		codersdk.FeatureMultipleExternalAuth:       true,
+		codersdk.FeatureTemplateRBAC:               true,
+		codersdk.FeatureExternalTokenEncryption:    true,
+		codersdk.FeatureExternalProvisionerDaemons: true,
+		codersdk.FeatureAdvancedTemplateScheduling: true,
+		codersdk.FeatureWorkspaceProxy:             true,
+		codersdk.FeatureUserRoleManagement:         true,
+		codersdk.FeatureAccessControl:              true,
+		codersdk.FeatureControlSharedPorts:         true,
+	}
+
+	legacyLicense := func() *coderdenttest.LicenseOptions {
+		return (&coderdenttest.LicenseOptions{
+			AccountType: "salesforce",
+			AccountID:   "Alice",
+			Trial:       false,
+			// Use the legacy boolean
+			AllFeatures: true,
+		}).Valid(time.Now())
+	}
+
+	enterpriseLicense := func() *coderdenttest.LicenseOptions {
+		return (&coderdenttest.LicenseOptions{
+			AccountType:   "salesforce",
+			AccountID:     "Bob",
+			DeploymentIDs: nil,
+			Trial:         false,
+			FeatureSet:    codersdk.FeatureSetEnterprise,
+			AllFeatures:   true,
+		}).Valid(time.Now())
+	}
+
+	premiumLicense := func() *coderdenttest.LicenseOptions {
+		return (&coderdenttest.LicenseOptions{
+			AccountType:   "salesforce",
+			AccountID:     "Charlie",
+			DeploymentIDs: nil,
+			Trial:         false,
+			FeatureSet:    codersdk.FeatureSetPremium,
+			AllFeatures:   true,
+		}).Valid(time.Now())
+	}
+
+	testCases := []struct {
+		Name        string
+		Licenses    []*coderdenttest.LicenseOptions
+		Enablements map[codersdk.FeatureName]bool
+		Arguments   license.FeatureArguments
+
+		ExpectedErrorContains string
+		AssertEntitlements    func(t *testing.T, entitlements codersdk.Entitlements)
+	}{
+		{
+			Name: "NoLicenses",
+			AssertEntitlements: func(t *testing.T, entitlements codersdk.Entitlements) {
+				assertNoErrors(t, entitlements)
+				assertNoWarnings(t, entitlements)
+				assert.False(t, entitlements.HasLicense)
+				assert.False(t, entitlements.Trial)
+			},
+		},
+		{
+			Name: "MixedUsedCounts",
+			Licenses: []*coderdenttest.LicenseOptions{
+				legacyLicense().UserLimit(100),
+				enterpriseLicense().UserLimit(500),
+			},
+			Enablements: defaultEnablements,
+			Arguments: license.FeatureArguments{
+				ActiveUserCount:   50,
+				ReplicaCount:      0,
+				ExternalAuthCount: 0,
+			},
+			AssertEntitlements: func(t *testing.T, entitlements codersdk.Entitlements) {
+				assertEnterpriseFeatures(t, entitlements)
+				assertNoErrors(t, entitlements)
+				assertNoWarnings(t, entitlements)
+				userFeature := entitlements.Features[codersdk.FeatureUserLimit]
+				assert.Equalf(t, int64(500), *userFeature.Limit, "user limit")
+				assert.Equalf(t, int64(50), *userFeature.Actual, "user count")
+			},
+		},
+		{
+			Name: "MixedUsedCountsWithExpired",
+			Licenses: []*coderdenttest.LicenseOptions{
+				// This license is ignored
+				enterpriseLicense().UserLimit(500).Expired(time.Now()),
+				enterpriseLicense().UserLimit(100),
+			},
+			Enablements: defaultEnablements,
+			Arguments: license.FeatureArguments{
+				ActiveUserCount:   200,
+				ReplicaCount:      0,
+				ExternalAuthCount: 0,
+			},
+			AssertEntitlements: func(t *testing.T, entitlements codersdk.Entitlements) {
+				assertEnterpriseFeatures(t, entitlements)
+				userFeature := entitlements.Features[codersdk.FeatureUserLimit]
+				assert.Equalf(t, int64(100), *userFeature.Limit, "user limit")
+				assert.Equalf(t, int64(200), *userFeature.Actual, "user count")
+
+				require.Len(t, entitlements.Errors, 1, "invalid license error")
+				require.Len(t, entitlements.Warnings, 1, "user count exceeds warning")
+				require.Contains(t, entitlements.Errors[0], "Invalid license")
+				require.Contains(t, entitlements.Warnings[0], "active users but is only licensed for")
+			},
+		},
+		{
+			// The new license does not have enough seats to cover the active user count.
+			// The old license is in it's grace period.
+			Name: "MixedUsedCountsWithGrace",
+			Licenses: []*coderdenttest.LicenseOptions{
+				enterpriseLicense().UserLimit(500).GracePeriod(time.Now()),
+				enterpriseLicense().UserLimit(100),
+			},
+			Enablements: defaultEnablements,
+			Arguments: license.FeatureArguments{
+				ActiveUserCount:   200,
+				ReplicaCount:      0,
+				ExternalAuthCount: 0,
+			},
+			AssertEntitlements: func(t *testing.T, entitlements codersdk.Entitlements) {
+				userFeature := entitlements.Features[codersdk.FeatureUserLimit]
+				assert.Equalf(t, int64(500), *userFeature.Limit, "user limit")
+				assert.Equalf(t, int64(200), *userFeature.Actual, "user count")
+				assert.Equal(t, userFeature.Entitlement, codersdk.EntitlementGracePeriod)
+			},
+		},
+		{
+			// Legacy license uses the "AllFeatures" boolean
+			Name: "LegacyLicense",
+			Licenses: []*coderdenttest.LicenseOptions{
+				legacyLicense().UserLimit(100),
+			},
+			Enablements: defaultEnablements,
+			Arguments: license.FeatureArguments{
+				ActiveUserCount:   50,
+				ReplicaCount:      0,
+				ExternalAuthCount: 0,
+			},
+			AssertEntitlements: func(t *testing.T, entitlements codersdk.Entitlements) {
+				assertEnterpriseFeatures(t, entitlements)
+				assertNoErrors(t, entitlements)
+				assertNoWarnings(t, entitlements)
+				userFeature := entitlements.Features[codersdk.FeatureUserLimit]
+				assert.Equalf(t, int64(100), *userFeature.Limit, "user limit")
+				assert.Equalf(t, int64(50), *userFeature.Actual, "user count")
+			},
+		},
+		{
+			Name: "EnterpriseDisabledMultiOrg",
+			Licenses: []*coderdenttest.LicenseOptions{
+				enterpriseLicense().UserLimit(100),
+			},
+			Enablements:           defaultEnablements,
+			Arguments:             license.FeatureArguments{},
+			ExpectedErrorContains: "",
+			AssertEntitlements: func(t *testing.T, entitlements codersdk.Entitlements) {
+				assert.False(t, entitlements.Features[codersdk.FeatureMultipleOrganizations].Enabled, "multi-org only enabled for premium")
+			},
+		},
+		{
+			Name: "PremiumEnabledMultiOrg",
+			Licenses: []*coderdenttest.LicenseOptions{
+				premiumLicense().UserLimit(100),
+			},
+			Enablements:           defaultEnablements,
+			Arguments:             license.FeatureArguments{},
+			ExpectedErrorContains: "",
+			AssertEntitlements: func(t *testing.T, entitlements codersdk.Entitlements) {
+				assert.True(t, entitlements.Features[codersdk.FeatureMultipleOrganizations].Enabled, "multi-org enabled for premium")
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+
+		t.Run(tc.Name, func(t *testing.T) {
+			t.Parallel()
+
+			generatedLicenses := make([]database.License, 0, len(tc.Licenses))
+			for i, lo := range tc.Licenses {
+				generatedLicenses = append(generatedLicenses, database.License{
+					ID:         int32(i),
+					UploadedAt: time.Now().Add(time.Hour * -1),
+					JWT:        lo.Generate(t),
+					Exp:        lo.GraceAt,
+					UUID:       uuid.New(),
+				})
+			}
+
+			entitlements, err := license.LicensesEntitlements(time.Now(), generatedLicenses, tc.Enablements, coderdenttest.Keys, tc.Arguments)
+			if tc.ExpectedErrorContains != "" {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), tc.ExpectedErrorContains)
+			} else {
+				require.NoError(t, err)
+				tc.AssertEntitlements(t, entitlements)
+			}
+		})
+	}
+}
+
+func assertNoErrors(t *testing.T, entitlements codersdk.Entitlements) {
+	assert.Empty(t, entitlements.Errors, "no errors")
+}
+
+func assertNoWarnings(t *testing.T, entitlements codersdk.Entitlements) {
+	assert.Empty(t, entitlements.Warnings, "no warnings")
+}
+
+func assertEnterpriseFeatures(t *testing.T, entitlements codersdk.Entitlements) {
+	for _, expected := range codersdk.FeatureSetEnterprise.Features() {
+		f := entitlements.Features[expected]
+		assert.Equalf(t, codersdk.EntitlementEntitled, f.Entitlement, "%s entitled", expected)
+		assert.Equalf(t, true, f.Enabled, "%s enabled", expected)
+	}
+}

site/src/api/typesGenerated.ts

@@ -2104,6 +2104,10 @@ export const FeatureNames: FeatureName[] = [
   "workspace_proxy",
 ];
 
+// From codersdk/deployment.go
+export type FeatureSet = "" | "enterprise" | "premium";
+export const FeatureSets: FeatureSet[] = ["", "enterprise", "premium"];
+
 // From codersdk/groups.go
 export type GroupSource = "oidc" | "user";
 export const GroupSources: GroupSource[] = ["oidc", "user"];