bugfix

internal/jobs/handlers/secret_reencrypt_test.go

@@ -0,0 +1,173 @@
+package handlers
+
+import (
+	"context"
+	"crypto/rand"
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/yan1h/agent-coding-workflow/internal/infra/crypto"
+	"github.com/yan1h/agent-coding-workflow/internal/jobs"
+)
+
+// memRow models one encrypted-column row across versions.
+type memRow struct {
+	blob       []byte
+	keyVersion int
+	plaintext  []byte // ground truth for assertions
+	null       bool
+}
+
+// memStore is an in-memory ReencryptStore over a single logical table, with a
+// failNext hook to model a crash mid-run for the resumability test.
+type memStore struct {
+	rows     []*memRow
+	failNext bool
+}
+
+func (m *memStore) Tables() []string { return []string{"t"} }
+
+func (m *memStore) StaleCount(_ context.Context, _ string, active int) (int, error) {
+	n := 0
+	for _, r := range m.rows {
+		if !r.null && r.keyVersion < active {
+			n++
+		}
+	}
+	return n, nil
+}
+
+func (m *memStore) ReencryptTable(ctx context.Context, enc *crypto.KeyedEncryptor, _ string, active, batch int) (int, error) {
+	done := 0
+	for _, r := range m.rows {
+		if done >= batch {
+			break
+		}
+		if r.null || r.keyVersion >= active {
+			continue
+		}
+		if m.failNext {
+			m.failNext = false
+			return done, context.Canceled // simulate a crash after some rows
+		}
+		plain, err := enc.DecryptVersion(ctx, r.blob, r.keyVersion)
+		if err != nil {
+			return done, err
+		}
+		sealed, err := enc.SealWithVersion(ctx, active, plain)
+		if err != nil {
+			return done, err
+		}
+		r.blob = sealed
+		r.keyVersion = active
+		done++
+	}
+	return done, nil
+}
+
+func key(t *testing.T) []byte {
+	t.Helper()
+	k := make([]byte, 32)
+	_, err := rand.Read(k)
+	require.NoError(t, err)
+	return k
+}
+
+// twoVersionEnc returns a KeyedEncryptor with v1+v2 keys and a store reporting
+// active=2, plus a helper to seal under v1.
+func twoVersionEnc(t *testing.T) *crypto.KeyedEncryptor {
+	prov := crypto.NewEnvProvider(nil)
+	prov.SetKey(1, key(t))
+	prov.SetKey(2, key(t))
+	return crypto.NewKeyedEncryptor(prov, crypto.NewStaticKeyStore(2))
+}
+
+func TestSecretReencrypt_AdvancesAllRows(t *testing.T) {
+	ctx := context.Background()
+	enc := twoVersionEnc(t)
+
+	// Seed 3 rows encrypted under v1 + 1 NULL row.
+	store := &memStore{}
+	for i := 0; i < 3; i++ {
+		pt := []byte{byte(i), byte(i + 1)}
+		ct, err := enc.SealWithVersion(ctx, 1, pt)
+		require.NoError(t, err)
+		store.rows = append(store.rows, &memRow{blob: ct, keyVersion: 1, plaintext: pt})
+	}
+	store.rows = append(store.rows, &memRow{null: true, keyVersion: 1})
+
+	h := NewSecretReencryptHandler(store, enc, nil)
+	payload, _ := json.Marshal(SecretReencryptPayload{ActiveVersion: 2, BatchSize: 2})
+	require.NoError(t, h.Handle(ctx, jobs.Job{Payload: payload}))
+
+	// All non-null rows now on v2 and decrypt to identical plaintext.
+	for _, r := range store.rows {
+		if r.null {
+			require.Equal(t, 1, r.keyVersion, "NULL rows are skipped, version unchanged")
+			continue
+		}
+		require.Equal(t, 2, r.keyVersion)
+		got, err := enc.DecryptVersion(ctx, r.blob, 2)
+		require.NoError(t, err)
+		require.Equal(t, r.plaintext, got)
+	}
+	left, _ := store.StaleCount(ctx, "t", 2)
+	require.Equal(t, 0, left)
+}
+
+func TestSecretReencrypt_IdempotentRerun(t *testing.T) {
+	ctx := context.Background()
+	enc := twoVersionEnc(t)
+	store := &memStore{}
+	pt := []byte("hello")
+	ct, err := enc.SealWithVersion(ctx, 1, pt)
+	require.NoError(t, err)
+	store.rows = append(store.rows, &memRow{blob: ct, keyVersion: 1, plaintext: pt})
+
+	h := NewSecretReencryptHandler(store, enc, nil)
+	payload, _ := json.Marshal(SecretReencryptPayload{ActiveVersion: 2})
+	require.NoError(t, h.Handle(ctx, jobs.Job{Payload: payload}))
+	blobAfterFirst := store.rows[0].blob
+
+	// Second run is a no-op: row already on active version, blob unchanged.
+	require.NoError(t, h.Handle(ctx, jobs.Job{Payload: payload}))
+	require.Equal(t, blobAfterFirst, store.rows[0].blob)
+}
+
+func TestSecretReencrypt_ResumableAfterCrash(t *testing.T) {
+	ctx := context.Background()
+	enc := twoVersionEnc(t)
+	store := &memStore{failNext: true}
+	for i := 0; i < 5; i++ {
+		pt := []byte{byte(i)}
+		ct, err := enc.SealWithVersion(ctx, 1, pt)
+		require.NoError(t, err)
+		store.rows = append(store.rows, &memRow{blob: ct, keyVersion: 1, plaintext: pt})
+	}
+
+	h := NewSecretReencryptHandler(store, enc, nil)
+	payload, _ := json.Marshal(SecretReencryptPayload{ActiveVersion: 2, BatchSize: 10})
+	// First run crashes (ctx.Canceled bubbles up).
+	require.Error(t, h.Handle(ctx, jobs.Job{Payload: payload}))
+	left, _ := store.StaleCount(ctx, "t", 2)
+	require.Greater(t, left, 0, "crash left some rows on old version")
+
+	// Re-run completes the rest.
+	require.NoError(t, h.Handle(ctx, jobs.Job{Payload: payload}))
+	left, _ = store.StaleCount(ctx, "t", 2)
+	require.Equal(t, 0, left)
+	for _, r := range store.rows {
+		got, err := enc.DecryptVersion(ctx, r.blob, 2)
+		require.NoError(t, err)
+		require.Equal(t, r.plaintext, got)
+	}
+}
+
+func TestSecretReencrypt_BadPayloadIsPermanent(t *testing.T) {
+	h := NewSecretReencryptHandler(&memStore{}, twoVersionEnc(t), nil)
+	err := h.Handle(context.Background(), jobs.Job{Payload: []byte("not json")})
+	require.Error(t, err)
+	require.True(t, jobs.IsPermanent(err))
+}