Skip to content

Commit 2de6cc3

Browse files
johnstcnjohnstcn
committed
beef up unit tests, refactor cli
1 parent ebf4eef commit 2de6cc3

File tree

3 files changed

+312
-126
lines changed

3 files changed

+312
-126
lines changed

‎enterprise/cli/server_dbcrypt.go

Lines changed: 180 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -3,16 +3,16 @@
33
package cli
44

55
import (
6-
"bytes"
76
"context"
87
"encoding/base64"
8+
"fmt"
9+
"strings"
910

1011
"cdr.dev/slog"
1112
"cdr.dev/slog/sloggers/sloghuman"
1213
"github.com/coder/coder/v2/cli"
1314
"github.com/coder/coder/v2/cli/clibase"
1415
"github.com/coder/coder/v2/cli/cliui"
15-
"github.com/coder/coder/v2/codersdk"
1616
"github.com/coder/coder/v2/enterprise/dbcrypt"
1717

1818
"golang.org/x/xerrors"
@@ -35,20 +35,10 @@ func (r *RootCmd) dbcryptCmd() *clibase.Cmd {
3535
}
3636

3737
func (*RootCmd) dbcryptRotateCmd() *clibase.Cmd {
38-
var (
39-
vals = new(codersdk.DeploymentValues)
40-
opts = vals.Options()
41-
)
38+
var flags rotateFlags
4239
cmd := &clibase.Cmd{
4340
Use: "rotate",
4441
Short: "Rotate database encryption keys.",
45-
Options: clibase.OptionSet{
46-
*opts.ByName("Postgres Connection URL"),
47-
*opts.ByName("External Token Encryption Keys"),
48-
},
49-
Middleware: clibase.Chain(
50-
clibase.RequireNArgs(0),
51-
),
5242
Handler: func(inv *clibase.Invocation) error {
5343
ctx, cancel := context.WithCancel(inv.Context())
5444
defer cancel()
@@ -57,38 +47,47 @@ func (*RootCmd) dbcryptRotateCmd() *clibase.Cmd {
5747
logger = logger.Leveled(slog.LevelDebug)
5848
}
5949

60-
if vals.PostgresURL == "" {
61-
return xerrors.Errorf("no database configured")
50+
if err := flags.valid(); err != nil {
51+
return err
6252
}
6353

64-
switch len(vals.ExternalTokenEncryptionKeys) {
65-
case 0:
66-
return xerrors.Errorf("no external token encryption keys provided")
67-
case 1:
68-
logger.Info(ctx, "only one key provided, data will be re-encrypted with the same key")
54+
ks := [][]byte{}
55+
dk, err := base64.StdEncoding.DecodeString(flags.New)
56+
if err != nil {
57+
return xerrors.Errorf("decode new key: %w", err)
6958
}
59+
ks = append(ks, dk)
7060

71-
keys := make([][]byte, 0, len(vals.ExternalTokenEncryptionKeys))
72-
var newKey []byte
73-
for idx, ek := range vals.ExternalTokenEncryptionKeys {
74-
dk, err := base64.StdEncoding.DecodeString(ek)
61+
for _, k := range flags.Old {
62+
dk, err := base64.StdEncoding.DecodeString(k)
7563
if err != nil {
76-
return xerrors.Errorf("key must be base64-encoded")
77-
}
78-
if idx == 0 {
79-
newKey = dk
80-
} else if bytes.Equal(dk, newKey) {
81-
return xerrors.Errorf("old key at index %d is the same as the new key", idx)
64+
return xerrors.Errorf("decode old key: %w", err)
8265
}
83-
keys = append(keys, dk)
66+
ks = append(ks, dk)
8467
}
8568

86-
ciphers, err := dbcrypt.NewCiphers(keys...)
69+
ciphers, err := dbcrypt.NewCiphers(ks...)
8770
if err != nil {
8871
return xerrors.Errorf("create ciphers: %w", err)
8972
}
9073

91-
sqlDB, err := cli.ConnectToPostgres(inv.Context(), logger, "postgres", vals.PostgresURL.Value())
74+
newDigest := ciphers[0].HexDigest()
75+
oldDigests := make([]string, 0, len(ciphers)-1)
76+
for _, c := range ciphers[1:] {
77+
oldDigests = append(oldDigests, c.HexDigest())
78+
}
79+
if len(oldDigests) == 0 {
80+
oldDigests = append(oldDigests, "none")
81+
}
82+
msg := fmt.Sprintf(`Rotate external token encryptions keys?\n- New key: %s\n- Old keys: %s`,
83+
newDigest,
84+
strings.Join(oldDigests, ", "),
85+
)
86+
if _, err := cliui.Prompt(inv, cliui.PromptOptions{Text: msg, IsConfirm: true}); err != nil {
87+
return err
88+
}
89+
90+
sqlDB, err := cli.ConnectToPostgres(inv.Context(), logger, "postgres", flags.PostgresURL)
9291
if err != nil {
9392
return xerrors.Errorf("connect to postgres: %w", err)
9493
}
@@ -103,24 +102,15 @@ func (*RootCmd) dbcryptRotateCmd() *clibase.Cmd {
103102
return nil
104103
},
105104
}
105+
flags.attach(&cmd.Options)
106106
return cmd
107107
}
108108

109109
func (*RootCmd) dbcryptDecryptCmd() *clibase.Cmd {
110-
var (
111-
vals = new(codersdk.DeploymentValues)
112-
opts = vals.Options()
113-
)
110+
var flags decryptFlags
114111
cmd := &clibase.Cmd{
115112
Use: "decrypt",
116113
Short: "Decrypt a previously encrypted database.",
117-
Options: clibase.OptionSet{
118-
*opts.ByName("Postgres Connection URL"),
119-
*opts.ByName("External Token Encryption Keys"),
120-
},
121-
Middleware: clibase.Chain(
122-
clibase.RequireNArgs(0),
123-
),
124114
Handler: func(inv *clibase.Invocation) error {
125115
ctx, cancel := context.WithCancel(inv.Context())
126116
defer cancel()
@@ -129,38 +119,32 @@ func (*RootCmd) dbcryptDecryptCmd() *clibase.Cmd {
129119
logger = logger.Leveled(slog.LevelDebug)
130120
}
131121

132-
if vals.PostgresURL == "" {
133-
return xerrors.Errorf("no database configured")
134-
}
135-
136-
switch len(vals.ExternalTokenEncryptionKeys) {
137-
case 0:
138-
return xerrors.Errorf("no external token encryption keys provided")
139-
case 1:
140-
logger.Info(ctx, "only one key provided, data will be re-encrypted with the same key")
122+
if err := flags.valid(); err != nil {
123+
return err
141124
}
142125

143-
keys := make([][]byte, 0, len(vals.ExternalTokenEncryptionKeys))
144-
var newKey []byte
145-
for idx, ek := range vals.ExternalTokenEncryptionKeys {
146-
dk, err := base64.StdEncoding.DecodeString(ek)
126+
ks := make([][]byte, 0, len(flags.Keys))
127+
for _, k := range flags.Keys {
128+
dk, err := base64.StdEncoding.DecodeString(k)
147129
if err != nil {
148-
return xerrors.Errorf("key must be base64-encoded")
130+
return xerrors.Errorf("decode key: %w", err)
149131
}
150-
if idx == 0 {
151-
newKey = dk
152-
} else if bytes.Equal(dk, newKey) {
153-
return xerrors.Errorf("old key at index %d is the same as the new key", idx)
154-
}
155-
keys = append(keys, dk)
132+
ks = append(ks, dk)
156133
}
157134

158-
ciphers, err := dbcrypt.NewCiphers(keys...)
135+
ciphers, err := dbcrypt.NewCiphers(ks...)
159136
if err != nil {
160137
return xerrors.Errorf("create ciphers: %w", err)
161138
}
162139

163-
sqlDB, err := cli.ConnectToPostgres(inv.Context(), logger, "postgres", vals.PostgresURL.Value())
140+
if _, err := cliui.Prompt(inv, cliui.PromptOptions{
141+
Text: "This will decrypt all encrypted data in the database. Are you sure you want to continue?",
142+
IsConfirm: true,
143+
}); err != nil {
144+
return err
145+
}
146+
147+
sqlDB, err := cli.ConnectToPostgres(inv.Context(), logger, "postgres", flags.PostgresURL)
164148
if err != nil {
165149
return xerrors.Errorf("connect to postgres: %w", err)
166150
}
@@ -175,23 +159,15 @@ func (*RootCmd) dbcryptDecryptCmd() *clibase.Cmd {
175159
return nil
176160
},
177161
}
162+
flags.attach(&cmd.Options)
178163
return cmd
179164
}
180165

181166
func (*RootCmd) dbcryptDeleteCmd() *clibase.Cmd {
182-
var (
183-
vals = new(codersdk.DeploymentValues)
184-
opts = vals.Options()
185-
)
167+
var flags deleteFlags
186168
cmd := &clibase.Cmd{
187169
Use: "delete",
188170
Short: "Delete all encrypted data from the database. THIS IS A DESTRUCTIVE OPERATION.",
189-
Options: clibase.OptionSet{
190-
*opts.ByName("Postgres Connection URL"),
191-
},
192-
Middleware: clibase.Chain(
193-
clibase.RequireNArgs(0),
194-
),
195171
Handler: func(inv *clibase.Invocation) error {
196172
ctx, cancel := context.WithCancel(inv.Context())
197173
defer cancel()
@@ -200,8 +176,8 @@ func (*RootCmd) dbcryptDeleteCmd() *clibase.Cmd {
200176
logger = logger.Leveled(slog.LevelDebug)
201177
}
202178

203-
if vals.PostgresURL == "" {
204-
return xerrors.Errorf("no database configured")
179+
if err := flags.valid(); err != nil {
180+
return err
205181
}
206182

207183
if _, err := cliui.Prompt(inv, cliui.PromptOptions{
@@ -211,7 +187,7 @@ func (*RootCmd) dbcryptDeleteCmd() *clibase.Cmd {
211187
return err
212188
}
213189

214-
sqlDB, err := cli.ConnectToPostgres(inv.Context(), logger, "postgres", vals.PostgresURL.Value())
190+
sqlDB, err := cli.ConnectToPostgres(inv.Context(), logger, "postgres", flags.PostgresURL)
215191
if err != nil {
216192
return xerrors.Errorf("connect to postgres: %w", err)
217193
}
@@ -226,5 +202,130 @@ func (*RootCmd) dbcryptDeleteCmd() *clibase.Cmd {
226202
return nil
227203
},
228204
}
205+
flags.attach(&cmd.Options)
229206
return cmd
230207
}
208+
209+
type rotateFlags struct {
210+
PostgresURL string
211+
New string
212+
Old []string
213+
}
214+
215+
func (f *rotateFlags) attach(opts *clibase.OptionSet) {
216+
*opts = append(
217+
*opts,
218+
clibase.Option{
219+
Flag: "postgres-url",
220+
Env: "CODER_PG_CONNECTION_URL",
221+
Description: "The connection URL for the Postgres database.",
222+
Value: clibase.StringOf(&f.PostgresURL),
223+
},
224+
clibase.Option{
225+
Flag: "new-key",
226+
Env: "CODER_EXTERNAL_TOKEN_ENCRYPTION_ENCRYPT_NEW_KEY",
227+
Description: "The new external token encryption key. Must be base64-encoded.",
228+
Value: clibase.StringOf(&f.New),
229+
},
230+
clibase.Option{
231+
Flag: "old-keys",
232+
Env: "CODER_EXTERNAL_TOKEN_ENCRYPTION_ENCRYPT_OLD_KEYS",
233+
Description: "The old external token encryption keys. Must be a comma-separated list of base64-encoded keys.",
234+
Value: clibase.StringArrayOf(&f.Old),
235+
},
236+
cliui.SkipPromptOption(),
237+
)
238+
}
239+
240+
func (f *rotateFlags) valid() error {
241+
if f.New == "" {
242+
return xerrors.Errorf("no new key provided")
243+
}
244+
245+
if val, err := base64.StdEncoding.DecodeString(f.New); err != nil {
246+
return xerrors.Errorf("new key must be base64-encoded")
247+
} else if len(val) != 32 {
248+
return xerrors.Errorf("new key must be exactly 32 bytes in length")
249+
}
250+
251+
for i, k := range f.Old {
252+
if val, err := base64.StdEncoding.DecodeString(k); err != nil {
253+
return xerrors.Errorf("old key at index %d must be base64-encoded", i)
254+
} else if len(val) != 32 {
255+
return xerrors.Errorf("old key at index %d must be exactly 32 bytes in length", i)
256+
}
257+
258+
// Pedantic, but typos here will ruin your day.
259+
if k == f.New {
260+
return xerrors.Errorf("old key at index %d is the same as the new key", i)
261+
}
262+
}
263+
264+
return nil
265+
}
266+
267+
type decryptFlags struct {
268+
PostgresURL string
269+
Keys []string
270+
}
271+
272+
func (f *decryptFlags) attach(opts *clibase.OptionSet) {
273+
*opts = append(
274+
*opts,
275+
clibase.Option{
276+
Flag: "postgres-url",
277+
Env: "CODER_PG_CONNECTION_URL",
278+
Description: "The connection URL for the Postgres database.",
279+
Value: clibase.StringOf(&f.PostgresURL),
280+
},
281+
clibase.Option{
282+
Flag: "keys",
283+
Env: "CODER_EXTERNAL_TOKEN_ENCRYPTION_DECRYPT_KEYS",
284+
Description: "Keys required to decrypt existing data. Must be a comma-separated list of base64-encoded keys.",
285+
Value: clibase.StringArrayOf(&f.Keys),
286+
},
287+
cliui.SkipPromptOption(),
288+
)
289+
}
290+
291+
func (f *decryptFlags) valid() error {
292+
if len(f.Keys) == 0 {
293+
return xerrors.Errorf("no keys provided")
294+
}
295+
296+
for i, k := range f.Keys {
297+
if val, err := base64.StdEncoding.DecodeString(k); err != nil {
298+
return xerrors.Errorf("key at index %d must be base64-encoded", i)
299+
} else if len(val) != 32 {
300+
return xerrors.Errorf("key at index %d must be exactly 32 bytes in length", i)
301+
}
302+
}
303+
304+
return nil
305+
}
306+
307+
type deleteFlags struct {
308+
PostgresURL string
309+
Confirm bool
310+
}
311+
312+
func (f *deleteFlags) attach(opts *clibase.OptionSet) {
313+
*opts = append(
314+
*opts,
315+
clibase.Option{
316+
Flag: "postgres-url",
317+
Env: "CODER_EXTERNAL_TOKEN_ENCRYPTION_POSTGRES_URL",
318+
Description: "The connection URL for the Postgres database.",
319+
Value: clibase.StringOf(&f.PostgresURL),
320+
},
321+
cliui.SkipPromptOption(),
322+
)
323+
}
324+
325+
func (f *deleteFlags) valid() error {
326+
if f.PostgresURL == "" {
327+
return xerrors.Errorf("no database configured")
328+
}
329+
330+
return nil
331+
}

0 commit comments

Comments
 (0)