refactor: add Ciphers to abstract over multiple ciphers

johnstcn · johnstcn · commit 5a0161c4789f · 2023-08-29T08:58:36.000Z
diff --git a/enterprise/dbcrypt/cipher.go b/enterprise/dbcrypt/cipher.go
@@ -30,13 +30,14 @@ func CipherAES256(key []byte) (Cipher, error) {
 	if err != nil {
 		return nil, err
 	}
-	digest := sha256.Sum256(key)
-	return &aes256{aead: aead, digest: digest[:]}, nil
+	digest := fmt.Sprintf("%x", sha256.Sum256(key))[:7]
+	return &aes256{aead: aead, digest: digest}, nil
 }
 
 type aes256 struct {
-	aead   cipher.AEAD
-	digest []byte
+	aead cipher.AEAD
+	// digest is the first 7 bytes of the hex-encoded SHA-256 digest of aead.
+	digest string
 }
 
 func (a *aes256) Encrypt(plaintext []byte) ([]byte, error) {
@@ -60,5 +61,54 @@ func (a *aes256) Decrypt(ciphertext []byte) ([]byte, error) {
 }
 
 func (a *aes256) HexDigest() string {
-	return fmt.Sprintf("%x", a.digest)
+	return a.digest
+}
+
+type CipherDigest string
+type Ciphers struct {
+	primary string
+	m       map[string]Cipher
+}
+
+// CiphersAES256 returns a new Ciphers instance with the given ciphers.
+// The first cipher in the list is the primary cipher. Any ciphers after the
+// first are considered secondary ciphers and are only used for decryption.
+func CiphersAES256(cs ...Cipher) Ciphers {
+	var primary string
+	m := make(map[string]Cipher)
+	for idx, c := range cs {
+		m[c.HexDigest()] = c
+		if idx == 0 {
+			primary = c.HexDigest()
+		}
+	}
+	return Ciphers{primary: primary, m: m}
+}
+
+// Encrypt encrypts the given plaintext using the primary cipher and returns the
+// ciphertext. The ciphertext is prefixed with the primary cipher's digest.
+func (cs Ciphers) Encrypt(plaintext []byte) ([]byte, error) {
+	c, ok := cs.m[cs.primary]
+	if !ok {
+		return nil, xerrors.Errorf("no ciphers configured")
+	}
+	prefix := []byte(c.HexDigest() + "-")
+	crypted, err := c.Encrypt(plaintext)
+	if err != nil {
+		return nil, err
+	}
+	return append(prefix, crypted...), nil
+}
+
+// Decrypt decrypts the given ciphertext using the cipher indicated by the
+// ciphertext's prefix. The prefix is the first 7 bytes of the hex-encoded
+// SHA-256 digest of the cipher's key. Decryption will fail if the prefix
+// does not match any of the configured ciphers.
+func (cs Ciphers) Decrypt(ciphertext []byte) ([]byte, error) {
+	requiredPrefix := string(ciphertext[:7])
+	c, ok := cs.m[requiredPrefix]
+	if !ok {
+		return nil, xerrors.Errorf("missing required decryption cipher %s", requiredPrefix)
+	}
+	return c.Decrypt(ciphertext[8:])
 }
diff --git a/enterprise/dbcrypt/cipher_test.go b/enterprise/dbcrypt/cipher_test.go
@@ -43,3 +43,43 @@ func TestCipherAES256(t *testing.T) {
 		require.ErrorContains(t, err, "key must be 32 bytes")
 	})
 }
+
+func TestCiphersAES256(t *testing.T) {
+	t.Parallel()
+
+	// Given: two ciphers
+	key1 := bytes.Repeat([]byte{'a'}, 32)
+	key2 := bytes.Repeat([]byte{'b'}, 32)
+	cipher1, err := dbcrypt.CipherAES256(key1)
+	require.NoError(t, err)
+	cipher2, err := dbcrypt.CipherAES256(key2)
+	require.NoError(t, err)
+
+	ciphers := dbcrypt.CiphersAES256(
+		cipher1,
+		cipher2,
+	)
+
+	// Then: it should encrypt with the cipher1
+	output, err := ciphers.Encrypt([]byte("hello world"))
+	require.NoError(t, err)
+	// The first 7 bytes of the output should be the hex digest of cipher1
+	require.Equal(t, cipher1.HexDigest(), string(output[:7]))
+
+	// And: it should decrypt successfully
+	decrypted, err := ciphers.Decrypt(output)
+	require.NoError(t, err)
+	require.Equal(t, "hello world", string(decrypted))
+
+	// Decryption of the above should fail with cipher2
+	_, err = cipher2.Decrypt(output)
+	var decryptErr *dbcrypt.DecryptFailedError
+	require.ErrorAs(t, err, &decryptErr)
+
+	// Decryption of data encrypted with cipher2 should succeed
+	output2, err := cipher2.Encrypt([]byte("hello world"))
+	require.NoError(t, err)
+	decrypted2, err := ciphers.Decrypt(bytes.Join([][]byte{[]byte(cipher2.HexDigest()), output2}, []byte{'-'}))
+	require.NoError(t, err)
+	require.Equal(t, "hello world", string(decrypted2))
+}
