Skip to content

Commit e91fd31

Browse files
sreyasreya
committed
Add crypto key generation and refactor rotation tests
1 parent 1c7bfc8 commit e91fd31

File tree

6 files changed

+227
-104
lines changed

6 files changed

+227
-104
lines changed

‎coderd/database/dbgen/dbgen.go

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -893,6 +893,40 @@ func CustomRole(t testing.TB, db database.Store, seed database.CustomRole) datab
893893
return role
894894
}
895895

896+
func CryptoKey(t testing.TB, db database.Store, seed database.CryptoKey) database.CryptoKey {
897+
t.Helper()
898+
899+
secret, err := cryptorand.String(10)
900+
require.NoError(t, err, "generate secret")
901+
902+
sequence, err := cryptorand.Intn(1<<31 - 1)
903+
require.NoError(t, err, "generate sequence")
904+
905+
feature, err := cryptorand.Element(database.AllCryptoKeyFeatureValues())
906+
require.NoError(t, err, "generate feature")
907+
908+
key, err := db.InsertCryptoKey(genCtx, database.InsertCryptoKeyParams{
909+
Sequence: takeFirst(seed.Sequence, int32(sequence)),
910+
Secret: takeFirst(seed.Secret, sql.NullString{
911+
String: secret,
912+
Valid: true,
913+
}),
914+
Feature: takeFirst(seed.Feature, feature),
915+
StartsAt: takeFirst(seed.StartsAt, time.Now()),
916+
})
917+
require.NoError(t, err, "insert crypto key")
918+
919+
if seed.DeletesAt.Valid {
920+
key, err = db.UpdateCryptoKeyDeletesAt(genCtx, database.UpdateCryptoKeyDeletesAtParams{
921+
Feature: key.Feature,
922+
Sequence: key.Sequence,
923+
DeletesAt: sql.NullTime{Time: seed.DeletesAt.Time, Valid: true},
924+
})
925+
require.NoError(t, err, "update crypto key deletes_at")
926+
}
927+
return key
928+
}
929+
896930
func must[V any](v V, err error) V {
897931
if err != nil {
898932
panic(err)

‎coderd/database/queries.sql.go

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/keys.sql

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,8 @@ FROM crypto_keys
1717
WHERE feature = $1
1818
AND sequence = $2
1919
AND secret IS NOT NULL
20-
AND $3 >= starts_at
21-
AND ($3 < deletes_at OR deletes_at IS NULL);
20+
AND @time >= starts_at
21+
AND (@time < deletes_at OR deletes_at IS NULL);
2222

2323
-- name: DeleteCryptoKey :one
2424
UPDATE crypto_keys

‎coderd/keyrotate/rotate.go

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ func (k *KeyRotator) rotateKeys(ctx context.Context) ([]database.CryptoKey, erro
7272

7373
// Groups the keys by feature so that we can
7474
// ensure we have at least one key for each feature.
75-
keysByFeature := keysByFeature(keys)
75+
keysByFeature := keysByFeature(keys, k.features)
7676

7777
now := dbtime.Time(k.Clock.Now())
7878
for feature, keys := range keysByFeature {
@@ -232,11 +232,10 @@ func shouldRotateKey(key database.CryptoKey, keyDuration time.Duration, now time
232232
return now.Add(time.Hour).After(expirationTime)
233233
}
234234

235-
func keysByFeature(keys []database.CryptoKey) map[database.CryptoKeyFeature][]database.CryptoKey {
236-
m := map[database.CryptoKeyFeature][]database.CryptoKey{
237-
database.CryptoKeyFeatureWorkspaceApps: {},
238-
database.CryptoKeyFeatureOidcConvert: {},
239-
database.CryptoKeyFeaturePeerReconnect: {},
235+
func keysByFeature(keys []database.CryptoKey, features []database.CryptoKeyFeature) map[database.CryptoKeyFeature][]database.CryptoKey {
236+
m := map[database.CryptoKeyFeature][]database.CryptoKey{}
237+
for _, feature := range features {
238+
m[feature] = []database.CryptoKey{}
240239
}
241240
for _, key := range keys {
242241
m[key.Feature] = append(m[key.Feature], key)

‎coderd/keyrotate/rotate_internal_test.go

Lines changed: 184 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,184 @@
1+
package keyrotate
2+
3+
import (
4+
"database/sql"
5+
"encoding/hex"
6+
"testing"
7+
"time"
8+
9+
"github.com/stretchr/testify/require"
10+
11+
"cdr.dev/slog"
12+
"cdr.dev/slog/sloggers/slogtest"
13+
"github.com/coder/coder/v2/coderd/database"
14+
"github.com/coder/coder/v2/coderd/database/dbgen"
15+
"github.com/coder/coder/v2/coderd/database/dbtestutil"
16+
"github.com/coder/coder/v2/coderd/database/dbtime"
17+
"github.com/coder/coder/v2/testutil"
18+
"github.com/coder/quartz"
19+
)
20+
21+
func Test_rotateKeys(t *testing.T) {
22+
t.Parallel()
23+
24+
t.Run("NoExistingKeys", func(t *testing.T) {
25+
t.Parallel()
26+
27+
var (
28+
db, _ = dbtestutil.NewDB(t)
29+
clock = quartz.NewMock(t)
30+
logger = slogtest.Make(t, nil).Leveled(slog.LevelDebug)
31+
ctx = testutil.Context(t, testutil.WaitShort)
32+
resultsCh = make(chan []database.CryptoKey, 1)
33+
)
34+
35+
kr := &KeyRotator{
36+
DB: db,
37+
KeyDuration: 0,
38+
Clock: clock,
39+
Logger: logger,
40+
ScanInterval: 0,
41+
ResultsCh: resultsCh,
42+
}
43+
44+
now := clock.Now()
45+
keys, err := kr.rotateKeys(ctx)
46+
require.NoError(t, err)
47+
require.Len(t, keys, len(database.AllCryptoKeyFeatureValues()))
48+
49+
// Fetch the keys from the database and ensure they
50+
// are as expected.
51+
dbkeys, err := db.GetCryptoKeys(ctx)
52+
require.NoError(t, err)
53+
require.Equal(t, keys, dbkeys)
54+
requireContainsAllFeatures(t, keys)
55+
for _, key := range keys {
56+
requireKey(t, key, key.Feature, now, time.Time{}, 1)
57+
}
58+
})
59+
60+
t.Run("RotatesKeysNearExpiration", func(t *testing.T) {
61+
t.Parallel()
62+
63+
var (
64+
db, _ = dbtestutil.NewDB(t)
65+
clock = quartz.NewMock(t)
66+
keyDuration = time.Hour * 24 * 7
67+
logger = slogtest.Make(t, nil).Leveled(slog.LevelDebug)
68+
ctx = testutil.Context(t, testutil.WaitShort)
69+
resultsCh = make(chan []database.CryptoKey, 1)
70+
)
71+
72+
kr := &KeyRotator{
73+
DB: db,
74+
KeyDuration: keyDuration,
75+
Clock: clock,
76+
Logger: logger,
77+
ScanInterval: 0,
78+
ResultsCh: resultsCh,
79+
features: []database.CryptoKeyFeature{
80+
database.CryptoKeyFeatureWorkspaceApps,
81+
},
82+
}
83+
84+
now := dbtime.Time(clock.Now())
85+
86+
// Seed the database with an existing key.
87+
oldKey := dbgen.CryptoKey(t, db, database.CryptoKey{
88+
Feature: database.CryptoKeyFeatureWorkspaceApps,
89+
StartsAt: now,
90+
Sequence: 15,
91+
})
92+
93+
// Advance the window to just inside rotation time.
94+
_ = clock.Advance(keyDuration - time.Minute*59)
95+
keys, err := kr.rotateKeys(ctx)
96+
require.NoError(t, err)
97+
require.Len(t, keys, 2)
98+
99+
now = dbtime.Time(clock.Now())
100+
expectedDeletesAt := now.Add(WorkspaceAppsTokenDuration + time.Hour)
101+
102+
// Fetch the old key, it should have an expires_at now.
103+
oldKey, err = db.GetCryptoKeyByFeatureAndSequence(ctx, database.GetCryptoKeyByFeatureAndSequenceParams{
104+
Feature: oldKey.Feature,
105+
Sequence: oldKey.Sequence,
106+
Time: now,
107+
})
108+
require.NoError(t, err)
109+
require.Equal(t, oldKey.DeletesAt.Time, expectedDeletesAt)
110+
111+
// The new key should be created and have a starts_at of the old key's expires_at.
112+
newKey, err := db.GetCryptoKeyByFeatureAndSequence(ctx, database.GetCryptoKeyByFeatureAndSequenceParams{
113+
Feature: database.CryptoKeyFeatureWorkspaceApps,
114+
Sequence: oldKey.Sequence + 1,
115+
})
116+
require.NoError(t, err)
117+
requireKey(t, newKey, database.CryptoKeyFeatureWorkspaceApps, oldKey.ExpiresAt(keyDuration), expectedDeletesAt, oldKey.Sequence+1)
118+
119+
clock.Advance(oldKey.DeletesAt.Time.Sub(now) + time.Second)
120+
121+
keys, err = kr.rotateKeys(ctx)
122+
require.NoError(t, err)
123+
require.Len(t, keys, 1)
124+
125+
// The old key should be deleted.
126+
_, err = db.GetCryptoKeyByFeatureAndSequence(ctx, database.GetCryptoKeyByFeatureAndSequenceParams{
127+
Feature: oldKey.Feature,
128+
Sequence: oldKey.Sequence,
129+
})
130+
require.ErrorIs(t, err, sql.ErrNoRows)
131+
})
132+
133+
t.Run("DoesNotRotateValidKeys", func(t *testing.T) {
134+
})
135+
136+
t.Run("DeletesExpiredKeys", func(t *testing.T) {
137+
t.Parallel()
138+
// TODO: Implement test for deleting expired keys
139+
})
140+
141+
t.Run("HandlesMultipleKeyTypes", func(t *testing.T) {
142+
t.Parallel()
143+
// TODO: Implement test for handling multiple key types
144+
})
145+
146+
t.Run("GracefullyHandlesErrors", func(t *testing.T) {
147+
t.Parallel()
148+
// TODO: Implement test for error handling
149+
})
150+
}
151+
152+
func requireKey(t *testing.T, key database.CryptoKey, feature database.CryptoKeyFeature, startsAt time.Time, deletesAt time.Time, sequence int32) {
153+
t.Helper()
154+
require.Equal(t, key.Feature, feature)
155+
require.Equal(t, key.StartsAt, startsAt)
156+
require.Equal(t, key.DeletesAt.Time, deletesAt)
157+
require.Equal(t, key.Sequence, sequence)
158+
159+
secret, err := hex.DecodeString(key.Secret.String)
160+
require.NoError(t, err)
161+
162+
switch key.Feature {
163+
case database.CryptoKeyFeatureOidcConvert:
164+
require.Len(t, secret, 32)
165+
case database.CryptoKeyFeatureWorkspaceApps:
166+
require.Len(t, secret, 96)
167+
case database.CryptoKeyFeaturePeerReconnect:
168+
require.Len(t, secret, 64)
169+
default:
170+
t.Fatalf("unknown key feature: %s", key.Feature)
171+
}
172+
}
173+
174+
func requireContainsAllFeatures(t *testing.T, keys []database.CryptoKey) {
175+
t.Helper()
176+
177+
features := make(map[database.CryptoKeyFeature]bool)
178+
for _, key := range keys {
179+
features[key.Feature] = true
180+
}
181+
require.True(t, features[database.CryptoKeyFeatureOidcConvert])
182+
require.True(t, features[database.CryptoKeyFeatureWorkspaceApps])
183+
require.True(t, features[database.CryptoKeyFeaturePeerReconnect])
184+
}

‎coderd/keyrotate/rotate_test.go

Lines changed: 0 additions & 94 deletions
This file was deleted.

0 commit comments

Comments
 (0)