Skip to content

Commit 6c28ce5

Browse files
johnstcnjohnstcn
committed
Merge branch 'cj/dbcrypt_redux_1' into cj/dbcrypt_redux_2
2 parents 3b8140b + 9c18168 commit 6c28ce5

File tree

5 files changed

+134
-54
lines changed

5 files changed

+134
-54
lines changed

‎coderd/database/dbfake/dbfake.go

Lines changed: 18 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -672,26 +672,6 @@ func (q *FakeQuerier) isEveryoneGroup(id uuid.UUID) bool {
672672
return false
673673
}
674674

675-
func (q *FakeQuerier) insertDBCryptKeyNoLock(_ context.Context, arg database.InsertDBCryptKeyParams) error {
676-
err := validateDatabaseType(arg)
677-
if err != nil {
678-
return err
679-
}
680-
681-
for _, key := range q.dbcryptKeys {
682-
if key.Number == arg.Number {
683-
return errDuplicateKey
684-
}
685-
}
686-
687-
q.dbcryptKeys = append(q.dbcryptKeys, database.DBCryptKey{
688-
Number: arg.Number,
689-
ActiveKeyDigest: sql.NullString{String: arg.ActiveKeyDigest, Valid: true},
690-
Test: arg.Test,
691-
})
692-
return nil
693-
}
694-
695675
func (q *FakeQuerier) GetActiveDBCryptKeys(_ context.Context) ([]database.DBCryptKey, error) {
696676
q.mutex.RLock()
697677
defer q.mutex.RUnlock()
@@ -3918,9 +3898,24 @@ func (q *FakeQuerier) InsertAuditLog(_ context.Context, arg database.InsertAudit
39183898
return alog, nil
39193899
}
39203900

3921-
func (q *FakeQuerier) InsertDBCryptKey(ctx context.Context, arg database.InsertDBCryptKeyParams) error {
3922-
// This only ever gets called inside a transaction, so we need to not lock.
3923-
return q.insertDBCryptKeyNoLock(ctx, arg)
3901+
func (q *FakeQuerier) InsertDBCryptKey(_ context.Context, arg database.InsertDBCryptKeyParams) error {
3902+
err := validateDatabaseType(arg)
3903+
if err != nil {
3904+
return err
3905+
}
3906+
3907+
for _, key := range q.dbcryptKeys {
3908+
if key.Number == arg.Number {
3909+
return errDuplicateKey
3910+
}
3911+
}
3912+
3913+
q.dbcryptKeys = append(q.dbcryptKeys, database.DBCryptKey{
3914+
Number: arg.Number,
3915+
ActiveKeyDigest: sql.NullString{String: arg.ActiveKeyDigest, Valid: true},
3916+
Test: arg.Test,
3917+
})
3918+
return nil
39243919
}
39253920

39263921
func (q *FakeQuerier) InsertDERPMeshKey(_ context.Context, id string) error {

‎coderd/database/queries.sql.go

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

‎coderd/database/queries/gitauth.sql

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,9 @@ INSERT INTO git_auth_links (
1111
created_at,
1212
updated_at,
1313
oauth_access_token,
14-
oauth_access_token_key_id,
14+
oauth_access_token_key_id,
1515
oauth_refresh_token,
16-
oauth_refresh_token_key_id,
16+
oauth_refresh_token_key_id,
1717
oauth_expiry
1818
) VALUES (
1919
$1,
@@ -23,16 +23,16 @@ INSERT INTO git_auth_links (
2323
$5,
2424
$6,
2525
$7,
26-
$8,
27-
$9
26+
$8,
27+
$9
2828
) RETURNING *;
2929

3030
-- name: UpdateGitAuthLink :one
3131
UPDATE git_auth_links SET
3232
updated_at = $3,
3333
oauth_access_token = $4,
34-
oauth_access_token_key_id = $5,
34+
oauth_access_token_key_id = $5,
3535
oauth_refresh_token = $6,
36-
oauth_refresh_token_key_id = $7,
36+
oauth_refresh_token_key_id = $7,
3737
oauth_expiry = $8
3838
WHERE provider_id = $1 AND user_id = $2 RETURNING *;

‎enterprise/dbcrypt/dbcrypt.go

Lines changed: 37 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,6 @@ import (
55
"database/sql"
66
"encoding/base64"
77

8-
"github.com/lib/pq"
9-
108
"github.com/coder/coder/v2/coderd/database"
119
"github.com/coder/coder/v2/coderd/database/dbauthz"
1210

@@ -49,7 +47,8 @@ func New(ctx context.Context, db database.Store, ciphers ...Cipher) (database.St
4947
Store: db,
5048
}
5149
// nolint: gocritic // This is allowed.
52-
if err := dbc.ensureEncrypted(dbauthz.AsSystemRestricted(ctx)); err != nil {
50+
authCtx := dbauthz.AsSystemRestricted(ctx)
51+
if err := dbc.ensureEncryptedWithRetry(authCtx); err != nil {
5352
return nil, xerrors.Errorf("ensure encrypted database fields: %w", err)
5453
}
5554
return dbc, nil
@@ -271,7 +270,10 @@ func (db *dbCrypt) encryptField(field *string, digest *sql.NullString) error {
271270
}
272271

273272
if field == nil {
274-
return nil
273+
return xerrors.Errorf("developer error: encryptField called with nil field")
274+
}
275+
if digest == nil {
276+
return xerrors.Errorf("developer error: encryptField called with nil digest")
275277
}
276278

277279
encrypted, err := db.ciphers[db.primaryCipherDigest].Encrypt([]byte(*field))
@@ -334,6 +336,23 @@ func (db *dbCrypt) decryptField(field *string, digest sql.NullString) error {
334336
return nil
335337
}
336338

339+
func (db *dbCrypt) ensureEncryptedWithRetry(ctx context.Context) error {
340+
var err error
341+
for i := 0; i < 3; i++ {
342+
err = db.ensureEncrypted(ctx)
343+
if err == nil {
344+
return nil
345+
}
346+
// If we get a serialization error, then we need to retry.
347+
if !database.IsSerializedError(err) {
348+
return err
349+
}
350+
// otherwise, retry
351+
}
352+
// If we get here, then we ran out of retries
353+
return err
354+
}
355+
337356
func (db *dbCrypt) ensureEncrypted(ctx context.Context) error {
338357
return db.InTx(func(s database.Store) error {
339358
// Attempt to read the encrypted test fields of the currently active keys.
@@ -343,32 +362,31 @@ func (db *dbCrypt) ensureEncrypted(ctx context.Context) error {
343362
}
344363

345364
var highestNumber int32
365+
var activeCipherFound bool
346366
for _, k := range ks {
367+
// If our primary key has been revoked, then we can't do anything.
368+
if k.RevokedKeyDigest.Valid && k.RevokedKeyDigest.String == db.primaryCipherDigest {
369+
return xerrors.Errorf("primary encryption key %q has been revoked", db.primaryCipherDigest)
370+
}
371+
347372
if k.ActiveKeyDigest.Valid && k.ActiveKeyDigest.String == db.primaryCipherDigest {
348-
// This is our currently active key. We don't need to do anything further.
349-
return nil
373+
activeCipherFound = true
350374
}
375+
351376
if k.Number > highestNumber {
352377
highestNumber = k.Number
353378
}
354379
}
355380

381+
if activeCipherFound {
382+
return nil
383+
}
384+
356385
// If we get here, then we have a new key that we need to insert.
357-
// If this conflicts with another transaction, we do not need to retry as
358-
// the other transaction will have inserted the key for us.
359-
if err := db.InsertDBCryptKey(ctx, database.InsertDBCryptKeyParams{
386+
return db.InsertDBCryptKey(ctx, database.InsertDBCryptKeyParams{
360387
Number: highestNumber + 1,
361388
ActiveKeyDigest: db.primaryCipherDigest,
362389
Test: testValue,
363-
}); err != nil {
364-
var pqErr *pq.Error
365-
if xerrors.As(err, &pqErr) && pqErr.Code == "23505" {
366-
// Unique constraint violation -> another transaction has inserted the key for us.
367-
return nil
368-
}
369-
return err
370-
}
371-
372-
return nil
390+
})
373391
}, &sql.TxOptions{Isolation: sql.LevelRepeatableRead})
374392
}

‎enterprise/dbcrypt/dbcrypt_internal_test.go

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,13 @@ import (
88
"io"
99
"testing"
1010

11+
"github.com/golang/mock/gomock"
12+
"github.com/lib/pq"
1113
"github.com/stretchr/testify/require"
1214

1315
"github.com/coder/coder/v2/coderd/database"
1416
"github.com/coder/coder/v2/coderd/database/dbgen"
17+
"github.com/coder/coder/v2/coderd/database/dbmock"
1518
"github.com/coder/coder/v2/coderd/database/dbtestutil"
1619
)
1720

@@ -446,6 +449,70 @@ func TestNew(t *testing.T) {
446449
require.NoError(t, err, "no error should be returned")
447450
require.Empty(t, keys, "no keys should be present")
448451
})
452+
453+
t.Run("PrimaryRevoked", func(t *testing.T) {
454+
t.Parallel()
455+
// Given: a cipher is loaded
456+
cipher := initCipher(t)
457+
ctx, cancel := context.WithCancel(context.Background())
458+
t.Cleanup(cancel)
459+
rawDB, _ := dbtestutil.NewDB(t)
460+
461+
// And: the cipher is revoked before we init the crypt db
462+
err := rawDB.InsertDBCryptKey(ctx, database.InsertDBCryptKeyParams{
463+
Number: 1,
464+
ActiveKeyDigest: cipher.HexDigest(),
465+
Test: fakeBase64RandomData(t, 32),
466+
})
467+
require.NoError(t, err, "no error should be returned")
468+
err = rawDB.RevokeDBCryptKey(ctx, cipher.HexDigest())
469+
require.NoError(t, err, "no error should be returned")
470+
471+
// Then: when we init the crypt db, we error because the key is revoked
472+
_, err = New(ctx, rawDB, cipher)
473+
require.Error(t, err)
474+
require.ErrorContains(t, err, "has been revoked")
475+
})
476+
477+
t.Run("Retry", func(t *testing.T) {
478+
t.Parallel()
479+
// Given: a cipher is loaded
480+
cipher := initCipher(t)
481+
ctx, cancel := context.WithCancel(context.Background())
482+
testVal, err := cipher.Encrypt([]byte("coder"))
483+
key := database.DBCryptKey{
484+
Number: 1,
485+
ActiveKeyDigest: sql.NullString{String: cipher.HexDigest(), Valid: true},
486+
Test: b64encode(testVal),
487+
}
488+
require.NoError(t, err)
489+
t.Cleanup(cancel)
490+
491+
// And: a database that returns an error once when we try to serialize a key
492+
ctrl := gomock.NewController(t)
493+
mockDB := dbmock.NewMockStore(ctrl)
494+
495+
gomock.InOrder(
496+
// First try: we get a serialization error.
497+
expectInTx(mockDB),
498+
mockDB.EXPECT().GetDBCryptKeys(gomock.Any()).Times(1).Return([]database.DBCryptKey{}, nil),
499+
mockDB.EXPECT().InsertDBCryptKey(gomock.Any(), gomock.Any()).Times(1).Return(&pq.Error{Code: "40001"}),
500+
// Second try: we get the key we wanted to insert initially.
501+
expectInTx(mockDB),
502+
mockDB.EXPECT().GetDBCryptKeys(gomock.Any()).Times(1).Return([]database.DBCryptKey{key}, nil),
503+
)
504+
505+
_, err = New(ctx, mockDB, cipher)
506+
require.NoError(t, err)
507+
})
508+
}
509+
510+
func expectInTx(mdb *dbmock.MockStore) *gomock.Call {
511+
return mdb.EXPECT().InTx(gomock.Any(), gomock.Any()).Times(1).DoAndReturn(
512+
func(f func(store database.Store) error, _ *sql.TxOptions) error {
513+
return f(mdb)
514+
},
515+
)
449516
}
450517

451518
func requireEncryptedEquals(t *testing.T, c Cipher, value, expected string) {

0 commit comments

Comments
 (0)