Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 355f3c0

Browse files
ammarioammario
committed
refactor: deduplicate / type license feature code
1 parent 592ce3b commit 355f3c0

File tree

6 files changed

+103
-69
lines changed

6 files changed

+103
-69
lines changed

‎codersdk/features.go

Lines changed: 26 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"context"
55
"encoding/json"
66
"net/http"
7+
"strings"
78
)
89

910
type Entitlement string
@@ -14,19 +15,28 @@ const (
1415
EntitlementNotEntitled Entitlement = "not_entitled"
1516
)
1617

18+
// To add a new feature, modify this set of enums as well as the FeatureNames
19+
// array.
20+
type FeatureName string
21+
1722
const (
18-
FeatureUserLimit = "user_limit"
19-
FeatureAuditLog = "audit_log"
20-
FeatureBrowserOnly = "browser_only"
21-
FeatureSCIM = "scim"
22-
FeatureTemplateRBAC = "template_rbac"
23-
FeatureHighAvailability = "high_availability"
24-
FeatureMultipleGitAuth = "multiple_git_auth"
25-
FeatureExternalProvisionerDaemons = "external_provisioner_daemons"
26-
FeatureAppearance = "appearance"
23+
FeatureUserLimit FeatureName = "user_limit"
24+
FeatureAuditLog FeatureName = "audit_log"
25+
FeatureBrowserOnly FeatureName = "browser_only"
26+
FeatureSCIM FeatureName = "scim"
27+
FeatureTemplateRBAC FeatureName = "template_rbac"
28+
FeatureHighAvailability FeatureName = "high_availability"
29+
FeatureMultipleGitAuth FeatureName = "multiple_git_auth"
30+
FeatureExternalProvisionerDaemons FeatureName = "external_provisioner_daemons"
31+
FeatureAppearance FeatureName = "appearance"
2732
)
2833

29-
var FeatureNames = []string{
34+
// Humanize returns the feature name in a human-readable format.
35+
func (n FeatureName) Humanize() string {
36+
return strings.Title(strings.ReplaceAll(string(n), "_", " "))
37+
}
38+
39+
var FeatureNames = []FeatureName{
3040
FeatureUserLimit,
3141
FeatureAuditLog,
3242
FeatureBrowserOnly,
@@ -46,12 +56,12 @@ type Feature struct {
4656
}
4757

4858
type Entitlements struct {
49-
Features map[string]Feature `json:"features"`
50-
Warnings []string `json:"warnings"`
51-
Errors []string `json:"errors"`
52-
HasLicense bool `json:"has_license"`
53-
Experimental bool `json:"experimental"`
54-
Trial bool `json:"trial"`
59+
Features map[FeatureName]Feature `json:"features"`
60+
Warnings []string `json:"warnings"`
61+
Errors []string `json:"errors"`
62+
HasLicense bool `json:"has_license"`
63+
Experimental bool `json:"experimental"`
64+
Trial bool `json:"trial"`
5565
}
5666

5767
func (c *Client) Entitlements(ctx context.Context) (Entitlements, error) {

‎codersdk/licenses.go

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"time"
99

1010
"github.com/google/uuid"
11+
"golang.org/x/xerrors"
1112
)
1213

1314
type AddLicenseRequest struct {
@@ -25,6 +26,30 @@ type License struct {
2526
Claims map[string]interface{} `json:"claims"`
2627
}
2728

29+
// Features provides the feature claims in license.
30+
func (l *License) Features() (map[FeatureName]int64, error) {
31+
strMap, ok := l.Claims["features"].(map[string]interface{})
32+
if !ok {
33+
return nil, xerrors.New("features key is unexpected type")
34+
}
35+
fMap := make(map[FeatureName]int64)
36+
for k, v := range strMap {
37+
jn, ok := v.(json.Number)
38+
if !ok {
39+
return nil, xerrors.Errorf("feature %q has unexpected type", k)
40+
}
41+
42+
n, err := jn.Int64()
43+
if err != nil {
44+
return nil, err
45+
}
46+
47+
fMap[FeatureName(k)] = n
48+
}
49+
50+
return fMap, nil
51+
}
52+
2853
func (c *Client) AddLicense(ctx context.Context, r AddLicenseRequest) (License, error) {
2954
res, err := c.Request(ctx, http.MethodPost, "/api/v2/licenses", r)
3055
if err != nil {

‎enterprise/coderd/coderd.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -238,7 +238,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
238238
api.entitlementsMu.Lock()
239239
defer api.entitlementsMu.Unlock()
240240

241-
entitlements, err := license.Entitlements(ctx, api.Database, api.Logger, len(api.replicaManager.All()), len(api.GitAuthConfigs), api.Keys, map[string]bool{
241+
entitlements, err := license.Entitlements(ctx, api.Database, api.Logger, len(api.replicaManager.All()), len(api.GitAuthConfigs), api.Keys, map[codersdk.FeatureName]bool{
242242
codersdk.FeatureAuditLog: api.AuditLogging,
243243
codersdk.FeatureBrowserOnly: api.BrowserOnly,
244244
codersdk.FeatureSCIM: len(api.SCIMAPIKey) != 0,
@@ -252,7 +252,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
252252
}
253253
entitlements.Experimental = api.DeploymentConfig.Experimental.Value
254254

255-
featureChanged := func(featureName string) (changed bool, enabled bool) {
255+
featureChanged := func(featureName codersdk.FeatureName) (changed bool, enabled bool) {
256256
if api.entitlements.Features == nil {
257257
return true, entitlements.Features[featureName].Enabled
258258
}

‎enterprise/coderd/license/license.go

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import (
44
"context"
55
"crypto/ed25519"
66
"fmt"
7-
"strings"
87
"time"
98

109
"github.com/golang-jwt/jwt/v4"
@@ -24,12 +23,12 @@ func Entitlements(
2423
replicaCount int,
2524
gitAuthCount int,
2625
keys map[string]ed25519.PublicKey,
27-
enablements map[string]bool,
26+
enablements map[codersdk.FeatureName]bool,
2827
) (codersdk.Entitlements, error) {
2928
now := time.Now()
3029
// Default all entitlements to be disabled.
3130
entitlements := codersdk.Entitlements{
32-
Features: map[string]codersdk.Feature{},
31+
Features: map[codersdk.FeatureName]codersdk.Feature{},
3332
Warnings: []string{},
3433
Errors: []string{},
3534
}
@@ -171,7 +170,7 @@ func Entitlements(
171170
if !feature.Enabled {
172171
continue
173172
}
174-
niceName := strings.Title(strings.ReplaceAll(featureName, "_", " "))
173+
niceName := featureName.Humanize()
175174
switch feature.Entitlement {
176175
case codersdk.EntitlementNotEntitled:
177176
entitlements.Warnings = append(entitlements.Warnings,

‎enterprise/coderd/license/license_test.go

Lines changed: 16 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ package license_test
33
import (
44
"context"
55
"fmt"
6-
"strings"
76
"testing"
87
"time"
98

@@ -19,17 +18,13 @@ import (
1918

2019
func TestEntitlements(t *testing.T) {
2120
t.Parallel()
22-
all := map[string]bool{
23-
codersdk.FeatureAuditLog: true,
24-
codersdk.FeatureBrowserOnly: true,
25-
codersdk.FeatureSCIM: true,
26-
codersdk.FeatureHighAvailability: true,
27-
codersdk.FeatureTemplateRBAC: true,
28-
codersdk.FeatureMultipleGitAuth: true,
29-
codersdk.FeatureExternalProvisionerDaemons: true,
30-
codersdk.FeatureAppearance: true,
21+
all := make(map[codersdk.FeatureName]bool)
22+
for _, n := range codersdk.FeatureNames {
23+
all[n] = true
3124
}
3225

26+
empty := map[codersdk.FeatureName]bool{}
27+
3328
t.Run("Defaults", func(t *testing.T) {
3429
t.Parallel()
3530
db := databasefake.New()
@@ -49,7 +44,7 @@ func TestEntitlements(t *testing.T) {
4944
JWT: coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{}),
5045
Exp: time.Now().Add(time.Hour),
5146
})
52-
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, map[string]bool{})
47+
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, empty)
5348
require.NoError(t, err)
5449
require.True(t, entitlements.HasLicense)
5550
require.False(t, entitlements.Trial)
@@ -75,7 +70,7 @@ func TestEntitlements(t *testing.T) {
7570
}),
7671
Exp: time.Now().Add(time.Hour),
7772
})
78-
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, map[string]bool{})
73+
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, empty)
7974
require.NoError(t, err)
8075
require.True(t, entitlements.HasLicense)
8176
require.False(t, entitlements.Trial)
@@ -115,7 +110,7 @@ func TestEntitlements(t *testing.T) {
115110
if featureName == codersdk.FeatureMultipleGitAuth {
116111
continue
117112
}
118-
niceName := strings.Title(strings.ReplaceAll(featureName, "_", " "))
113+
niceName := featureName.Humanize()
119114
require.Equal(t, codersdk.EntitlementGracePeriod, entitlements.Features[featureName].Entitlement)
120115
require.Contains(t, entitlements.Warnings, fmt.Sprintf("%s is enabled but your license for this feature is expired.", niceName))
121116
}
@@ -141,7 +136,7 @@ func TestEntitlements(t *testing.T) {
141136
if featureName == codersdk.FeatureMultipleGitAuth {
142137
continue
143138
}
144-
niceName := strings.Title(strings.ReplaceAll(featureName, "_", " "))
139+
niceName := featureName.Humanize()
145140
// Ensures features that are not entitled are properly disabled.
146141
require.False(t, entitlements.Features[featureName].Enabled)
147142
require.Equal(t, codersdk.EntitlementNotEntitled, entitlements.Features[featureName].Entitlement)
@@ -163,7 +158,7 @@ func TestEntitlements(t *testing.T) {
163158
}),
164159
Exp: time.Now().Add(time.Hour),
165160
})
166-
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, map[string]bool{})
161+
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, empty)
167162
require.NoError(t, err)
168163
require.True(t, entitlements.HasLicense)
169164
require.Contains(t, entitlements.Warnings, "Your deployment has 2 active users but is only licensed for 1.")
@@ -185,7 +180,7 @@ func TestEntitlements(t *testing.T) {
185180
}),
186181
Exp: time.Now().Add(time.Hour),
187182
})
188-
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, map[string]bool{})
183+
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, empty)
189184
require.NoError(t, err)
190185
require.True(t, entitlements.HasLicense)
191186
require.Empty(t, entitlements.Warnings)
@@ -208,7 +203,7 @@ func TestEntitlements(t *testing.T) {
208203
}),
209204
})
210205

211-
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, map[string]bool{})
206+
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, empty)
212207
require.NoError(t, err)
213208
require.True(t, entitlements.HasLicense)
214209
require.False(t, entitlements.Trial)
@@ -255,7 +250,7 @@ func TestEntitlements(t *testing.T) {
255250
AuditLog: true,
256251
}),
257252
})
258-
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 2, 1, coderdenttest.Keys, map[string]bool{
253+
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 2, 1, coderdenttest.Keys, map[codersdk.FeatureName]bool{
259254
codersdk.FeatureHighAvailability: true,
260255
})
261256
require.NoError(t, err)
@@ -275,7 +270,7 @@ func TestEntitlements(t *testing.T) {
275270
}),
276271
Exp: time.Now().Add(time.Hour),
277272
})
278-
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 2, 1, coderdenttest.Keys, map[string]bool{
273+
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 2, 1, coderdenttest.Keys, map[codersdk.FeatureName]bool{
279274
codersdk.FeatureHighAvailability: true,
280275
})
281276
require.NoError(t, err)
@@ -303,7 +298,7 @@ func TestEntitlements(t *testing.T) {
303298
AuditLog: true,
304299
}),
305300
})
306-
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 2, coderdenttest.Keys, map[string]bool{
301+
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 2, coderdenttest.Keys, map[codersdk.FeatureName]bool{
307302
codersdk.FeatureMultipleGitAuth: true,
308303
})
309304
require.NoError(t, err)
@@ -323,7 +318,7 @@ func TestEntitlements(t *testing.T) {
323318
}),
324319
Exp: time.Now().Add(time.Hour),
325320
})
326-
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 2, coderdenttest.Keys, map[string]bool{
321+
entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 2, coderdenttest.Keys, map[codersdk.FeatureName]bool{
327322
codersdk.FeatureMultipleGitAuth: true,
328323
})
329324
require.NoError(t, err)

‎enterprise/coderd/licenses_test.go

Lines changed: 31 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ package coderd_test
22

33
import (
44
"context"
5-
"encoding/json"
65
"net/http"
76
"testing"
87

@@ -32,9 +31,9 @@ func TestPostLicense(t *testing.T) {
3231
assert.GreaterOrEqual(t, respLic.ID, int32(0))
3332
// just a couple spot checks for sanity
3433
assert.Equal(t, "testing", respLic.Claims["account_id"])
35-
features, ok := respLic.Claims["features"].(map[string]interface{})
36-
require.True(t, ok)
37-
assert.Equal(t, json.Number("1"), features[codersdk.FeatureAuditLog])
34+
features, err := respLic.Features()
35+
require.NoError(t, err)
36+
assert.EqualValues(t, 1, features[codersdk.FeatureAuditLog])
3837
})
3938

4039
t.Run("Unauthorized", func(t *testing.T) {
@@ -100,31 +99,37 @@ func TestGetLicense(t *testing.T) {
10099
require.Len(t, licenses, 2)
101100
assert.Equal(t, int32(1), licenses[0].ID)
102101
assert.Equal(t, "testing", licenses[0].Claims["account_id"])
103-
assert.Equal(t, map[string]interface{}{
104-
codersdk.FeatureUserLimit: json.Number("0"),
105-
codersdk.FeatureAuditLog: json.Number("1"),
106-
codersdk.FeatureSCIM: json.Number("1"),
107-
codersdk.FeatureBrowserOnly: json.Number("1"),
108-
codersdk.FeatureHighAvailability: json.Number("0"),
109-
codersdk.FeatureTemplateRBAC: json.Number("1"),
110-
codersdk.FeatureMultipleGitAuth: json.Number("0"),
111-
codersdk.FeatureExternalProvisionerDaemons: json.Number("0"),
112-
codersdk.FeatureAppearance: json.Number("0"),
113-
}, licenses[0].Claims["features"])
102+
103+
features, err := licenses[0].Features()
104+
require.NoError(t, err)
105+
assert.Equal(t, map[codersdk.FeatureName]int64{
106+
codersdk.FeatureUserLimit: 0,
107+
codersdk.FeatureAuditLog: 1,
108+
codersdk.FeatureSCIM: 1,
109+
codersdk.FeatureBrowserOnly: 1,
110+
codersdk.FeatureHighAvailability: 0,
111+
codersdk.FeatureTemplateRBAC: 1,
112+
codersdk.FeatureMultipleGitAuth: 0,
113+
codersdk.FeatureExternalProvisionerDaemons: 0,
114+
codersdk.FeatureAppearance: 0,
115+
}, features)
114116
assert.Equal(t, int32(2), licenses[1].ID)
115117
assert.Equal(t, "testing2", licenses[1].Claims["account_id"])
116118
assert.Equal(t, true, licenses[1].Claims["trial"])
117-
assert.Equal(t, map[string]interface{}{
118-
codersdk.FeatureUserLimit: json.Number("200"),
119-
codersdk.FeatureAuditLog: json.Number("1"),
120-
codersdk.FeatureSCIM: json.Number("1"),
121-
codersdk.FeatureBrowserOnly: json.Number("1"),
122-
codersdk.FeatureHighAvailability: json.Number("0"),
123-
codersdk.FeatureTemplateRBAC: json.Number("0"),
124-
codersdk.FeatureMultipleGitAuth: json.Number("0"),
125-
codersdk.FeatureExternalProvisionerDaemons: json.Number("0"),
126-
codersdk.FeatureAppearance: json.Number("0"),
127-
}, licenses[1].Claims["features"])
119+
120+
features, err = licenses[1].Features()
121+
require.NoError(t, err)
122+
assert.Equal(t, map[codersdk.FeatureName]int64{
123+
codersdk.FeatureUserLimit: 200,
124+
codersdk.FeatureAuditLog: 1,
125+
codersdk.FeatureSCIM: 1,
126+
codersdk.FeatureBrowserOnly: 1,
127+
codersdk.FeatureHighAvailability: 0,
128+
codersdk.FeatureTemplateRBAC: 0,
129+
codersdk.FeatureMultipleGitAuth: 0,
130+
codersdk.FeatureExternalProvisionerDaemons: 0,
131+
codersdk.FeatureAppearance: 0,
132+
}, features)
128133
})
129134
}
130135

0 commit comments

Comments
 (0)