-
Notifications
You must be signed in to change notification settings - Fork 937
chore: refactor entitlements to be a safe object to use #14406
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
33b8199
66b363c
65fd791
8d90083
3e44716
04d9b03
aaea397
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Potential follow-up: we could replace this with There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. True. Because before we had access to the whole struct, our usage of it seemed a bit arbitrary at times. Sometimes we grab it and check entitled, most times just enabled. I'm not trying to fix all our usage right now, but it would be good to audit at some times. |
||
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) | ||
} | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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)) | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I implemented the methods as I saw them used. There might be a way to reduce the number of methods on this struct.