beef up unit tests, refactor cli

coder · johnstcn · Sep 7, 2023 · Sep 4, 2023 · Sep 4, 2023 · Sep 5, 2023
commit 2de6cc3cb334dffb51c78faa3ba4963998b4fea3
diff --git a/enterprise/cli/server_dbcrypt.go b/enterprise/cli/server_dbcrypt.go
@@ -3,16 +3,16 @@
 package cli
 
 import (
-	"bytes"
 	"context"
 	"encoding/base64"
+	"fmt"
+	"strings"
 
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/sloghuman"
 	"github.com/coder/coder/v2/cli"
 	"github.com/coder/coder/v2/cli/clibase"
 	"github.com/coder/coder/v2/cli/cliui"
-	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/enterprise/dbcrypt"
 
 	"golang.org/x/xerrors"
@@ -35,20 +35,10 @@ func (r *RootCmd) dbcryptCmd() *clibase.Cmd {
 }
 
 func (*RootCmd) dbcryptRotateCmd() *clibase.Cmd {
-	var (
-		vals = new(codersdk.DeploymentValues)
-		opts = vals.Options()
-	)
+	var flags rotateFlags
 	cmd := &clibase.Cmd{
 		Use:   "rotate",
 		Short: "Rotate database encryption keys.",
-		Options: clibase.OptionSet{
-			*opts.ByName("Postgres Connection URL"),
-			*opts.ByName("External Token Encryption Keys"),
-		},
-		Middleware: clibase.Chain(
-			clibase.RequireNArgs(0),
-		),
 		Handler: func(inv *clibase.Invocation) error {
 			ctx, cancel := context.WithCancel(inv.Context())
 			defer cancel()
@@ -57,38 +47,47 @@ func (*RootCmd) dbcryptRotateCmd() *clibase.Cmd {
 				logger = logger.Leveled(slog.LevelDebug)
 			}
 
-			if vals.PostgresURL == "" {
-				return xerrors.Errorf("no database configured")
+			if err := flags.valid(); err != nil {
+				return err
 			}
 
-			switch len(vals.ExternalTokenEncryptionKeys) {
-			case 0:
-				return xerrors.Errorf("no external token encryption keys provided")
-			case 1:
-				logger.Info(ctx, "only one key provided, data will be re-encrypted with the same key")
+			ks := [][]byte{}
+			dk, err := base64.StdEncoding.DecodeString(flags.New)
+			if err != nil {
+				return xerrors.Errorf("decode new key: %w", err)
 			}
+			ks = append(ks, dk)
 
-			keys := make([][]byte, 0, len(vals.ExternalTokenEncryptionKeys))
-			var newKey []byte
-			for idx, ek := range vals.ExternalTokenEncryptionKeys {
-				dk, err := base64.StdEncoding.DecodeString(ek)
+			for _, k := range flags.Old {
+				dk, err := base64.StdEncoding.DecodeString(k)
 				if err != nil {
-					return xerrors.Errorf("key must be base64-encoded")
-				}
-				if idx == 0 {
-					newKey = dk
-				} else if bytes.Equal(dk, newKey) {
-					return xerrors.Errorf("old key at index %d is the same as the new key", idx)
+					return xerrors.Errorf("decode old key: %w", err)
 				}
-				keys = append(keys, dk)
+				ks = append(ks, dk)
 			}
 
-			ciphers, err := dbcrypt.NewCiphers(keys...)
+			ciphers, err := dbcrypt.NewCiphers(ks...)
 			if err != nil {
 				return xerrors.Errorf("create ciphers: %w", err)
 			}
 
-			sqlDB, err := cli.ConnectToPostgres(inv.Context(), logger, "postgres", vals.PostgresURL.Value())
+			newDigest := ciphers[0].HexDigest()
+			oldDigests := make([]string, 0, len(ciphers)-1)
+			for _, c := range ciphers[1:] {
+				oldDigests = append(oldDigests, c.HexDigest())
+			}
+			if len(oldDigests) == 0 {
+				oldDigests = append(oldDigests, "none")
+			}
+			msg := fmt.Sprintf(`Rotate external token encryptions keys?\n- New key: %s\n- Old keys: %s`,
+				newDigest,
+				strings.Join(oldDigests, ", "),
+			)
+			if _, err := cliui.Prompt(inv, cliui.PromptOptions{Text: msg, IsConfirm: true}); err != nil {
+				return err
+			}
+
+			sqlDB, err := cli.ConnectToPostgres(inv.Context(), logger, "postgres", flags.PostgresURL)
 			if err != nil {
 				return xerrors.Errorf("connect to postgres: %w", err)
 			}
@@ -103,24 +102,15 @@ func (*RootCmd) dbcryptRotateCmd() *clibase.Cmd {
 			return nil
 		},
 	}
+	flags.attach(&cmd.Options)
 	return cmd
 }
 
 func (*RootCmd) dbcryptDecryptCmd() *clibase.Cmd {
-	var (
-		vals = new(codersdk.DeploymentValues)
-		opts = vals.Options()
-	)
+	var flags decryptFlags
 	cmd := &clibase.Cmd{
 		Use:   "decrypt",
 		Short: "Decrypt a previously encrypted database.",
-		Options: clibase.OptionSet{
-			*opts.ByName("Postgres Connection URL"),
-			*opts.ByName("External Token Encryption Keys"),
-		},
-		Middleware: clibase.Chain(
-			clibase.RequireNArgs(0),
-		),
 		Handler: func(inv *clibase.Invocation) error {
 			ctx, cancel := context.WithCancel(inv.Context())
 			defer cancel()
@@ -129,38 +119,32 @@ func (*RootCmd) dbcryptDecryptCmd() *clibase.Cmd {
 				logger = logger.Leveled(slog.LevelDebug)
 			}
 
-			if vals.PostgresURL == "" {
-				return xerrors.Errorf("no database configured")
-			}
-
-			switch len(vals.ExternalTokenEncryptionKeys) {
-			case 0:
-				return xerrors.Errorf("no external token encryption keys provided")
-			case 1:
-				logger.Info(ctx, "only one key provided, data will be re-encrypted with the same key")
+			if err := flags.valid(); err != nil {
+				return err
 			}
 
-			keys := make([][]byte, 0, len(vals.ExternalTokenEncryptionKeys))
-			var newKey []byte
-			for idx, ek := range vals.ExternalTokenEncryptionKeys {
-				dk, err := base64.StdEncoding.DecodeString(ek)
+			ks := make([][]byte, 0, len(flags.Keys))
+			for _, k := range flags.Keys {
+				dk, err := base64.StdEncoding.DecodeString(k)
 				if err != nil {
-					return xerrors.Errorf("key must be base64-encoded")
+					return xerrors.Errorf("decode key: %w", err)
 				}
-				if idx == 0 {
-					newKey = dk
-				} else if bytes.Equal(dk, newKey) {
-					return xerrors.Errorf("old key at index %d is the same as the new key", idx)
-				}
-				keys = append(keys, dk)
+				ks = append(ks, dk)
 			}
 
-			ciphers, err := dbcrypt.NewCiphers(keys...)
+			ciphers, err := dbcrypt.NewCiphers(ks...)
 			if err != nil {
 				return xerrors.Errorf("create ciphers: %w", err)
 			}
 
-			sqlDB, err := cli.ConnectToPostgres(inv.Context(), logger, "postgres", vals.PostgresURL.Value())
+			if _, err := cliui.Prompt(inv, cliui.PromptOptions{
+				Text:      "This will decrypt all encrypted data in the database. Are you sure you want to continue?",
+				IsConfirm: true,
+			}); err != nil {
+				return err
+			}
+
+			sqlDB, err := cli.ConnectToPostgres(inv.Context(), logger, "postgres", flags.PostgresURL)
 			if err != nil {
 				return xerrors.Errorf("connect to postgres: %w", err)
 			}
@@ -175,23 +159,15 @@ func (*RootCmd) dbcryptDecryptCmd() *clibase.Cmd {
 			return nil
 		},
 	}
+	flags.attach(&cmd.Options)
 	return cmd
 }
 
 func (*RootCmd) dbcryptDeleteCmd() *clibase.Cmd {
-	var (
-		vals = new(codersdk.DeploymentValues)
-		opts = vals.Options()
-	)
+	var flags deleteFlags
 	cmd := &clibase.Cmd{
 		Use:   "delete",
 		Short: "Delete all encrypted data from the database. THIS IS A DESTRUCTIVE OPERATION.",
-		Options: clibase.OptionSet{
-			*opts.ByName("Postgres Connection URL"),
-		},
-		Middleware: clibase.Chain(
-			clibase.RequireNArgs(0),
-		),
 		Handler: func(inv *clibase.Invocation) error {
 			ctx, cancel := context.WithCancel(inv.Context())
 			defer cancel()
@@ -200,8 +176,8 @@ func (*RootCmd) dbcryptDeleteCmd() *clibase.Cmd {
 				logger = logger.Leveled(slog.LevelDebug)
 			}
 
-			if vals.PostgresURL == "" {
-				return xerrors.Errorf("no database configured")
+			if err := flags.valid(); err != nil {
+				return err
 			}
 
 			if _, err := cliui.Prompt(inv, cliui.PromptOptions{
@@ -211,7 +187,7 @@ func (*RootCmd) dbcryptDeleteCmd() *clibase.Cmd {
 				return err
 			}
 
-			sqlDB, err := cli.ConnectToPostgres(inv.Context(), logger, "postgres", vals.PostgresURL.Value())
+			sqlDB, err := cli.ConnectToPostgres(inv.Context(), logger, "postgres", flags.PostgresURL)
 			if err != nil {
 				return xerrors.Errorf("connect to postgres: %w", err)
 			}
@@ -226,5 +202,130 @@ func (*RootCmd) dbcryptDeleteCmd() *clibase.Cmd {
 			return nil
 		},
 	}
+	flags.attach(&cmd.Options)
 	return cmd
 }
+
+type rotateFlags struct {
+	PostgresURL string
+	New         string
+	Old         []string
+}
+
+func (f *rotateFlags) attach(opts *clibase.OptionSet) {
+	*opts = append(
+		*opts,
+		clibase.Option{
+			Flag:        "postgres-url",
+			Env:         "CODER_PG_CONNECTION_URL",
+			Description: "The connection URL for the Postgres database.",
+			Value:       clibase.StringOf(&f.PostgresURL),
+		},
+		clibase.Option{
+			Flag:        "new-key",
+			Env:         "CODER_EXTERNAL_TOKEN_ENCRYPTION_ENCRYPT_NEW_KEY",
+			Description: "The new external token encryption key. Must be base64-encoded.",
+			Value:       clibase.StringOf(&f.New),
+		},
+		clibase.Option{
+			Flag:        "old-keys",
+			Env:         "CODER_EXTERNAL_TOKEN_ENCRYPTION_ENCRYPT_OLD_KEYS",
+			Description: "The old external token encryption keys. Must be a comma-separated list of base64-encoded keys.",
+			Value:       clibase.StringArrayOf(&f.Old),
+		},
+		cliui.SkipPromptOption(),
+	)
+}
+
+func (f *rotateFlags) valid() error {
+	if f.New == "" {
+		return xerrors.Errorf("no new key provided")
+	}
+
+	if val, err := base64.StdEncoding.DecodeString(f.New); err != nil {
+		return xerrors.Errorf("new key must be base64-encoded")
+	} else if len(val) != 32 {
+		return xerrors.Errorf("new key must be exactly 32 bytes in length")
+	}
+
+	for i, k := range f.Old {
+		if val, err := base64.StdEncoding.DecodeString(k); err != nil {
+			return xerrors.Errorf("old key at index %d must be base64-encoded", i)
+		} else if len(val) != 32 {
+			return xerrors.Errorf("old key at index %d must be exactly 32 bytes in length", i)
+		}
+
+		// Pedantic, but typos here will ruin your day.
+		if k == f.New {
+			return xerrors.Errorf("old key at index %d is the same as the new key", i)
+		}
+	}
+
+	return nil
+}
+
+type decryptFlags struct {
+	PostgresURL string
+	Keys        []string
+}
+
+func (f *decryptFlags) attach(opts *clibase.OptionSet) {
+	*opts = append(
+		*opts,
+		clibase.Option{
+			Flag:        "postgres-url",
+			Env:         "CODER_PG_CONNECTION_URL",
+			Description: "The connection URL for the Postgres database.",
+			Value:       clibase.StringOf(&f.PostgresURL),
+		},
+		clibase.Option{
+			Flag:        "keys",
+			Env:         "CODER_EXTERNAL_TOKEN_ENCRYPTION_DECRYPT_KEYS",
+			Description: "Keys required to decrypt existing data. Must be a comma-separated list of base64-encoded keys.",
+			Value:       clibase.StringArrayOf(&f.Keys),
+		},
+		cliui.SkipPromptOption(),
+	)
+}
+
+func (f *decryptFlags) valid() error {
+	if len(f.Keys) == 0 {
+		return xerrors.Errorf("no keys provided")
+	}
+
+	for i, k := range f.Keys {
+		if val, err := base64.StdEncoding.DecodeString(k); err != nil {
+			return xerrors.Errorf("key at index %d must be base64-encoded", i)
+		} else if len(val) != 32 {
+			return xerrors.Errorf("key at index %d must be exactly 32 bytes in length", i)
+		}
+	}
+
+	return nil
+}
+
+type deleteFlags struct {
+	PostgresURL string
+	Confirm     bool
+}
+
+func (f *deleteFlags) attach(opts *clibase.OptionSet) {
+	*opts = append(
+		*opts,
+		clibase.Option{
+			Flag:        "postgres-url",
+			Env:         "CODER_EXTERNAL_TOKEN_ENCRYPTION_POSTGRES_URL",
+			Description: "The connection URL for the Postgres database.",
+			Value:       clibase.StringOf(&f.PostgresURL),
+		},
+		cliui.SkipPromptOption(),
+	)
+}
+
+func (f *deleteFlags) valid() error {
+	if f.PostgresURL == "" {
+		return xerrors.Errorf("no database configured")
+	}
+
+	return nil
+}