Skip to content

Commit 87cb577

Browse files
sreyasreya
committed
Refactor CryptoKey usage to use codersdk package
This change streamlines the CryptoKey handling by utilizing the codersdk package, thereby reducing code duplication and potentially simplifying maintenance efforts in the future.
1 parent 5f4beaa commit 87cb577

File tree

3 files changed

+57
-85
lines changed

3 files changed

+57
-85
lines changed

‎enterprise/wsproxy/keycache.go

Lines changed: 21 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99

1010
"cdr.dev/slog"
1111

12+
"github.com/coder/coder/v2/codersdk"
1213
"github.com/coder/coder/v2/enterprise/wsproxy/wsproxysdk"
1314
"github.com/coder/quartz"
1415
)
@@ -19,8 +20,8 @@ type CryptoKeyCache struct {
1920
Clock quartz.Clock
2021

2122
keysMu sync.RWMutex
22-
keys map[int32]wsproxysdk.CryptoKey
23-
latest wsproxysdk.CryptoKey
23+
keys map[int32]codersdk.CryptoKey
24+
latest codersdk.CryptoKey
2425
}
2526

2627
func NewCryptoKeyCache(ctx context.Context, log slog.Logger, client *wsproxysdk.Client, opts ...func(*CryptoKeyCache)) (*CryptoKeyCache, error) {
@@ -45,37 +46,37 @@ func NewCryptoKeyCache(ctx context.Context, log slog.Logger, client *wsproxysdk.
4546
return cache, nil
4647
}
4748

48-
func (k *CryptoKeyCache) Latest(ctx context.Context) (wsproxysdk.CryptoKey, error) {
49+
func (k *CryptoKeyCache) Latest(ctx context.Context) (codersdk.CryptoKey, error) {
4950
k.keysMu.RLock()
5051
latest := k.latest
5152
k.keysMu.RUnlock()
5253

5354
now := k.Clock.Now().UTC()
54-
if latest.Active(now) {
55+
if latest.CanSign(now) {
5556
return latest, nil
5657
}
5758

5859
k.keysMu.Lock()
5960
defer k.keysMu.Unlock()
6061

61-
if k.latest.Active(now) {
62+
if k.latest.CanSign(now) {
6263
return k.latest, nil
6364
}
6465

6566
var err error
6667
k.keys, k.latest, err = k.fetch(ctx)
6768
if err != nil {
68-
return wsproxysdk.CryptoKey{}, xerrors.Errorf("fetch: %w", err)
69+
return codersdk.CryptoKey{}, xerrors.Errorf("fetch: %w", err)
6970
}
7071

71-
if !k.latest.Active(now) {
72-
return wsproxysdk.CryptoKey{}, xerrors.Errorf("no active keys found")
72+
if !k.latest.CanSign(now) {
73+
return codersdk.CryptoKey{}, xerrors.Errorf("no active keys found")
7374
}
7475

7576
return k.latest, nil
7677
}
7778

78-
func (k *CryptoKeyCache) Version(ctx context.Context, sequence int32) (wsproxysdk.CryptoKey, error) {
79+
func (k *CryptoKeyCache) Version(ctx context.Context, sequence int32) (codersdk.CryptoKey, error) {
7980
now := k.Clock.Now().UTC()
8081
k.keysMu.RLock()
8182
key, ok := k.keys[sequence]
@@ -94,12 +95,12 @@ func (k *CryptoKeyCache) Version(ctx context.Context, sequence int32) (wsproxysd
9495
var err error
9596
k.keys, k.latest, err = k.fetch(ctx)
9697
if err != nil {
97-
return wsproxysdk.CryptoKey{}, xerrors.Errorf("fetch: %w", err)
98+
return codersdk.CryptoKey{}, xerrors.Errorf("fetch: %w", err)
9899
}
99100

100101
key, ok = k.keys[sequence]
101102
if !ok {
102-
return wsproxysdk.CryptoKey{}, xerrors.Errorf("key %d not found", sequence)
103+
return codersdk.CryptoKey{}, xerrors.Errorf("key %d not found", sequence)
103104
}
104105

105106
return validKey(key, now)
@@ -121,31 +122,31 @@ func (k *CryptoKeyCache) refresh(ctx context.Context) {
121122
})
122123
}
123124

124-
func (k *CryptoKeyCache) fetch(ctx context.Context) (map[int32]wsproxysdk.CryptoKey, wsproxysdk.CryptoKey, error) {
125+
func (k *CryptoKeyCache) fetch(ctx context.Context) (map[int32]codersdk.CryptoKey, codersdk.CryptoKey, error) {
125126
keys, err := k.client.CryptoKeys(ctx)
126127
if err != nil {
127-
return nil, wsproxysdk.CryptoKey{}, xerrors.Errorf("get security keys: %w", err)
128+
return nil, codersdk.CryptoKey{}, xerrors.Errorf("get security keys: %w", err)
128129
}
129130

130131
kmap, latest := toKeyMap(keys.CryptoKeys, k.Clock.Now().UTC())
131132
return kmap, latest, nil
132133
}
133134

134-
func toKeyMap(keys []wsproxysdk.CryptoKey, now time.Time) (map[int32]wsproxysdk.CryptoKey, wsproxysdk.CryptoKey) {
135-
m := make(map[int32]wsproxysdk.CryptoKey)
136-
var latest wsproxysdk.CryptoKey
135+
func toKeyMap(keys []codersdk.CryptoKey, now time.Time) (map[int32]codersdk.CryptoKey, codersdk.CryptoKey) {
136+
m := make(map[int32]codersdk.CryptoKey)
137+
var latest codersdk.CryptoKey
137138
for _, key := range keys {
138139
m[key.Sequence] = key
139-
if key.Sequence > latest.Sequence && key.Active(now) {
140+
if key.Sequence > latest.Sequence && key.CanSign(now) {
140141
latest = key
141142
}
142143
}
143144
return m, latest
144145
}
145146

146-
func validKey(key wsproxysdk.CryptoKey, now time.Time) (wsproxysdk.CryptoKey, error) {
147-
if key.Invalid(now) {
148-
return wsproxysdk.CryptoKey{}, xerrors.Errorf("key %d is invalid", key.Sequence)
147+
func validKey(key codersdk.CryptoKey, now time.Time) (codersdk.CryptoKey, error) {
148+
if !key.CanSign(now) {
149+
return codersdk.CryptoKey{}, xerrors.Errorf("key %d is invalid", key.Sequence)
149150
}
150151

151152
return key, nil

‎enterprise/wsproxy/keycache_test.go

Lines changed: 36 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212

1313
"cdr.dev/slog/sloggers/slogtest"
1414

15+
"github.com/coder/coder/v2/codersdk"
1516
"github.com/coder/coder/v2/enterprise/wsproxy"
1617
"github.com/coder/coder/v2/enterprise/wsproxy/wsproxysdk"
1718
"github.com/coder/coder/v2/testutil"
@@ -33,23 +34,23 @@ func TestCryptoKeyCache(t *testing.T) {
3334
)
3435

3536
now := clock.Now().UTC()
36-
expected := wsproxysdk.CryptoKey{
37-
Feature: wsproxysdk.CryptoKeyFeatureWorkspaceApp,
37+
expected := codersdk.CryptoKey{
38+
Feature: codersdk.CryptoKeyFeatureWorkspaceApp,
3839
Secret: "key2",
3940
Sequence: 2,
4041
StartsAt: now,
4142
}
4243

43-
fc := newFakeCoderd(t, []wsproxysdk.CryptoKey{
44+
fc := newFakeCoderd(t, []codersdk.CryptoKey{
4445
{
45-
Feature: wsproxysdk.CryptoKeyFeatureWorkspaceApp,
46+
Feature: codersdk.CryptoKeyFeatureWorkspaceApp,
4647
Secret: "key1",
4748
Sequence: 1,
4849
StartsAt: now,
4950
},
5051
// Should be ignored since it hasn't breached its starts_at time yet.
5152
{
52-
Feature: wsproxysdk.CryptoKeyFeatureWorkspaceApp,
53+
Feature: codersdk.CryptoKeyFeatureWorkspaceApp,
5354
Secret: "key3",
5455
Sequence: 3,
5556
StartsAt: now.Add(time.Second * 2),
@@ -74,18 +75,18 @@ func TestCryptoKeyCache(t *testing.T) {
7475
clock = quartz.NewMock(t)
7576
)
7677

77-
fc := newFakeCoderd(t, []wsproxysdk.CryptoKey{})
78+
fc := newFakeCoderd(t, []codersdk.CryptoKey{})
7879

7980
cache, err := wsproxy.NewCryptoKeyCache(ctx, logger, wsproxysdk.New(fc.url), withClock(clock))
8081
require.NoError(t, err)
8182

82-
expected := wsproxysdk.CryptoKey{
83-
Feature: wsproxysdk.CryptoKeyFeatureWorkspaceApp,
83+
expected := codersdk.CryptoKey{
84+
Feature: codersdk.CryptoKeyFeatureWorkspaceApp,
8485
Secret: "key1",
8586
Sequence: 12,
8687
StartsAt: clock.Now().UTC(),
8788
}
88-
fc.keys = []wsproxysdk.CryptoKey{expected}
89+
fc.keys = []codersdk.CryptoKey{expected}
8990

9091
got, err := cache.Latest(ctx)
9192
require.NoError(t, err)
@@ -110,17 +111,17 @@ func TestCryptoKeyCache(t *testing.T) {
110111
clock = quartz.NewMock(t)
111112
)
112113
now := clock.Now().UTC()
113-
expected := wsproxysdk.CryptoKey{
114-
Feature: wsproxysdk.CryptoKeyFeatureWorkspaceApp,
114+
expected := codersdk.CryptoKey{
115+
Feature: codersdk.CryptoKeyFeatureWorkspaceApp,
115116
Secret: "key1",
116117
Sequence: 1,
117118
StartsAt: clock.Now().UTC(),
118119
}
119120

120-
fc := newFakeCoderd(t, []wsproxysdk.CryptoKey{
121+
fc := newFakeCoderd(t, []codersdk.CryptoKey{
121122
expected,
122123
{
123-
Feature: wsproxysdk.CryptoKeyFeatureWorkspaceApp,
124+
Feature: codersdk.CryptoKeyFeatureWorkspaceApp,
124125
Secret: "key2",
125126
Sequence: 2,
126127
StartsAt: now.Add(-time.Second),
@@ -151,16 +152,16 @@ func TestCryptoKeyCache(t *testing.T) {
151152
)
152153

153154
now := clock.Now().UTC()
154-
expected := wsproxysdk.CryptoKey{
155-
Feature: wsproxysdk.CryptoKeyFeatureWorkspaceApp,
155+
expected := codersdk.CryptoKey{
156+
Feature: codersdk.CryptoKeyFeatureWorkspaceApp,
156157
Secret: "key1",
157158
Sequence: 12,
158159
StartsAt: now,
159160
}
160-
fc := newFakeCoderd(t, []wsproxysdk.CryptoKey{
161+
fc := newFakeCoderd(t, []codersdk.CryptoKey{
161162
expected,
162163
{
163-
Feature: wsproxysdk.CryptoKeyFeatureWorkspaceApp,
164+
Feature: codersdk.CryptoKeyFeatureWorkspaceApp,
164165
Secret: "key2",
165166
Sequence: 13,
166167
StartsAt: now,
@@ -184,18 +185,18 @@ func TestCryptoKeyCache(t *testing.T) {
184185
clock = quartz.NewMock(t)
185186
)
186187

187-
fc := newFakeCoderd(t, []wsproxysdk.CryptoKey{})
188+
fc := newFakeCoderd(t, []codersdk.CryptoKey{})
188189

189190
cache, err := wsproxy.NewCryptoKeyCache(ctx, logger, wsproxysdk.New(fc.url), withClock(clock))
190191
require.NoError(t, err)
191192

192-
expected := wsproxysdk.CryptoKey{
193-
Feature: wsproxysdk.CryptoKeyFeatureWorkspaceApp,
193+
expected := codersdk.CryptoKey{
194+
Feature: codersdk.CryptoKeyFeatureWorkspaceApp,
194195
Secret: "key1",
195196
Sequence: 12,
196197
StartsAt: clock.Now().UTC(),
197198
}
198-
fc.keys = []wsproxysdk.CryptoKey{expected}
199+
fc.keys = []codersdk.CryptoKey{expected}
199200

200201
got, err := cache.Version(ctx, expected.Sequence)
201202
require.NoError(t, err)
@@ -219,14 +220,14 @@ func TestCryptoKeyCache(t *testing.T) {
219220
)
220221

221222
now := clock.Now().UTC()
222-
expected := wsproxysdk.CryptoKey{
223-
Feature: wsproxysdk.CryptoKeyFeatureWorkspaceApp,
223+
expected := codersdk.CryptoKey{
224+
Feature: codersdk.CryptoKeyFeatureWorkspaceApp,
224225
Secret: "key1",
225226
Sequence: 12,
226227
StartsAt: now.Add(-time.Second),
227228
}
228229

229-
fc := newFakeCoderd(t, []wsproxysdk.CryptoKey{
230+
fc := newFakeCoderd(t, []codersdk.CryptoKey{
230231
expected,
231232
})
232233

@@ -249,15 +250,15 @@ func TestCryptoKeyCache(t *testing.T) {
249250
)
250251

251252
now := clock.Now().UTC()
252-
expected := wsproxysdk.CryptoKey{
253-
Feature: wsproxysdk.CryptoKeyFeatureWorkspaceApp,
253+
expected := codersdk.CryptoKey{
254+
Feature: codersdk.CryptoKeyFeatureWorkspaceApp,
254255
Secret: "key1",
255256
Sequence: 12,
256257
StartsAt: now.Add(-time.Second),
257258
DeletesAt: now,
258259
}
259260

260-
fc := newFakeCoderd(t, []wsproxysdk.CryptoKey{
261+
fc := newFakeCoderd(t, []codersdk.CryptoKey{
261262
expected,
262263
})
263264

@@ -282,14 +283,14 @@ func TestCryptoKeyCache(t *testing.T) {
282283
trap := clock.Trap().TickerFunc()
283284

284285
now := clock.Now().UTC()
285-
expected := wsproxysdk.CryptoKey{
286-
Feature: wsproxysdk.CryptoKeyFeatureWorkspaceApp,
286+
expected := codersdk.CryptoKey{
287+
Feature: codersdk.CryptoKeyFeatureWorkspaceApp,
287288
Secret: "key1",
288289
Sequence: 12,
289290
StartsAt: now,
290291
DeletesAt: now.Add(time.Minute * 10),
291292
}
292-
fc := newFakeCoderd(t, []wsproxysdk.CryptoKey{
293+
fc := newFakeCoderd(t, []codersdk.CryptoKey{
293294
expected,
294295
})
295296

@@ -303,13 +304,13 @@ func TestCryptoKeyCache(t *testing.T) {
303304

304305
wait := trap.MustWait(ctx)
305306

306-
newKey := wsproxysdk.CryptoKey{
307-
Feature: wsproxysdk.CryptoKeyFeatureWorkspaceApp,
307+
newKey := codersdk.CryptoKey{
308+
Feature: codersdk.CryptoKeyFeatureWorkspaceApp,
308309
Secret: "key2",
309310
Sequence: 13,
310311
StartsAt: now,
311312
}
312-
fc.keys = []wsproxysdk.CryptoKey{newKey}
313+
fc.keys = []codersdk.CryptoKey{newKey}
313314

314315
wait.Release()
315316

@@ -332,12 +333,12 @@ func TestCryptoKeyCache(t *testing.T) {
332333

333334
type fakeCoderd struct {
334335
server *httptest.Server
335-
keys []wsproxysdk.CryptoKey
336+
keys []codersdk.CryptoKey
336337
called int
337338
url *url.URL
338339
}
339340

340-
func newFakeCoderd(t *testing.T, keys []wsproxysdk.CryptoKey) *fakeCoderd {
341+
func newFakeCoderd(t *testing.T, keys []codersdk.CryptoKey) *fakeCoderd {
341342
t.Helper()
342343

343344
c := &fakeCoderd{

‎enterprise/wsproxy/wsproxysdk/wsproxysdk.go

Lines changed: 0 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -204,37 +204,7 @@ type RegisterWorkspaceProxyRequest struct {
204204
Version string `json:"version"`
205205
}
206206

207-
type CryptoKeyFeature string
208-
209-
const (
210-
CryptoKeyFeatureWorkspaceApp CryptoKeyFeature = "workspace_apps"
211-
CryptoKeyFeatureOIDCConvert CryptoKeyFeature = "oidc_convert"
212-
CryptoKeyFeatureTailnetResume CryptoKeyFeature = "tailnet_resume"
213-
)
214-
215-
type CryptoKey struct {
216-
Feature CryptoKeyFeature `json:"feature"`
217-
Secret string `json:"secret"`
218-
DeletesAt time.Time `json:"deletes_at"`
219-
Sequence int32 `json:"sequence"`
220-
StartsAt time.Time `json:"starts_at"`
221-
}
222-
223-
func (c CryptoKey) Active(now time.Time) bool {
224-
now = now.UTC()
225-
isAfterStartsAt := !c.StartsAt.IsZero() && !now.Before(c.StartsAt)
226-
return isAfterStartsAt && !c.Invalid(now)
227-
}
228-
229-
func (c CryptoKey) Invalid(now time.Time) bool {
230-
now = now.UTC()
231-
noSecret := c.Secret == ""
232-
afterDelete := !c.DeletesAt.IsZero() && !now.Before(c.DeletesAt.UTC())
233-
return noSecret || afterDelete
234-
}
235-
236207
type RegisterWorkspaceProxyResponse struct {
237-
Keys []CryptoKey `json:"keys"`
238208
AppSecurityKey string `json:"app_security_key"`
239209
DERPMeshKey string `json:"derp_mesh_key"`
240210
DERPRegionID int32 `json:"derp_region_id"`

0 commit comments

Comments
 (0)