feat: implement premium vs enterprise licenses (#13907)

* feat: implement premium vs enterprise licenses Implement different sets of licensed features.
2026-09-24 15:04:27 +08:00 · 2024-07-24 12:07:59 -05:00
parent 0d9615b4fd
commit 15fda232b7
8 changed files with 976 additions and 155 deletions
@@ -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 {
@@ -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")
+}