Skip to content

Commit a457307

Browse files
johnstcnjohnstcn
committed
flesh out unit test
1 parent 0c01b36 commit a457307

File tree

1 file changed

+114
-15
lines changed

1 file changed

+114
-15
lines changed

‎enterprise/cli/dbcrypt_rotate_test.go

Lines changed: 114 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,36 +1,135 @@
11
package cli_test
22

33
import (
4+
"context"
5+
"database/sql"
6+
"encoding/base64"
7+
"fmt"
8+
9+
"sync/atomic"
410
"testing"
511

6-
"github.com/coder/coder/v2/cli/clitest"
7-
"github.com/coder/coder/v2/codersdk"
8-
"github.com/coder/coder/v2/enterprise/coderd/coderdenttest"
9-
"github.com/coder/coder/v2/enterprise/coderd/license"
12+
"cdr.dev/slog/sloggers/slogtest"
13+
14+
"github.com/coder/coder/v2/coderd/database"
15+
"github.com/coder/coder/v2/coderd/database/dbgen"
16+
"github.com/coder/coder/v2/coderd/database/dbtestutil"
17+
"github.com/coder/coder/v2/coderd/database/postgres"
18+
"github.com/coder/coder/v2/cryptorand"
19+
"github.com/coder/coder/v2/enterprise/dbcrypt"
1020
"github.com/coder/coder/v2/pty/ptytest"
21+
1122
"github.com/stretchr/testify/require"
1223
)
1324

1425
func TestDBCryptRotate(t *testing.T) {
15-
t.Parallel()
26+
//nolint: paralleltest // use of t.Setenv
27+
if !dbtestutil.WillUsePostgres() {
28+
t.Skip("this test requires a postgres instance")
29+
}
30+
31+
ctx, cancel := context.WithCancel(context.Background())
32+
t.Cleanup(cancel)
33+
34+
// Setup a postgres database.
35+
connectionURL, closePg, err := postgres.Open()
36+
require.NoError(t, err)
37+
t.Cleanup(closePg)
38+
39+
sqlDB, err := sql.Open("postgres", connectionURL)
40+
require.NoError(t, err)
41+
t.Cleanup(func() {
42+
_ = sqlDB.Close()
43+
})
44+
db := database.New(sqlDB)
1645

17-
// TODO: create a test database and populate some encrypted data with cipher A
46+
// Setup an initial cipher
47+
keyA := mustString(t, 32)
48+
cA, err := dbcrypt.CipherAES256([]byte(keyA))
49+
require.NoError(t, err)
50+
cipherA := &atomic.Pointer[dbcrypt.Cipher]{}
51+
cipherB := &atomic.Pointer[dbcrypt.Cipher]{}
52+
cipherA.Store(&cA)
53+
54+
// Create an encrypted database
55+
log := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
56+
cryptdb, err := dbcrypt.New(ctx, db, &dbcrypt.Options{
57+
PrimaryCipher: cipherA,
58+
SecondaryCipher: cipherB,
59+
Logger: log,
60+
})
61+
require.NoError(t, err)
1862

19-
client, _ := coderdenttest.New(t, &coderdenttest.Options{LicenseOptions: &coderdenttest.LicenseOptions{
20-
Features: license.Features{
21-
codersdk.FeatureExternalTokenEncryption: 1,
22-
},
23-
}})
63+
// Populate the database with some data encrypted with cipher A.
64+
var users []database.User
65+
for i := 0; i < 10; i++ {
66+
usr := dbgen.User(t, cryptdb, database.User{
67+
LoginType: database.LoginTypeOIDC,
68+
})
69+
_ = dbgen.UserLink(t, cryptdb, database.UserLink{
70+
UserID: usr.ID,
71+
LoginType: usr.LoginType,
72+
OAuthAccessToken: mustString(t, 16),
73+
OAuthRefreshToken: mustString(t, 16),
74+
})
75+
_ = dbgen.GitAuthLink(t, cryptdb, database.GitAuthLink{
76+
UserID: usr.ID,
77+
ProviderID: "fake",
78+
OAuthAccessToken: mustString(t, 16),
79+
OAuthRefreshToken: mustString(t, 16),
80+
})
81+
users = append(users, usr)
82+
}
2483

2584
// Run the cmd with ciphers B,A
85+
keyB := mustString(t, 32)
86+
cB, err := dbcrypt.CipherAES256([]byte(keyB))
87+
require.NoError(t, err)
88+
externalTokensArg := fmt.Sprintf(
89+
"%s,%s",
90+
base64.StdEncoding.EncodeToString([]byte(keyB)),
91+
base64.StdEncoding.EncodeToString([]byte(keyA)),
92+
)
2693

27-
inv, conf := newCLI(t, "dbcrypt-rotate") // TODO: env?
94+
inv, _ := newCLI(t, "dbcrypt-rotate",
95+
"--postgres-url", connectionURL,
96+
"--external-token-encryption-keys", externalTokensArg,
97+
)
2898
pty := ptytest.New(t)
2999
inv.Stdout = pty.Output()
30-
clitest.SetupConfig(t, client, conf)
31100

32-
err := inv.Run()
101+
err = inv.Run()
33102
require.NoError(t, err)
34103

35-
// TODO: validate that all data has been updated with the checksum of the new cipher.
104+
// Validate that all data has been updated with the checksum of the new cipher.
105+
expectedPrefixA := fmt.Sprintf("dbcrypt-%s-", cA.HexDigest()[:7])
106+
expectedPrefixB := fmt.Sprintf("dbcrypt-%s-", cB.HexDigest()[:7])
107+
for _, usr := range users {
108+
ul, err := db.GetUserLinkByUserIDLoginType(ctx, database.GetUserLinkByUserIDLoginTypeParams{
109+
UserID: usr.ID,
110+
LoginType: usr.LoginType,
111+
})
112+
require.NoError(t, err, "failed to get user link for user %s", usr.ID)
113+
require.NotContains(t, ul.OAuthAccessToken, expectedPrefixA, "user_link.oauth_access_token should not contain the old cipher checksum")
114+
require.NotContains(t, ul.OAuthRefreshToken, expectedPrefixA, "user_link.oauth_refresh_token should not contain the old cipher checksum")
115+
require.Contains(t, ul.OAuthAccessToken, expectedPrefixB, "user_link.oauth_access_token should contain the new cipher checksum")
116+
require.Contains(t, ul.OAuthRefreshToken, expectedPrefixB, "user_link.oauth_refresh_token should contain the new cipher checksum")
117+
118+
gal, err := db.GetGitAuthLink(ctx, database.GetGitAuthLinkParams{
119+
UserID: usr.ID,
120+
ProviderID: "fake",
121+
})
122+
require.NoError(t, err, "failed to get git auth link for user %s", usr.ID)
123+
require.NotContains(t, gal.OAuthAccessToken, expectedPrefixA, "git_auth_link.oauth_access_token should not contain the old cipher checksum")
124+
require.NotContains(t, gal.OAuthRefreshToken, expectedPrefixA, "git_auth_link.oauth_refresh_token should not contain the old cipher checksum")
125+
require.Contains(t, gal.OAuthAccessToken, expectedPrefixB, "git_auth_link.oauth_access_token should contain the new cipher checksum")
126+
require.Contains(t, gal.OAuthRefreshToken, expectedPrefixB, "git_auth_link.oauth_refresh_token should contain the new cipher checksum")
127+
}
128+
}
129+
130+
func mustString(t *testing.T, n int) string {
131+
t.Helper()
132+
s, err := cryptorand.String(n)
133+
require.NoError(t, err)
134+
return s
36135
}

0 commit comments

Comments
 (0)