store_test.go

package cryptlite_test

import (
	"context"
	"errors"
	"os"
	"path/filepath"
	"testing"

	"github.com/berbyte/cryptlite"
	"github.com/berbyte/cryptlite/keyring"
)

func newTestStore(t *testing.T) *cryptlite.Store {
	t.Helper()
	dir := t.TempDir()
	cfg := cryptlite.Config{
		DBPath: filepath.Join(dir, "test.db"),
		Keychain: cryptlite.KeychainConfig{
			Service: "cryptlite-test",
			Account: "default",
		},
	}
	s, err := cryptlite.OpenWithKeyring(cfg, keyring.NewMemory())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { s.Close() })
	return s
}

func TestStoreEncryptDecryptRoundtrip(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	payload := []byte(`{"key":"value"}`)
	blob, err := s.Encrypt(ctx, payload, []byte("table/col/id"))
	if err != nil {
		t.Fatal(err)
	}
	got, err := s.Decrypt(ctx, *blob, []byte("table/col/id"))
	if err != nil {
		t.Fatal(err)
	}
	if string(got) != string(payload) {
		t.Fatalf("got %q want %q", got, payload)
	}
}

func TestStoreEncryptJSONRoundtrip(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	payload := []byte(`{"name":"Alice","age":30}`)
	blob, err := s.EncryptJSON(ctx, payload)
	if err != nil {
		t.Fatal(err)
	}
	got, err := s.DecryptJSON(ctx, *blob)
	if err != nil {
		t.Fatal(err)
	}
	if string(got) != string(payload) {
		t.Fatalf("got %q want %q", got, payload)
	}
}

func TestStoreDecryptWrongAAD(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	blob, _ := s.Encrypt(ctx, []byte("secret"), []byte("ctx-a"))
	if _, err := s.Decrypt(ctx, *blob, []byte("ctx-b")); err == nil {
		t.Fatal("expected failure with wrong AAD")
	}
}

func TestStoreKeyRotation(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()

	// Encrypt with key v1.
	blob, err := s.Encrypt(ctx, []byte("old"), nil)
	if err != nil {
		t.Fatal(err)
	}
	if blob.KeyVersion != 1 {
		t.Fatalf("expected key version 1, got %d", blob.KeyVersion)
	}

	// Rotate to v2.
	v, err := s.RotateKey(ctx)
	if err != nil {
		t.Fatal(err)
	}
	if v != 2 {
		t.Fatalf("expected new version 2, got %d", v)
	}

	// New blobs use v2.
	blob2, err := s.Encrypt(ctx, []byte("new"), nil)
	if err != nil {
		t.Fatal(err)
	}
	if blob2.KeyVersion != 2 {
		t.Fatalf("expected key version 2, got %d", blob2.KeyVersion)
	}

	// Old blob (v1) still decryptable.
	got, err := s.Decrypt(ctx, *blob, nil)
	if err != nil {
		t.Fatal("old blob should still decrypt:", err)
	}
	if string(got) != "old" {
		t.Fatalf("got %q", got)
	}
}

func TestStoreMigrationIdempotent(t *testing.T) {
	dir := t.TempDir()
	cfg := cryptlite.Config{
		DBPath:   filepath.Join(dir, "idem.db"),
		Keychain: cryptlite.KeychainConfig{Service: "svc", Account: "acc"},
	}
	kr := keyring.NewMemory()
	s1, err := cryptlite.OpenWithKeyring(cfg, kr)
	if err != nil {
		t.Fatal(err)
	}
	s1.Close()

	s2, err := cryptlite.OpenWithKeyring(cfg, kr)
	if err != nil {
		t.Fatal("second open failed:", err)
	}
	s2.Close()
}

func TestStorePersistKey(t *testing.T) {
	dir := t.TempDir()
	cfg := cryptlite.Config{
		DBPath:   filepath.Join(dir, "persist.db"),
		Keychain: cryptlite.KeychainConfig{Service: "svc", Account: "acc"},
	}
	kr := keyring.NewMemory()

	s1, _ := cryptlite.OpenWithKeyring(cfg, kr)
	ctx := context.Background()
	blob, _ := s1.Encrypt(ctx, []byte("hello"), nil)
	s1.Close()

	// Re-open and decrypt using the persisted keychain entry.
	s2, err := cryptlite.OpenWithKeyring(cfg, kr)
	if err != nil {
		t.Fatal(err)
	}
	defer s2.Close()
	got, err := s2.Decrypt(ctx, *blob, nil)
	if err != nil {
		t.Fatal(err)
	}
	if string(got) != "hello" {
		t.Fatalf("got %q", got)
	}
}

func TestStoreSQLiteIntegration(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	db := s.DB()

	_, err := db.Exec(`CREATE TABLE IF NOT EXISTS notes (
		id TEXT PRIMARY KEY,
		title TEXT NOT NULL,
		body_ciphertext BLOB NOT NULL,
		body_nonce BLOB NOT NULL,
		body_key_version INTEGER NOT NULL
	)`)
	if err != nil {
		t.Fatal(err)
	}

	payload := []byte(`{"content":"secret note"}`)
	blob, err := s.Encrypt(ctx, payload, []byte("notes/body/1"))
	if err != nil {
		t.Fatal(err)
	}

	_, err = db.Exec(
		`INSERT INTO notes (id, title, body_ciphertext, body_nonce, body_key_version) VALUES (?,?,?,?,?)`,
		"1", "My Note", blob.Ciphertext, blob.Nonce, blob.KeyVersion,
	)
	if err != nil {
		t.Fatal(err)
	}

	var loaded cryptlite.EncryptedBlob
	err = db.QueryRow(`SELECT body_ciphertext, body_nonce, body_key_version FROM notes WHERE id = '1'`).
		Scan(&loaded.Ciphertext, &loaded.Nonce, &loaded.KeyVersion)
	if err != nil {
		t.Fatal(err)
	}

	got, err := s.Decrypt(ctx, loaded, []byte("notes/body/1"))
	if err != nil {
		t.Fatal(err)
	}
	if string(got) != string(payload) {
		t.Fatalf("got %q", got)
	}
}

func TestStoreBatchInsert(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	db := s.DB()

	_, _ = db.Exec(`CREATE TABLE items (id INTEGER PRIMARY KEY, val_ciphertext BLOB, val_nonce BLOB, val_key_version INTEGER)`)

	tx, err := db.Begin()
	if err != nil {
		t.Fatal(err)
	}
	for i := 0; i < 50; i++ {
		blob, err := s.Encrypt(ctx, []byte("item"), nil)
		if err != nil {
			tx.Rollback()
			t.Fatal(err)
		}
		if _, err = tx.Exec(`INSERT INTO items VALUES (?,?,?,?)`, i, blob.Ciphertext, blob.Nonce, blob.KeyVersion); err != nil {
			tx.Rollback()
			t.Fatal(err)
		}
	}
	if err := tx.Commit(); err != nil {
		t.Fatal(err)
	}

	var count int
	db.QueryRow(`SELECT COUNT(*) FROM items`).Scan(&count)
	if count != 50 {
		t.Fatalf("expected 50, got %d", count)
	}
}

func TestInvalidConfig(t *testing.T) {
	_, err := cryptlite.Open(cryptlite.Config{})
	if err != cryptlite.ErrInvalidConfig {
		t.Fatalf("expected ErrInvalidConfig, got %v", err)
	}
}

// TestLoadOrCreateActiveKey_InconsistentState covers the data-loss bug where a
// database with no active-key metadata (e.g. restored from backup, or a
// partially-applied migration) used to silently mint a brand-new key and
// overwrite whatever was already in the keychain for version 1 — orphaning
// every row encrypted under the real key. It must now fail loudly instead.
func TestLoadOrCreateActiveKey_InconsistentState(t *testing.T) {
	dir := t.TempDir()
	cfg := cryptlite.Config{
		DBPath:   filepath.Join(dir, "inconsistent.db"),
		Keychain: cryptlite.KeychainConfig{Service: "svc", Account: "default"},
	}
	kr := keyring.NewMemory()
	// Simulate a pre-existing real key for v1 with no corresponding db metadata
	// (fresh/restored db file, divergent keychain).
	if err := kr.Set("svc", "default/v1", "deadbeef"); err != nil {
		t.Fatal(err)
	}

	_, err := cryptlite.OpenWithKeyring(cfg, kr)
	if !errors.Is(err, cryptlite.ErrInconsistentKeyState) {
		t.Fatalf("expected ErrInconsistentKeyState, got %v", err)
	}
}

// TestStoreKeyVersion_RefusesOverwrite covers RotateKey racing (or being
// retried) such that a key already exists in the keychain for the version
// about to be minted — it must refuse rather than silently clobber it.
func TestStoreKeyVersion_RefusesOverwrite(t *testing.T) {
	dir := t.TempDir()
	cfg := cryptlite.Config{
		DBPath:   filepath.Join(dir, "rotate.db"),
		Keychain: cryptlite.KeychainConfig{Service: "svc", Account: "default"},
	}
	kr := keyring.NewMemory()
	s, err := cryptlite.OpenWithKeyring(cfg, kr)
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()

	// Pre-seed the keychain with whatever the *next* version's account would
	// be, simulating a key that's already there (e.g. a previous RotateKey
	// call whose db write didn't commit).
	if err := kr.Set("svc", "default/v2", "deadbeef"); err != nil {
		t.Fatal(err)
	}

	if _, err := s.RotateKey(context.Background()); !errors.Is(err, cryptlite.ErrKeyAlreadyExists) {
		t.Fatalf("expected ErrKeyAlreadyExists, got %v", err)
	}
}

func init() {
	// Ensure temp dir cleanup.
	_ = os.MkdirTemp
}