chore: refactor entitlements to be a safe object to use (#14406)

* chore: refactor entitlements to be passable as an argument

Previously, all usage of entitlements requires mutex usage on the
api struct directly. This prevents passing the entitlements to
a sub package. It also creates the possibility for misuse.

coderd/entitlements/entitlements.go

@@ -0,0 +1,109 @@
+package entitlements
+
+import (
+	"encoding/json"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/coder/coder/v2/codersdk"
+)
+
+type Set struct {
+	entitlementsMu sync.RWMutex
+	entitlements   codersdk.Entitlements
+}
+
+func New() *Set {
+	return &Set{
+		// Some defaults for an unlicensed instance.
+		// These will be updated when coderd is initialized.
+		entitlements: codersdk.Entitlements{
+			Features:         map[codersdk.FeatureName]codersdk.Feature{},
+			Warnings:         nil,
+			Errors:           nil,
+			HasLicense:       false,
+			Trial:            false,
+			RequireTelemetry: false,
+			RefreshedAt:      time.Time{},
+		},
+	}
+}
+
+// AllowRefresh returns whether the entitlements are allowed to be refreshed.
+// If it returns false, that means it was recently refreshed and the caller should
+// wait the returned duration before trying again.
+func (l *Set) AllowRefresh(now time.Time) (bool, time.Duration) {
+	l.entitlementsMu.RLock()
+	defer l.entitlementsMu.RUnlock()
+
+	diff := now.Sub(l.entitlements.RefreshedAt)
+	if diff < time.Minute {
+		return false, time.Minute - diff
+	}
+
+	return true, 0
+}
+
+func (l *Set) Feature(name codersdk.FeatureName) (codersdk.Feature, bool) {
+	l.entitlementsMu.RLock()
+	defer l.entitlementsMu.RUnlock()
+
+	f, ok := l.entitlements.Features[name]
+	return f, ok
+}
+
+func (l *Set) Enabled(feature codersdk.FeatureName) bool {
+	l.entitlementsMu.RLock()
+	defer l.entitlementsMu.RUnlock()
+
+	f, ok := l.entitlements.Features[feature]
+	if !ok {
+		return false
+	}
+	return f.Enabled
+}
+
+// AsJSON is used to return this to the api without exposing the entitlements for
+// mutation.
+func (l *Set) AsJSON() json.RawMessage {
+	l.entitlementsMu.RLock()
+	defer l.entitlementsMu.RUnlock()
+
+	b, _ := json.Marshal(l.entitlements)
+	return b
+}
+
+func (l *Set) Replace(entitlements codersdk.Entitlements) {
+	l.entitlementsMu.Lock()
+	defer l.entitlementsMu.Unlock()
+
+	l.entitlements = entitlements
+}
+
+func (l *Set) Update(do func(entitlements *codersdk.Entitlements)) {
+	l.entitlementsMu.Lock()
+	defer l.entitlementsMu.Unlock()
+
+	do(&l.entitlements)
+}
+
+func (l *Set) FeatureChanged(featureName codersdk.FeatureName, newFeature codersdk.Feature) (initial, changed, enabled bool) {
+	l.entitlementsMu.RLock()
+	defer l.entitlementsMu.RUnlock()
+
+	oldFeature := l.entitlements.Features[featureName]
+	if oldFeature.Enabled != newFeature.Enabled {
+		return false, true, newFeature.Enabled
+	}
+	return false, false, newFeature.Enabled
+}
+
+func (l *Set) WriteEntitlementWarningHeaders(header http.Header) {
+	l.entitlementsMu.RLock()
+	defer l.entitlementsMu.RUnlock()
+
+	for _, warning := range l.entitlements.Warnings {
+		header.Add(codersdk.EntitlementsWarningHeader, warning)
+	}
+}

coderd/entitlements/entitlements_test.go

@@ -0,0 +1,63 @@
+package entitlements_test
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/coderd/entitlements"
+	"github.com/coder/coder/v2/codersdk"
+)
+
+func TestUpdate(t *testing.T) {
+	t.Parallel()
+
+	set := entitlements.New()
+	require.False(t, set.Enabled(codersdk.FeatureMultipleOrganizations))
+
+	set.Update(func(entitlements *codersdk.Entitlements) {
+		entitlements.Features[codersdk.FeatureMultipleOrganizations] = codersdk.Feature{
+			Enabled:     true,
+			Entitlement: codersdk.EntitlementEntitled,
+		}
+	})
+	require.True(t, set.Enabled(codersdk.FeatureMultipleOrganizations))
+}
+
+func TestAllowRefresh(t *testing.T) {
+	t.Parallel()
+
+	now := time.Now()
+	set := entitlements.New()
+	set.Update(func(entitlements *codersdk.Entitlements) {
+		entitlements.RefreshedAt = now
+	})
+
+	ok, wait := set.AllowRefresh(now)
+	require.False(t, ok)
+	require.InDelta(t, time.Minute.Seconds(), wait.Seconds(), 5)
+
+	set.Update(func(entitlements *codersdk.Entitlements) {
+		entitlements.RefreshedAt = now.Add(time.Minute * -2)
+	})
+
+	ok, wait = set.AllowRefresh(now)
+	require.True(t, ok)
+	require.Equal(t, time.Duration(0), wait)
+}
+
+func TestReplace(t *testing.T) {
+	t.Parallel()
+
+	set := entitlements.New()
+	require.False(t, set.Enabled(codersdk.FeatureMultipleOrganizations))
+	set.Replace(codersdk.Entitlements{
+		Features: map[codersdk.FeatureName]codersdk.Feature{
+			codersdk.FeatureMultipleOrganizations: {
+				Enabled: true,
+			},
+		},
+	})
+	require.True(t, set.Enabled(codersdk.FeatureMultipleOrganizations))
+}