chore: refactor entitlements to keep it in just the options

Emyrk · Emyrk · commit ad4f901d1455 · 2024-08-26T10:32:20.000-05:00
Duplicating the reference did not feel valuable, just confusing
diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go
@@ -75,6 +75,9 @@ func New(ctx context.Context, options *Options) (_ *API, err error) {
 		// from when an additional replica was started.
 		options.ReplicaErrorGracePeriod = time.Minute
 	}
+	if options.Entitlements == nil {
+		options.Entitlements = entitlements.New()
+	}
 
 	ctx, cancelFunc := context.WithCancel(ctx)
 
@@ -105,25 +108,22 @@ func New(ctx context.Context, options *Options) (_ *API, err error) {
 		return nil, xerrors.Errorf("init database encryption: %w", err)
 	}
 
-	entitlementsSet := entitlements.New()
 	options.Database = cryptDB
 	api := &API{
-		ctx:          ctx,
-		cancel:       cancelFunc,
-		Options:      options,
-		entitlements: entitlementsSet,
+		ctx:     ctx,
+		cancel:  cancelFunc,
+		Options: options,
 		provisionerDaemonAuth: &provisionerDaemonAuth{
 			psk:        options.ProvisionerDaemonPSK,
 			authorizer: options.Authorizer,
 			db:         options.Database,
 		},
 		licenseMetricsCollector: &license.MetricsCollector{
-			Entitlements: entitlementsSet,
+			Entitlements: options.Entitlements,
 		},
 	}
 	// This must happen before coderd initialization!
 	options.PostAuthAdditionalHeadersFunc = api.writeEntitlementWarningsHeader
-	options.Options.Entitlements = api.entitlements
 	api.AGPL = coderd.New(options.Options)
 	defer func() {
 		if err != nil {
@@ -561,8 +561,6 @@ type API struct {
 	// ProxyHealth checks the reachability of all workspace proxies.
 	ProxyHealth *proxyhealth.ProxyHealth
 
-	entitlements *entitlements.Set
-
 	provisionerDaemonAuth *provisionerDaemonAuth
 
 	licenseMetricsCollector *license.MetricsCollector
@@ -595,7 +593,7 @@ func (api *API) writeEntitlementWarningsHeader(a rbac.Subject, header http.Heade
 		return
 	}
 
-	api.entitlements.WriteEntitlementWarningHeaders(header)
+	api.Entitlements.WriteEntitlementWarningHeaders(header)
 }
 
 func (api *API) Close() error {
@@ -658,7 +656,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 		//
 		// We don't simply append to entitlement.Errors since we don't want any
 		// enterprise features enabled.
-		api.entitlements.Update(func(entitlements *codersdk.Entitlements) {
+		api.Entitlements.Update(func(entitlements *codersdk.Entitlements) {
 			entitlements.Errors = []string{
 				"License requires telemetry but telemetry is disabled",
 			}
@@ -669,7 +667,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 	}
 
 	featureChanged := func(featureName codersdk.FeatureName) (initial, changed, enabled bool) {
-		return api.entitlements.FeatureChanged(featureName, reloadedEntitlements.Features[featureName])
+		return api.Entitlements.FeatureChanged(featureName, reloadedEntitlements.Features[featureName])
 	}
 
 	shouldUpdate := func(initial, changed, enabled bool) bool {
@@ -835,7 +833,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 	}
 	reloadedEntitlements.Features[codersdk.FeatureExternalTokenEncryption] = featureExternalTokenEncryption
 
-	api.entitlements.Replace(reloadedEntitlements)
+	api.Entitlements.Replace(reloadedEntitlements)
 	return nil
 }
 
@@ -1015,7 +1013,7 @@ func derpMapper(logger slog.Logger, proxyHealth *proxyhealth.ProxyHealth) func(*
 // @Router /entitlements [get]
 func (api *API) serveEntitlements(rw http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
-	httpapi.Write(ctx, rw, http.StatusOK, api.entitlements.AsJSON())
+	httpapi.Write(ctx, rw, http.StatusOK, api.Entitlements.AsJSON())
 }
 
 func (api *API) runEntitlementsLoop(ctx context.Context) {
diff --git a/enterprise/coderd/coderd_test.go b/enterprise/coderd/coderd_test.go
@@ -3,6 +3,7 @@ package coderd_test
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"net/http"
 	"reflect"
 	"strings"
@@ -46,14 +47,17 @@ func TestEntitlements(t *testing.T) {
 	t.Parallel()
 	t.Run("NoLicense", func(t *testing.T) {
 		t.Parallel()
-		adminClient, adminUser := coderdenttest.New(t, &coderdenttest.Options{
+		adminClient, _, api, adminUser := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
 			DontAddLicense: true,
 		})
 		anotherClient, _ := coderdtest.CreateAnotherUser(t, adminClient, adminUser.OrganizationID)
 		res, err := anotherClient.Entitlements(context.Background())
 		require.NoError(t, err)
 		require.False(t, res.HasLicense)
 		require.Empty(t, res.Warnings)
+
+		// Ensure the entitlements are the same reference
+		require.Equal(t, fmt.Sprintf("%p", api.Entitlements), fmt.Sprintf("%p", api.AGPL.Entitlements))
 	})
 	t.Run("FullLicense", func(t *testing.T) {
 		// PGCoordinator requires a real postgres
