fix: bug with subkeys, verify test fields on startup (#2444)

Fixes #2348
Fixes #2466

backend/controller/dal/dal.go

@@ -221,6 +221,9 @@ func New(ctx context.Context, conn *stdsql.DB, encryptionBuilder encryption.Buil
 		return nil, fmt.Errorf("failed to setup encryptor: %w", err)
 	}
 	d.encryptor = encryptor
+	if err = d.verifyEncryptor(ctx); err != nil {
+		return nil, fmt.Errorf("failed to verify encryption: %w", err)
+	}
 
 	return d, nil
 }

backend/controller/dal/dal_test.go

@@ -542,3 +542,81 @@ func TestDeleteOldEvents(t *testing.T) {
 		assert.Equal(t, int64(0), count)
 	})
 }
+
+func TestVerifyEncryption(t *testing.T) {
+	ctx := log.ContextWithNewDefaultLogger(context.Background())
+	conn := sqltest.OpenForTesting(ctx, t)
+	uri := "fake-kms://CK6YwYkBElQKSAowdHlwZS5nb29nbGVhcGlzLmNvbS9nb29nbGUuY3J5cHRvLnRpbmsuQWVzR2NtS2V5EhIaEJy4TIQgfCuwxA3ZZgChp_wYARABGK6YwYkBIAE"
+
+	t.Run("DeleteVerificationColumns", func(t *testing.T) {
+		dal, err := New(ctx, conn, encryption.NewBuilder().WithKMSURI(optional.Some(uri)))
+		assert.NoError(t, err)
+
+		// check that there are columns set in encryption_keys
+		row, err := dal.db.GetOnlyEncryptionKey(ctx)
+		assert.NoError(t, err)
+		assert.NotZero(t, row.VerifyTimeline.Ok())
+		assert.NotZero(t, row.VerifyAsync.Ok())
+
+		// delete the columns to see if they are recreated
+		err = dal.db.UpdateEncryptionVerification(ctx, optional.None[encryption.EncryptedTimelineColumn](), optional.None[encryption.EncryptedAsyncColumn]())
+		assert.NoError(t, err)
+
+		dal, err = New(ctx, conn, encryption.NewBuilder().WithKMSURI(optional.Some(uri)))
+		assert.NoError(t, err)
+
+		row, err = dal.db.GetOnlyEncryptionKey(ctx)
+		assert.NoError(t, err)
+		assert.NotZero(t, row.VerifyTimeline.Ok())
+		assert.NotZero(t, row.VerifyAsync.Ok())
+	})
+
+	t.Run("DifferentKey", func(t *testing.T) {
+		_, err := New(ctx, conn, encryption.NewBuilder().WithKMSURI(optional.Some(uri)))
+		assert.NoError(t, err)
+
+		differentKey := "fake-kms://CJP7ksIKElQKSAowdHlwZS5nb29nbGVhcGlzLmNvbS9nb29nbGUuY3J5cHRvLnRpbmsuQWVzR2NtS2V5EhIaEJWT3z-xdW23HO7hc9vF3YoYARABGJP7ksIKIAE"
+		_, err = New(ctx, conn, encryption.NewBuilder().WithKMSURI(optional.Some(differentKey)))
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "decryption failed")
+	})
+
+	t.Run("SameKeyButWrongTimelineVerification", func(t *testing.T) {
+		dal, err := New(ctx, conn, encryption.NewBuilder().WithKMSURI(optional.Some(uri)))
+		assert.NoError(t, err)
+
+		err = dal.db.UpdateEncryptionVerification(ctx, optional.Some[encryption.EncryptedTimelineColumn]([]byte("123")), optional.None[encryption.EncryptedAsyncColumn]())
+		assert.NoError(t, err)
+		_, err = New(ctx, conn, encryption.NewBuilder().WithKMSURI(optional.Some(uri)))
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "verification sanity")
+		assert.Contains(t, err.Error(), "verify timeline")
+
+		err = dal.db.UpdateEncryptionVerification(ctx, optional.None[encryption.EncryptedTimelineColumn](), optional.Some[encryption.EncryptedAsyncColumn]([]byte("123")))
+		assert.NoError(t, err)
+		_, err = New(ctx, conn, encryption.NewBuilder().WithKMSURI(optional.Some(uri)))
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "verification sanity")
+		assert.Contains(t, err.Error(), "verify async")
+	})
+
+	t.Run("SameKeyButEncryptWrongPlainText", func(t *testing.T) {
+		result, err := conn.Exec("DELETE FROM encryption_keys")
+		assert.NoError(t, err)
+		affected, err := result.RowsAffected()
+		assert.NoError(t, err)
+		assert.Equal(t, int64(1), affected)
+		dal, err := New(ctx, conn, encryption.NewBuilder().WithKMSURI(optional.Some(uri)))
+		assert.NoError(t, err)
+
+		encrypted := encryption.EncryptedColumn[encryption.TimelineSubKey]{}
+		err = dal.encrypt([]byte("123"), &encrypted)
+		assert.NoError(t, err)
+
+		err = dal.db.UpdateEncryptionVerification(ctx, optional.Some(encrypted), optional.None[encryption.EncryptedAsyncColumn]())
+		assert.NoError(t, err)
+		_, err = New(ctx, conn, encryption.NewBuilder().WithKMSURI(optional.Some(uri)))
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "string does not match")
+	})
+}

backend/controller/dal/encryption.go

@@ -5,6 +5,8 @@ import (
 	"encoding/json"
 	"fmt"
 
+	"github.com/alecthomas/types/optional"
+
 	"github.com/TBD54566975/ftl/backend/dal"
 	"github.com/TBD54566975/ftl/internal/encryption"
 	"github.com/TBD54566975/ftl/internal/log"
@@ -66,10 +68,11 @@ func (d *DAL) EnsureKey(ctx context.Context, generateKey func() ([]byte, error))
 	}
 	defer tx.CommitOrRollback(ctx, &err)
 
-	encryptedKey, err = tx.db.GetOnlyEncryptionKey(ctx)
+	var key []byte
+	row, err := tx.db.GetOnlyEncryptionKey(ctx)
 	if err != nil && dal.IsNotFound(err) {
 		logger.Debugf("No encryption key found, generating a new one")
-		key, err := generateKey()
+		key, err = generateKey()
 		if err != nil {
 			return nil, fmt.Errorf("failed to generate key: %w", err)
 		}
@@ -84,5 +87,85 @@ func (d *DAL) EnsureKey(ctx context.Context, generateKey func() ([]byte, error))
 	}
 
 	logger.Debugf("Encryption key found, using it")
-	return encryptedKey, nil
+
+	return row.Key, nil
+}
+
+const verification = "FTL - Towards a 𝝺-calculus for large-scale systems"
+
+func (d *DAL) verifyEncryptor(ctx context.Context) (err error) {
+	var tx *Tx
+	tx, err = d.Begin(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to begin transaction: %w", err)
+	}
+	defer tx.CommitOrRollback(ctx, &err)
+
+	row, err := tx.db.GetOnlyEncryptionKey(ctx)
+	if err != nil {
+		if dal.IsNotFound(err) {
+			// No encryption key found, probably using noop.
+			return nil
+		}
+		return fmt.Errorf("failed to get encryption row from the db: %w", err)
+	}
+
+	needsUpdate := false
+	newTimeline, err := verifySubkey(d.encryptor, row.VerifyTimeline)
+	if err != nil {
+		return fmt.Errorf("failed to verify timeline subkey: %w", err)
+	}
+	if newTimeline != nil {
+		needsUpdate = true
+		row.VerifyTimeline = optional.Some(newTimeline)
+	}
+
+	newAsync, err := verifySubkey(d.encryptor, row.VerifyAsync)
+	if err != nil {
+		return fmt.Errorf("failed to verify async subkey: %w", err)
+	}
+	if newAsync != nil {
+		needsUpdate = true
+		row.VerifyAsync = optional.Some(newAsync)
+	}
+
+	if !needsUpdate {
+		return nil
+	}
+
+	if !row.VerifyTimeline.Ok() || !row.VerifyAsync.Ok() {
+		panic("should be unreachable. verifySubkey should have set the subkey")
+	}
+
+	err = tx.db.UpdateEncryptionVerification(ctx, row.VerifyTimeline, row.VerifyAsync)
+	if err != nil {
+		return fmt.Errorf("failed to update encryption verification: %w", err)
+	}
+
+	return nil
+}
+
+// verifySubkey checks if the subkey is set and if not, sets it to a verification string.
+// returns (nil, nil) if verified and not changed
+func verifySubkey[SK encryption.SubKey](encryptor encryption.DataEncryptor, encrypted optional.Option[encryption.EncryptedColumn[SK]]) (encryption.EncryptedColumn[SK], error) {
+	verifyField, ok := encrypted.Get()
+	if !ok {
+		err := encryptor.Encrypt([]byte(verification), &verifyField)
+		if err != nil {
+			return nil, fmt.Errorf("failed to encrypt verification sanity string: %w", err)
+		}
+		return verifyField, nil
+	}
+
+	decrypted, err := encryptor.Decrypt(&verifyField)
+	if err != nil {
+		return nil, fmt.Errorf("failed to decrypt verification sanity string: %w", err)
+	}
+
+	if string(decrypted) != verification {
+		return nil, fmt.Errorf("decrypted verification string does not match expected value")
+	}
+
+	// verified, no need to update
+	return nil, nil
 }

backend/controller/sql/queries.sql

@@ -901,10 +901,16 @@ FROM topic_events
 WHERE id = $1::BIGINT;
 
 -- name: GetOnlyEncryptionKey :one
-SELECT key
+SELECT key, verify_timeline, verify_async
 FROM encryption_keys
 WHERE id = 1;
 
 -- name: CreateOnlyEncryptionKey :exec
 INSERT INTO encryption_keys (id, key)
 VALUES (1, $1);
+
+-- name: UpdateEncryptionVerification :exec
+UPDATE encryption_keys
+SET verify_timeline = $1,
+    verify_async = $2
+WHERE id = 1;

backend/controller/sql/schema/20240820011612_encryption_verification.sql

@@ -0,0 +1,7 @@
+-- migrate:up
+
+ALTER TABLE encryption_keys
+    ADD COLUMN verify_timeline encrypted_timeline,
+    ADD COLUMN verify_async encrypted_async;
+
+-- migrate:down

internal/encryption/database.go

@@ -50,7 +50,7 @@ type SubKey interface{ SubKey() string }
 // TimelineSubKey is a type that represents the subkey for logs.
 type TimelineSubKey struct{}
 
-func (TimelineSubKey) SubKey() string { return "logs" }
+func (TimelineSubKey) SubKey() string { return "timeline" }
 
 // AsyncSubKey is a type that represents the subkey for async.
 type AsyncSubKey struct{}

internal/encryption/encryption.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"sync"
 
 	"github.com/alecthomas/types/optional"
 	awsv1kms "github.com/aws/aws-sdk-go/service/kms"
@@ -99,6 +100,7 @@ type KMSEncryptor struct {
 	root            keyset.Handle
 	kekAEAD         tink.AEAD
 	encryptedKeyset []byte
+	cachedDerivedMu sync.RWMutex
 	cachedDerived   map[SubKey]tink.AEAD
 }
 
@@ -206,7 +208,10 @@ func deriveKeyset(root keyset.Handle, salt []byte) (*keyset.Handle, error) {
 }
 
 func (k *KMSEncryptor) getDerivedPrimitive(subKey SubKey) (tink.AEAD, error) {
-	if primitive, ok := k.cachedDerived[subKey]; ok {
+	k.cachedDerivedMu.RLock()
+	primitive, ok := k.cachedDerived[subKey]
+	k.cachedDerivedMu.RUnlock()
+	if ok {
 		return primitive, nil
 	}
 
@@ -215,12 +220,15 @@ func (k *KMSEncryptor) getDerivedPrimitive(subKey SubKey) (tink.AEAD, error) {
 		return nil, fmt.Errorf("failed to derive keyset: %w", err)
 	}
 
-	primitive, err := aead.New(derived)
+	primitive, err = aead.New(derived)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create primitive: %w", err)
 	}
 
+	k.cachedDerivedMu.Lock()
 	k.cachedDerived[subKey] = primitive
+	k.cachedDerivedMu.Unlock()
+
 	return primitive, nil
 }