store.go

package cryptlite

import (
	"context"
	"database/sql"
	"encoding/hex"
	"fmt"
	"strconv"
	"sync"

	"github.com/berbyte/cryptlite/keyring"
	_ "modernc.org/sqlite"
)

// Store is the main entry point for encrypted SQLite access.
type Store struct {
	cfg    Config
	db     *sql.DB
	kr     keyring.Keyring
	mu     sync.RWMutex
	keys   map[int][]byte // version → raw key bytes
	active int            // active key version for new writes
}

// dsnWithPragmas wraps a plain file path in a file: URI with WAL journal mode
// and a 5-second busy timeout. Using the DSN ensures every new connection from
// the pool gets these settings, which is required for busy_timeout (per-connection pragma).
func dsnWithPragmas(path string) string {
	return "file:" + path + "?_pragma=journal_mode(WAL)&_pragma=busy_timeout(5000)"
}

// Open opens (or creates) the SQLite database and loads the active encryption key.
func Open(cfg Config) (*Store, error) {
	if err := cfg.validate(); err != nil {
		return nil, err
	}
	db, err := sql.Open("sqlite", dsnWithPragmas(cfg.DBPath))
	if err != nil {
		return nil, err
	}
	s := &Store{
		cfg:  cfg,
		db:   db,
		kr:   keyring.NewOS(),
		keys: make(map[int][]byte),
	}
	if err := s.migrate(); err != nil {
		db.Close()
		return nil, err
	}
	if err := s.loadOrCreateActiveKey(); err != nil {
		db.Close()
		return nil, err
	}
	return s, nil
}

// OpenWithKeyring opens the store using the provided Keyring — intended for tests.
func OpenWithKeyring(cfg Config, kr keyring.Keyring) (*Store, error) {
	if err := cfg.validate(); err != nil {
		return nil, err
	}
	db, err := sql.Open("sqlite", dsnWithPragmas(cfg.DBPath))
	if err != nil {
		return nil, err
	}
	s := &Store{
		cfg:  cfg,
		db:   db,
		kr:   kr,
		keys: make(map[int][]byte),
	}
	if err := s.migrate(); err != nil {
		db.Close()
		return nil, err
	}
	if err := s.loadOrCreateActiveKey(); err != nil {
		db.Close()
		return nil, err
	}
	return s, nil
}

// Close closes the underlying database connection.
func (s *Store) Close() error { return s.db.Close() }

// DB returns the raw *sql.DB for application queries.
func (s *Store) DB() *sql.DB { return s.db }

// Encrypt encrypts plaintext using the active key and the provided AAD.
func (s *Store) Encrypt(_ context.Context, plaintext, aad []byte) (*EncryptedBlob, error) {
	s.mu.RLock()
	version := s.active
	key := s.keys[version]
	s.mu.RUnlock()

	ct, nonce, err := encrypt(key, plaintext, aad)
	if err != nil {
		return nil, err
	}
	return &EncryptedBlob{Ciphertext: ct, Nonce: nonce, KeyVersion: version}, nil
}

// Decrypt decrypts a blob using the key version stored in the blob.
func (s *Store) Decrypt(_ context.Context, blob EncryptedBlob, aad []byte) ([]byte, error) {
	s.mu.RLock()
	key, ok := s.keys[blob.KeyVersion]
	s.mu.RUnlock()
	if !ok {
		// Try to load the key from keychain on demand.
		k, err := s.loadKeyVersion(blob.KeyVersion)
		if err != nil {
			return nil, ErrUnsupportedKeyVersion
		}
		key = k
	}
	return decrypt(key, blob.Ciphertext, blob.Nonce, aad)
}

// EncryptJSON is a convenience wrapper for JSON payloads with no AAD.
func (s *Store) EncryptJSON(ctx context.Context, rawJSON []byte) (*EncryptedBlob, error) {
	return s.Encrypt(ctx, rawJSON, nil)
}

// DecryptJSON is a convenience wrapper for JSON payloads with no AAD.
func (s *Store) DecryptJSON(ctx context.Context, blob EncryptedBlob) ([]byte, error) {
	return s.Decrypt(ctx, blob, nil)
}

// RotateKey creates a new key version in the keychain and marks it active.
// Old keys remain available for reads.
func (s *Store) RotateKey(_ context.Context) (int, error) {
	newKey, err := generateKey()
	if err != nil {
		return 0, err
	}
	s.mu.Lock()
	defer s.mu.Unlock()

	newVersion := s.active + 1
	if err := s.storeKeyVersion(newVersion, newKey); err != nil {
		return 0, err
	}
	s.keys[newVersion] = newKey
	s.active = newVersion
	return newVersion, nil
}

// keychainAccount returns the versioned account name for a key version.
func (s *Store) keychainAccount(version int) string {
	return s.cfg.Keychain.Account + "/v" + strconv.Itoa(version)
}

func (s *Store) storeKeyVersion(version int, key []byte) error {
	account := s.keychainAccount(version)
	if _, err := s.kr.Get(s.cfg.Keychain.Service, account); err == nil {
		return fmt.Errorf("cryptlite: store key v%d: %w (%s/%s)", version, ErrKeyAlreadyExists, s.cfg.Keychain.Service, account)
	}
	if err := s.kr.Set(s.cfg.Keychain.Service, account, hex.EncodeToString(key)); err != nil {
		return fmt.Errorf("cryptlite: store key v%d: %w", version, err)
	}
	_, err := s.db.Exec(
		`INSERT OR IGNORE INTO encsqlite_keys (version, keychain_service, keychain_account, created_at, active)
		 VALUES (?, ?, ?, strftime('%s','now'), 0)`,
		version, s.cfg.Keychain.Service, account,
	)
	return err
}

func (s *Store) loadKeyVersion(version int) ([]byte, error) {
	account := s.keychainAccount(version)
	hexKey, err := s.kr.Get(s.cfg.Keychain.Service, account)
	if err != nil {
		return nil, ErrKeyNotFound
	}
	key, err := hex.DecodeString(hexKey)
	if err != nil {
		return nil, ErrKeyNotFound
	}
	s.mu.Lock()
	s.keys[version] = key
	s.mu.Unlock()
	return key, nil
}

func (s *Store) loadOrCreateActiveKey() error {
	// Find the highest active key from the metadata table.
	var version int
	err := s.db.QueryRow(
		`SELECT version FROM encsqlite_keys WHERE active = 1 ORDER BY version DESC LIMIT 1`,
	).Scan(&version)
	if err == sql.ErrNoRows {
		// No active-key metadata. This is only safe to treat as a fresh
		// install if the keychain is also empty for the default version — if
		// a key is already sitting there, the db metadata and keychain have
		// diverged (e.g. a restored or partially-migrated database) and
		// minting + storing a new key would silently overwrite the real one,
		// orphaning everything already encrypted under it.
		account := s.keychainAccount(1)
		if _, getErr := s.kr.Get(s.cfg.Keychain.Service, account); getErr == nil {
			return fmt.Errorf("cryptlite: %w (%s/%s)", ErrInconsistentKeyState, s.cfg.Keychain.Service, account)
		}
		return s.bootstrapKey()
	}
	if err != nil {
		return err
	}
	key, err := s.loadKeyVersion(version)
	if err != nil {
		return err
	}
	s.mu.Lock()
	s.active = version
	s.keys[version] = key
	s.mu.Unlock()
	return nil
}

func (s *Store) bootstrapKey() error {
	key, err := generateKey()
	if err != nil {
		return err
	}
	const version = 1
	if err := s.storeKeyVersion(version, key); err != nil {
		return err
	}
	_, err = s.db.Exec(`UPDATE encsqlite_keys SET active = 1 WHERE version = ?`, version)
	if err != nil {
		return err
	}
	s.mu.Lock()
	s.keys[version] = key
	s.active = version
	s.mu.Unlock()
	return nil
}