Skip to content

Commit 152f695

Browse files
kylecarbskylecarbs
committed
Improve tests for license enablement code
1 parent d10712d commit 152f695

File tree

12 files changed

+486
-336
lines changed

12 files changed

+486
-336
lines changed

‎.vscode/settings.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,8 @@
99
"cliflag",
1010
"cliui",
1111
"codecov",
12-
"Codespaces",
1312
"coderd",
13+
"coderdenttest",
1414
"coderdtest",
1515
"codersdk",
1616
"cronstrue",
@@ -24,6 +24,7 @@
2424
"drpcmux",
2525
"drpcserver",
2626
"Dsts",
27+
"enablements",
2728
"fatih",
2829
"Formik",
2930
"gitsshkey",

‎enterprise/cli/features_test.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ func TestFeaturesList(t *testing.T) {
5757
var entitlements codersdk.Entitlements
5858
err := json.Unmarshal(buf.Bytes(), &entitlements)
5959
require.NoError(t, err, "unmarshal JSON output")
60-
assert.Len(t, entitlements.Features, 4)
60+
assert.Len(t, entitlements.Features, 5)
6161
assert.Empty(t, entitlements.Warnings)
6262
assert.Equal(t, codersdk.EntitlementNotEntitled,
6363
entitlements.Features[codersdk.FeatureUserLimit].Entitlement)
@@ -67,6 +67,8 @@ func TestFeaturesList(t *testing.T) {
6767
entitlements.Features[codersdk.FeatureBrowserOnly].Entitlement)
6868
assert.Equal(t, codersdk.EntitlementNotEntitled,
6969
entitlements.Features[codersdk.FeatureWorkspaceQuota].Entitlement)
70+
assert.Equal(t, codersdk.EntitlementNotEntitled,
71+
entitlements.Features[codersdk.FeatureSCIM].Entitlement)
7072
assert.False(t, entitlements.HasLicense)
7173
assert.False(t, entitlements.Experimental)
7274
})

‎enterprise/coderd/coderd.go

Lines changed: 35 additions & 178 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ package coderd
33
import (
44
"context"
55
"crypto/ed25519"
6-
"fmt"
76
"net/http"
87
"sync"
98
"time"
@@ -15,11 +14,14 @@ import (
1514

1615
"cdr.dev/slog"
1716
"github.com/coder/coder/coderd"
17+
agplaudit "github.com/coder/coder/coderd/audit"
1818
"github.com/coder/coder/coderd/httpapi"
1919
"github.com/coder/coder/coderd/httpmw"
20+
"github.com/coder/coder/coderd/workspacequota"
2021
"github.com/coder/coder/codersdk"
2122
"github.com/coder/coder/enterprise/audit"
2223
"github.com/coder/coder/enterprise/audit/backends"
24+
"github.com/coder/coder/enterprise/coderd/license"
2325
)
2426

2527
// New constructs an Enterprise coderd API instance.
@@ -34,19 +36,8 @@ func New(ctx context.Context, options *Options) (*API, error) {
3436
}
3537
ctx, cancelFunc := context.WithCancel(ctx)
3638
api := &API{
37-
AGPL: coderd.New(options.Options),
38-
Options: options,
39-
40-
entitlements: entitlements{
41-
activeUsers: codersdk.Feature{
42-
Entitlement: codersdk.EntitlementNotEntitled,
43-
Enabled: false,
44-
},
45-
auditLogs: codersdk.EntitlementNotEntitled,
46-
browserOnly: codersdk.EntitlementNotEntitled,
47-
scim: codersdk.EntitlementNotEntitled,
48-
workspaceQuota: codersdk.EntitlementNotEntitled,
49-
},
39+
AGPL: coderd.New(options.Options),
40+
Options: options,
5041
cancelEntitlementsLoop: cancelFunc,
5142
}
5243
oauthConfigs := &httpmw.OAuth2Configs{
@@ -117,17 +108,7 @@ type API struct {
117108

118109
cancelEntitlementsLoop func()
119110
entitlementsMu sync.RWMutex
120-
entitlements entitlements
121-
}
122-
123-
type entitlements struct {
124-
hasLicense bool
125-
trial bool
126-
activeUsers codersdk.Feature
127-
auditLogs codersdk.Entitlement
128-
browserOnly codersdk.Entitlement
129-
scim codersdk.Entitlement
130-
workspaceQuota codersdk.Entitlement
111+
entitlements codersdk.Entitlements
131112
}
132113

133114
func (api *API) Close() error {
@@ -136,98 +117,57 @@ func (api *API) Close() error {
136117
}
137118

138119
func (api *API) updateEntitlements(ctx context.Context) error {
139-
licenses, err := api.Database.GetUnexpiredLicenses(ctx)
140-
if err != nil {
141-
return err
142-
}
143120
api.entitlementsMu.Lock()
144121
defer api.entitlementsMu.Unlock()
145-
now := time.Now()
146122

147-
// Default all entitlements to be disabled.
148-
entitlements := entitlements{
149-
hasLicense: false,
150-
activeUsers: codersdk.Feature{
151-
Enabled: false,
152-
Entitlement: codersdk.EntitlementNotEntitled,
153-
},
154-
auditLogs: codersdk.EntitlementNotEntitled,
155-
scim: codersdk.EntitlementNotEntitled,
156-
browserOnly: codersdk.EntitlementNotEntitled,
157-
workspaceQuota: codersdk.EntitlementNotEntitled,
158-
trial: true,
123+
entitlements, err := license.Entitlements(ctx, api.Database, api.Logger, api.Keys, map[string]bool{
124+
codersdk.FeatureAuditLog: api.AuditLogging,
125+
codersdk.FeatureBrowserOnly: api.BrowserOnly,
126+
codersdk.FeatureSCIM: len(api.SCIMAPIKey) != 0,
127+
codersdk.FeatureWorkspaceQuota: api.UserWorkspaceQuota != 0,
128+
})
129+
if err != nil {
130+
return err
159131
}
160132

161-
// Here we loop through licenses to detect enabled features.
162-
for _, l := range licenses {
163-
claims, err := validateDBLicense(l, api.Keys)
164-
if err != nil {
165-
api.Logger.Debug(ctx, "skipping invalid license",
166-
slog.F("id", l.ID), slog.Error(err))
167-
continue
168-
}
169-
entitlements.hasLicense = true
170-
entitlement := codersdk.EntitlementEntitled
171-
if now.After(claims.LicenseExpires.Time) {
172-
// if the grace period were over, the validation fails, so if we are after
173-
// LicenseExpires we must be in grace period.
174-
entitlement = codersdk.EntitlementGracePeriod
133+
featureChanged := func(featureName string) (changed bool, enabled bool) {
134+
if api.entitlements.Features == nil {
135+
return true, entitlements.Features[featureName].Enabled
175136
}
176-
if claims.Features.UserLimit > 0 {
177-
entitlements.activeUsers = codersdk.Feature{
178-
Enabled: true,
179-
Entitlement: entitlement,
180-
}
181-
currentLimit := int64(0)
182-
if entitlements.activeUsers.Limit != nil {
183-
currentLimit = *entitlements.activeUsers.Limit
184-
}
185-
limit := max(currentLimit, claims.Features.UserLimit)
186-
entitlements.activeUsers.Limit = &limit
187-
}
188-
if claims.Features.AuditLog > 0 {
189-
entitlements.auditLogs = entitlement
190-
}
191-
if claims.Features.BrowserOnly > 0 {
192-
entitlements.browserOnly = entitlement
193-
}
194-
if claims.Features.SCIM > 0 {
195-
entitlements.scim = entitlement
196-
}
197-
if claims.Features.WorkspaceQuota > 0 {
198-
entitlements.workspaceQuota = entitlement
199-
}
200-
if !claims.Trial {
201-
entitlements.trial = claims.Trial
137+
oldFeature := api.entitlements.Features[featureName]
138+
newFeature := entitlements.Features[featureName]
139+
if oldFeature.Enabled != newFeature.Enabled {
140+
return true, newFeature.Enabled
202141
}
142+
return false, newFeature.Enabled
203143
}
204144

205-
if entitlements.auditLogs != api.entitlements.auditLogs {
206-
// A flag could be added to the options that would allow disabling
207-
// enhanced audit logging here!
208-
if entitlements.auditLogs != codersdk.EntitlementNotEntitled && api.AuditLogging {
209-
auditor := audit.NewAuditor(
145+
if changed, enabled := featureChanged(codersdk.FeatureAuditLog); changed {
146+
auditor := agplaudit.NewNop()
147+
if enabled {
148+
auditor = audit.NewAuditor(
210149
audit.DefaultFilter,
211150
backends.NewPostgres(api.Database, true),
212151
backends.NewSlog(api.Logger),
213152
)
214-
api.AGPL.Auditor.Store(&auditor)
215153
}
154+
api.AGPL.Auditor.Store(&auditor)
216155
}
217156

218-
if entitlements.browserOnly != api.entitlements.browserOnly {
157+
if changed, enabled := featureChanged(codersdk.FeatureBrowserOnly); changed {
219158
var handler func(rw http.ResponseWriter) bool
220-
if entitlements.browserOnly != codersdk.EntitlementNotEntitled && api.BrowserOnly {
159+
if enabled {
221160
handler = api.shouldBlockNonBrowserConnections
222161
}
223162
api.AGPL.WorkspaceClientCoordinateOverride.Store(&handler)
224163
}
225164

226-
if entitlements.workspaceQuota != api.entitlements.workspaceQuota {
227-
if entitlements.workspaceQuota != codersdk.EntitlementNotEntitled && api.UserWorkspaceQuota > 0 {
228-
enforcer := NewEnforcer(api.Options.UserWorkspaceQuota)
229-
api.AGPL.WorkspaceQuotaEnforcer.Store(&enforcer)
165+
if changed, enabled := featureChanged(codersdk.FeatureWorkspaceQuota); changed {
166+
enforcer := workspacequota.NewNop()
167+
if enabled {
168+
enforcer = NewEnforcer(api.Options.UserWorkspaceQuota)
230169
}
170+
api.AGPL.WorkspaceQuotaEnforcer.Store(&enforcer)
231171
}
232172

233173
api.entitlements = entitlements
@@ -240,83 +180,7 @@ func (api *API) serveEntitlements(rw http.ResponseWriter, r *http.Request) {
240180
api.entitlementsMu.RLock()
241181
entitlements := api.entitlements
242182
api.entitlementsMu.RUnlock()
243-
244-
resp := codersdk.Entitlements{
245-
Features: make(map[string]codersdk.Feature),
246-
Warnings: make([]string, 0),
247-
HasLicense: entitlements.hasLicense,
248-
Trial: entitlements.trial,
249-
Experimental: api.Experimental,
250-
}
251-
252-
if entitlements.activeUsers.Limit != nil {
253-
activeUserCount, err := api.Database.GetActiveUserCount(ctx)
254-
if err != nil {
255-
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
256-
Message: "Unable to query database",
257-
Detail: err.Error(),
258-
})
259-
return
260-
}
261-
entitlements.activeUsers.Actual = &activeUserCount
262-
if activeUserCount > *entitlements.activeUsers.Limit {
263-
resp.Warnings = append(resp.Warnings,
264-
fmt.Sprintf(
265-
"Your deployment has %d active users but is only licensed for %d.",
266-
activeUserCount, *entitlements.activeUsers.Limit))
267-
}
268-
}
269-
resp.Features[codersdk.FeatureUserLimit] = entitlements.activeUsers
270-
271-
// Audit logs
272-
resp.Features[codersdk.FeatureAuditLog] = codersdk.Feature{
273-
Entitlement: entitlements.auditLogs,
274-
Enabled: api.AuditLogging,
275-
}
276-
// Audit logging is enabled by default. We don't want to display
277-
// a warning if they don't have a license.
278-
if entitlements.hasLicense && api.AuditLogging {
279-
if entitlements.auditLogs == codersdk.EntitlementNotEntitled {
280-
resp.Warnings = append(resp.Warnings,
281-
"Audit logging is enabled but your license is not entitled to this feature.")
282-
}
283-
if entitlements.auditLogs == codersdk.EntitlementGracePeriod {
284-
resp.Warnings = append(resp.Warnings,
285-
"Audit logging is enabled but your license for this feature is expired.")
286-
}
287-
}
288-
289-
resp.Features[codersdk.FeatureBrowserOnly] = codersdk.Feature{
290-
Entitlement: entitlements.browserOnly,
291-
Enabled: api.BrowserOnly,
292-
}
293-
if api.BrowserOnly {
294-
if entitlements.browserOnly == codersdk.EntitlementNotEntitled {
295-
resp.Warnings = append(resp.Warnings,
296-
"Browser only connections are enabled but your license is not entitled to this feature.")
297-
}
298-
if entitlements.browserOnly == codersdk.EntitlementGracePeriod {
299-
resp.Warnings = append(resp.Warnings,
300-
"Browser only connections are enabled but your license for this feature is expired.")
301-
}
302-
}
303-
304-
resp.Features[codersdk.FeatureWorkspaceQuota] = codersdk.Feature{
305-
Entitlement: entitlements.workspaceQuota,
306-
Enabled: api.UserWorkspaceQuota > 0,
307-
}
308-
if api.UserWorkspaceQuota > 0 {
309-
if entitlements.workspaceQuota == codersdk.EntitlementNotEntitled {
310-
resp.Warnings = append(resp.Warnings,
311-
"Workspace quotas are enabled but your license is not entitled to this feature.")
312-
}
313-
if entitlements.workspaceQuota == codersdk.EntitlementGracePeriod {
314-
resp.Warnings = append(resp.Warnings,
315-
"Workspace quotas are enabled but your license for this feature is expired.")
316-
}
317-
}
318-
319-
httpapi.Write(ctx, rw, http.StatusOK, resp)
183+
httpapi.Write(ctx, rw, http.StatusOK, entitlements)
320184
}
321185

322186
func (api *API) runEntitlementsLoop(ctx context.Context) {
@@ -380,10 +244,3 @@ func (api *API) runEntitlementsLoop(ctx context.Context) {
380244
}
381245
}
382246
}
383-
384-
func max(a, b int64) int64 {
385-
if a > b {
386-
return a
387-
}
388-
return b
389-
}

‎enterprise/coderd/coderd_test.go

Lines changed: 0 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -86,48 +86,6 @@ func TestEntitlements(t *testing.T) {
8686
assert.Equal(t, codersdk.EntitlementNotEntitled, al.Entitlement)
8787
assert.True(t, al.Enabled)
8888
})
89-
t.Run("Warnings", func(t *testing.T) {
90-
t.Parallel()
91-
client := coderdenttest.New(t, &coderdenttest.Options{
92-
AuditLogging: true,
93-
BrowserOnly: true,
94-
})
95-
first := coderdtest.CreateFirstUser(t, client)
96-
for i := 0; i < 4; i++ {
97-
coderdtest.CreateAnotherUser(t, client, first.OrganizationID)
98-
}
99-
coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
100-
UserLimit: 4,
101-
AuditLog: true,
102-
BrowserOnly: true,
103-
GraceAt: time.Now().Add(-time.Second),
104-
})
105-
res, err := client.Entitlements(context.Background())
106-
require.NoError(t, err)
107-
assert.True(t, res.HasLicense)
108-
ul := res.Features[codersdk.FeatureUserLimit]
109-
assert.Equal(t, codersdk.EntitlementGracePeriod, ul.Entitlement)
110-
assert.Equal(t, int64(4), *ul.Limit)
111-
assert.Equal(t, int64(5), *ul.Actual)
112-
assert.True(t, ul.Enabled)
113-
al := res.Features[codersdk.FeatureAuditLog]
114-
assert.Equal(t, codersdk.EntitlementGracePeriod, al.Entitlement)
115-
assert.True(t, al.Enabled)
116-
assert.Nil(t, al.Limit)
117-
assert.Nil(t, al.Actual)
118-
bo := res.Features[codersdk.FeatureBrowserOnly]
119-
assert.Equal(t, codersdk.EntitlementGracePeriod, bo.Entitlement)
120-
assert.True(t, bo.Enabled)
121-
assert.Nil(t, bo.Limit)
122-
assert.Nil(t, bo.Actual)
123-
assert.Len(t, res.Warnings, 3)
124-
assert.Contains(t, res.Warnings,
125-
"Your deployment has 5 active users but is only licensed for 4.")
126-
assert.Contains(t, res.Warnings,
127-
"Audit logging is enabled but your license for this feature is expired.")
128-
assert.Contains(t, res.Warnings,
129-
"Browser only connections are enabled but your license for this feature is expired.")
130-
})
13189
t.Run("Pubsub", func(t *testing.T) {
13290
t.Parallel()
13391
client, _, api := coderdenttest.NewWithAPI(t, nil)

0 commit comments

Comments
 (0)