chio-store-sqlite 0.1.2

//! Tenant-scoped encrypted BLOB persistence for tee capture payloads.
//!
//! The tee stores redacted request and response bodies as opaque
//! binary payloads. This module keeps those payloads encrypted at rest
//! with a tenant-provided 32-byte key and a per-blob ChaCha20-Poly1305
//! nonce. Authentication failures on read are surfaced as errors and do
//! not return plaintext.

use std::fs;
use std::path::Path;
use std::time::{SystemTime, UNIX_EPOCH};

use chacha20poly1305::aead::rand_core::RngCore;
use chacha20poly1305::aead::{Aead, KeyInit, OsRng, Payload};
use chacha20poly1305::{ChaCha20Poly1305, Key, Nonce};
use r2d2::Pool;
use r2d2_sqlite::SqliteConnectionManager;
use rusqlite::{params, OptionalExtension};
use uuid::Uuid;
use zeroize::{Zeroize, ZeroizeOnDrop};

/// Tenant identifier used to isolate encrypted BLOB rows.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct TenantId(String);

impl TenantId {
    /// Construct a tenant id from a string-like value.
    #[must_use]
    pub fn new(value: impl Into<String>) -> Self {
        Self(value.into())
    }

    /// Borrow the tenant id as stored in SQLite.
    #[must_use]
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl From<&str> for TenantId {
    fn from(value: &str) -> Self {
        Self::new(value)
    }
}

impl From<String> for TenantId {
    fn from(value: String) -> Self {
        Self::new(value)
    }
}

/// Tenant-scoped 256-bit AEAD key.
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct TenantKey([u8; 32]);

impl TenantKey {
    /// Construct from exactly 32 bytes.
    #[must_use]
    pub fn from_bytes(bytes: [u8; 32]) -> Self {
        Self(bytes)
    }

    /// Borrow key bytes for callers that already own secure key custody.
    #[must_use]
    pub fn as_bytes(&self) -> &[u8; 32] {
        &self.0
    }
}

impl std::fmt::Debug for TenantKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("TenantKey(<redacted>)")
    }
}

/// Ciphertext plus the unique nonce required for decryption.
#[derive(Clone, PartialEq, Eq)]
pub struct EncryptedBlob {
    /// ChaCha20-Poly1305 ciphertext and authentication tag.
    pub ciphertext: Vec<u8>,
    /// 96-bit nonce generated per encrypted blob.
    pub nonce: [u8; 12],
}

struct StoredEncryptedBlob {
    blob: EncryptedBlob,
    blob_id: String,
    tenant_id: TenantId,
    created_at: i64,
}

impl std::fmt::Debug for EncryptedBlob {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("EncryptedBlob")
            .field("ciphertext_len", &self.ciphertext.len())
            .field("nonce", &self.nonce)
            .finish()
    }
}

/// Opaque handle returned after a blob is written.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct BlobHandle {
    blob_id: String,
    tenant_id: TenantId,
}

impl BlobHandle {
    fn new(blob_id: String, tenant_id: TenantId) -> Self {
        Self { blob_id, tenant_id }
    }

    /// Stable blob id generated by the store.
    #[must_use]
    pub fn blob_id(&self) -> &str {
        &self.blob_id
    }

    /// Tenant scope bound into this handle.
    #[must_use]
    pub fn tenant_id(&self) -> &TenantId {
        &self.tenant_id
    }
}

/// Decryption failure. Authentication failure is intentionally coarse so
/// callers cannot distinguish a wrong key from ciphertext tampering.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DecryptError {
    /// AEAD authentication failed.
    AuthenticationFailed,
}

impl std::fmt::Display for DecryptError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::AuthenticationFailed => f.write_str("blob authentication failed"),
        }
    }
}

impl std::error::Error for DecryptError {}

/// Errors returned by SQLite encrypted BLOB persistence.
#[derive(Debug)]
pub enum BlobStoreError {
    /// SQLite connection pool acquisition failed.
    Pool(String),
    /// SQLite returned an error.
    Sqlite(String),
    /// Filesystem operation failed while opening a store.
    Io(String),
    /// Tenant ids must be present before a blob can be persisted.
    EmptyTenantId,
    /// Tenant ids must be exact storage scope keys.
    InvalidTenantId,
    /// Stored nonce length was not the required 12 bytes.
    InvalidNonceLength(usize),
    /// The handle did not identify a row in its tenant scope.
    NotFound,
    /// AEAD authentication failed.
    Decrypt(DecryptError),
}

impl std::fmt::Display for BlobStoreError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Pool(error) => write!(f, "sqlite encrypted blob pool error: {error}"),
            Self::Sqlite(error) => write!(f, "sqlite encrypted blob error: {error}"),
            Self::Io(error) => write!(f, "sqlite encrypted blob io error: {error}"),
            Self::EmptyTenantId => f.write_str("tenant id must not be empty"),
            Self::InvalidTenantId => {
                f.write_str("tenant id must be unpadded and contain no control characters")
            }
            Self::InvalidNonceLength(len) => {
                write!(f, "encrypted blob nonce must be 12 bytes, got {len}")
            }
            Self::NotFound => f.write_str("encrypted blob handle not found"),
            Self::Decrypt(error) => write!(f, "encrypted blob decrypt error: {error}"),
        }
    }
}

impl std::error::Error for BlobStoreError {}

impl From<rusqlite::Error> for BlobStoreError {
    fn from(error: rusqlite::Error) -> Self {
        Self::Sqlite(error.to_string())
    }
}

impl From<r2d2::Error> for BlobStoreError {
    fn from(error: r2d2::Error) -> Self {
        Self::Pool(error.to_string())
    }
}

impl From<std::io::Error> for BlobStoreError {
    fn from(error: std::io::Error) -> Self {
        Self::Io(error.to_string())
    }
}

impl From<DecryptError> for BlobStoreError {
    fn from(error: DecryptError) -> Self {
        Self::Decrypt(error)
    }
}

/// Encryption failure surfaced without panicking in public helpers.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EncryptError {
    /// AEAD encryption failed.
    EncryptionFailed,
}

impl std::fmt::Display for EncryptError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::EncryptionFailed => f.write_str("blob encryption failed"),
        }
    }
}

impl std::error::Error for EncryptError {}

/// Encrypt `plaintext` with the tenant key.
pub fn encrypt_blob(
    tenant_key: &TenantKey,
    plaintext: &[u8],
) -> Result<EncryptedBlob, EncryptError> {
    try_encrypt_blob(tenant_key, plaintext).map_err(|_| EncryptError::EncryptionFailed)
}

/// Decrypt `blob` with the tenant key.
pub fn decrypt_blob(tenant_key: &TenantKey, blob: &EncryptedBlob) -> Result<Vec<u8>, DecryptError> {
    let cipher = cipher_for_key(tenant_key);
    cipher
        .decrypt(Nonce::from_slice(&blob.nonce), blob.ciphertext.as_ref())
        .map_err(|_| DecryptError::AuthenticationFailed)
}

pub(crate) fn decrypt_blob_with_aad(
    tenant_key: &TenantKey,
    blob: &EncryptedBlob,
    aad: &[u8],
) -> Result<Vec<u8>, DecryptError> {
    let cipher = cipher_for_key(tenant_key);
    cipher
        .decrypt(
            Nonce::from_slice(&blob.nonce),
            Payload {
                msg: blob.ciphertext.as_ref(),
                aad,
            },
        )
        .map_err(|_| DecryptError::AuthenticationFailed)
}

fn try_encrypt_blob(tenant_key: &TenantKey, plaintext: &[u8]) -> Result<EncryptedBlob, ()> {
    try_encrypt_blob_with_aad(tenant_key, plaintext, &[])
}

pub(crate) fn try_encrypt_blob_with_aad(
    tenant_key: &TenantKey,
    plaintext: &[u8],
    aad: &[u8],
) -> Result<EncryptedBlob, ()> {
    let cipher = cipher_for_key(tenant_key);
    let mut nonce = [0_u8; 12];
    OsRng.fill_bytes(&mut nonce);
    let ciphertext = cipher
        .encrypt(
            Nonce::from_slice(&nonce),
            Payload {
                msg: plaintext,
                aad,
            },
        )
        .map_err(|_| ())?;
    Ok(EncryptedBlob { ciphertext, nonce })
}

fn cipher_for_key(tenant_key: &TenantKey) -> ChaCha20Poly1305 {
    ChaCha20Poly1305::new(Key::from_slice(tenant_key.as_bytes()))
}

/// SQLite-backed encrypted BLOB store.
pub struct SqliteEncryptedBlobStore {
    pool: Pool<SqliteConnectionManager>,
}

/// Encrypted-blob-store schema revision. Bump on every schema-affecting change.
const ENCRYPTED_BLOB_STORE_SUPPORTED_SCHEMA_VERSION: i32 = 0;
/// Stable key under which this store records its schema revision in the shared
/// keyed metadata table, distinct from any co-located store's key.
const ENCRYPTED_BLOB_STORE_SCHEMA_KEY: &str = "encrypted_blob";
/// Tables shipped before schema stamping existed, used to adopt a pre-stamping
/// encrypted-blob database rather than reject it as foreign.
const ENCRYPTED_BLOB_STORE_LEGACY_ANCHOR_TABLES: &[&str] = &["chio_encrypted_blobs"];

impl SqliteEncryptedBlobStore {
    /// Open the store at `path`, creating parent directories if needed.
    pub fn open(path: impl AsRef<Path>) -> Result<Self, BlobStoreError> {
        let path = path.as_ref();
        // Resolve any `file:` URI to its on-disk parent before creating it, so a
        // URI-configured store creates the real backing directory rather than a
        // bogus scheme-prefixed one.
        if let Some(parent) = crate::sqlite_parent_dir_to_create(path) {
            fs::create_dir_all(&parent)?;
        }
        let manager = SqliteConnectionManager::file(path);
        let pool = Pool::builder().max_size(8).build(manager)?;
        let store = Self { pool };
        store.run_migrations()?;
        Ok(store)
    }

    /// Open an in-memory store for tests.
    pub fn open_in_memory() -> Result<Self, BlobStoreError> {
        let manager = SqliteConnectionManager::memory();
        let pool = Pool::builder().max_size(1).build(manager)?;
        let store = Self { pool };
        store.run_migrations()?;
        Ok(store)
    }

    fn run_migrations(&self) -> Result<(), BlobStoreError> {
        let conn = self.pool.get()?;
        crate::check_schema_version(
            &conn,
            ENCRYPTED_BLOB_STORE_SCHEMA_KEY,
            ENCRYPTED_BLOB_STORE_SUPPORTED_SCHEMA_VERSION,
            ENCRYPTED_BLOB_STORE_LEGACY_ANCHOR_TABLES,
        )
        .map_err(|error| BlobStoreError::Sqlite(error.to_string()))?;
        conn.execute_batch(
            r#"
            PRAGMA journal_mode = WAL;
            PRAGMA synchronous = FULL;
            PRAGMA busy_timeout = 5000;

            CREATE TABLE IF NOT EXISTS chio_encrypted_blobs (
                blob_id TEXT PRIMARY KEY,
                tenant_id TEXT NOT NULL,
                nonce BLOB NOT NULL CHECK(length(nonce) = 12),
                ciphertext BLOB NOT NULL,
                created_at INTEGER NOT NULL
            );

            CREATE INDEX IF NOT EXISTS idx_chio_encrypted_blobs_tenant
                ON chio_encrypted_blobs(tenant_id, created_at);
            "#,
        )?;
        crate::stamp_schema_version(
            &conn,
            ENCRYPTED_BLOB_STORE_SCHEMA_KEY,
            ENCRYPTED_BLOB_STORE_SUPPORTED_SCHEMA_VERSION,
        )
        .map_err(|error| BlobStoreError::Sqlite(error.to_string()))?;
        Ok(())
    }

    /// Write `payload` encrypted under `key` and return an opaque handle.
    pub fn write_encrypted_blob(
        &self,
        tenant_id: &TenantId,
        key: &TenantKey,
        payload: &[u8],
    ) -> Result<BlobHandle, BlobStoreError> {
        validate_tenant_id(tenant_id)?;
        let blob_id = format!("blob-{}", Uuid::now_v7());
        let created_at = now_secs();
        let aad = blob_aad(&blob_id, tenant_id.as_str(), created_at);
        let encrypted = try_encrypt_blob_with_aad(key, payload, &aad)
            .map_err(|_| BlobStoreError::Decrypt(DecryptError::AuthenticationFailed))?;
        let conn = self.pool.get()?;
        conn.execute(
            r#"
            INSERT INTO chio_encrypted_blobs
                (blob_id, tenant_id, nonce, ciphertext, created_at)
            VALUES
                (?1, ?2, ?3, ?4, ?5)
            "#,
            params![
                blob_id,
                tenant_id.as_str(),
                encrypted.nonce.as_slice(),
                encrypted.ciphertext,
                created_at,
            ],
        )?;
        Ok(BlobHandle::new(blob_id, tenant_id.clone()))
    }

    /// Load encrypted material without decrypting it. Intended for
    /// audit, migration, and test verification paths.
    pub fn load_encrypted_blob(
        &self,
        handle: &BlobHandle,
    ) -> Result<EncryptedBlob, BlobStoreError> {
        let conn = self.pool.get()?;
        let row = conn
            .query_row(
                r#"
                SELECT nonce, ciphertext
                FROM chio_encrypted_blobs
                WHERE blob_id = ?1 AND tenant_id = ?2
                "#,
                params![handle.blob_id(), handle.tenant_id().as_str()],
                |row| {
                    let nonce: Vec<u8> = row.get("nonce")?;
                    let ciphertext: Vec<u8> = row.get("ciphertext")?;
                    Ok((nonce, ciphertext))
                },
            )
            .optional()?;

        let Some((nonce, ciphertext)) = row else {
            return Err(BlobStoreError::NotFound);
        };
        let nonce_len = nonce.len();
        let nonce = match <[u8; 12]>::try_from(nonce) {
            Ok(value) => value,
            Err(_) => return Err(BlobStoreError::InvalidNonceLength(nonce_len)),
        };
        Ok(EncryptedBlob { ciphertext, nonce })
    }

    /// Read and decrypt a blob. Wrong keys, wrong handles, and tampered
    /// ciphertext all fail closed by returning an error.
    pub fn read_encrypted_blob(
        &self,
        handle: &BlobHandle,
        key: &TenantKey,
    ) -> Result<Vec<u8>, BlobStoreError> {
        let stored = self.load_encrypted_blob_record(handle)?;
        let aad = blob_aad(
            &stored.blob_id,
            stored.tenant_id.as_str(),
            stored.created_at,
        );
        Ok(decrypt_blob_with_aad(key, &stored.blob, &aad)?)
    }

    fn load_encrypted_blob_record(
        &self,
        handle: &BlobHandle,
    ) -> Result<StoredEncryptedBlob, BlobStoreError> {
        let conn = self.pool.get()?;
        let row = conn
            .query_row(
                r#"
                SELECT blob_id, tenant_id, nonce, ciphertext, created_at
                FROM chio_encrypted_blobs
                WHERE blob_id = ?1 AND tenant_id = ?2
                "#,
                params![handle.blob_id(), handle.tenant_id().as_str()],
                |row| {
                    let blob_id: String = row.get("blob_id")?;
                    let tenant_id: String = row.get("tenant_id")?;
                    let nonce: Vec<u8> = row.get("nonce")?;
                    let ciphertext: Vec<u8> = row.get("ciphertext")?;
                    let created_at: i64 = row.get("created_at")?;
                    Ok((blob_id, tenant_id, nonce, ciphertext, created_at))
                },
            )
            .optional()?;

        let Some((blob_id, tenant_id, nonce, ciphertext, created_at)) = row else {
            return Err(BlobStoreError::NotFound);
        };
        let nonce_len = nonce.len();
        let nonce = match <[u8; 12]>::try_from(nonce) {
            Ok(value) => value,
            Err(_) => return Err(BlobStoreError::InvalidNonceLength(nonce_len)),
        };
        Ok(StoredEncryptedBlob {
            blob: EncryptedBlob { ciphertext, nonce },
            blob_id,
            tenant_id: TenantId::new(tenant_id),
            created_at,
        })
    }
}

fn blob_aad(blob_id: &str, tenant_id: &str, created_at: i64) -> Vec<u8> {
    format!("{blob_id}\0{tenant_id}\0{created_at}").into_bytes()
}

fn validate_tenant_id(tenant_id: &TenantId) -> Result<(), BlobStoreError> {
    let value = tenant_id.as_str();
    if value.trim().is_empty() {
        return Err(BlobStoreError::EmptyTenantId);
    }
    if value.trim() != value || value.chars().any(char::is_control) {
        return Err(BlobStoreError::InvalidTenantId);
    }
    Ok(())
}

fn now_secs() -> i64 {
    match SystemTime::now().duration_since(UNIX_EPOCH) {
        Ok(duration) => i64::try_from(duration.as_secs()).unwrap_or(i64::MAX),
        Err(_) => 0,
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;

    #[test]
    fn encrypt_decrypt_round_trip() {
        let key = TenantKey::from_bytes([7; 32]);
        let plaintext = b"redacted fixture payload";
        let blob = encrypt_blob(&key, plaintext).expect("encrypt with valid key");
        assert_ne!(blob.ciphertext, plaintext);
        let decrypted = decrypt_blob(&key, &blob).expect("decrypt with correct key");
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn decrypt_rejects_wrong_key() {
        let key = TenantKey::from_bytes([7; 32]);
        let wrong_key = TenantKey::from_bytes([8; 32]);
        let blob = encrypt_blob(&key, b"payload").expect("encrypt with valid key");
        let error = decrypt_blob(&wrong_key, &blob).expect_err("wrong key must fail");
        assert_eq!(error, DecryptError::AuthenticationFailed);
    }

    #[test]
    fn stored_blob_ciphertext_is_bound_to_handle_metadata() {
        let store = SqliteEncryptedBlobStore::open_in_memory().unwrap();
        let tenant = TenantId::new("tenant-a");
        let key = TenantKey::from_bytes([9; 32]);
        let source = store
            .write_encrypted_blob(&tenant, &key, b"source payload")
            .unwrap();
        let target = store
            .write_encrypted_blob(&tenant, &key, b"target payload")
            .unwrap();

        let conn = store.pool.get().unwrap();
        conn.execute(
            r#"
            UPDATE chio_encrypted_blobs
            SET nonce = (
                    SELECT nonce FROM chio_encrypted_blobs WHERE blob_id = ?1
                ),
                ciphertext = (
                    SELECT ciphertext FROM chio_encrypted_blobs WHERE blob_id = ?1
                )
            WHERE blob_id = ?2
            "#,
            params![source.blob_id(), target.blob_id()],
        )
        .unwrap();
        drop(conn);

        let error = store
            .read_encrypted_blob(&target, &key)
            .expect_err("transplanted ciphertext should fail metadata-bound AEAD");
        assert!(matches!(
            error,
            BlobStoreError::Decrypt(DecryptError::AuthenticationFailed)
        ));
    }

    #[test]
    fn write_rejects_padded_or_control_tenant_id_before_persistence() {
        let store = SqliteEncryptedBlobStore::open_in_memory().unwrap();
        let key = TenantKey::from_bytes([9; 32]);

        for tenant in [" tenant-a", "tenant-a ", "tenant-a\n"] {
            let error = store
                .write_encrypted_blob(&TenantId::new(tenant), &key, b"payload")
                .expect_err("malformed tenant id must fail closed");
            assert!(
                error.to_string().contains("tenant id"),
                "unexpected tenant validation error: {error}"
            );
        }

        let count: i64 = store
            .pool
            .get()
            .unwrap()
            .query_row("SELECT COUNT(*) FROM chio_encrypted_blobs", [], |row| {
                row.get(0)
            })
            .unwrap();
        assert_eq!(count, 0);
    }

    #[test]
    fn store_schema_does_not_publish_plaintext_hashes() {
        let store = SqliteEncryptedBlobStore::open_in_memory().unwrap();
        let conn = store.pool.get().unwrap();
        let mut statement = conn
            .prepare("PRAGMA table_info(chio_encrypted_blobs)")
            .unwrap();
        let names = statement
            .query_map([], |row| row.get::<_, String>(1))
            .unwrap()
            .collect::<Result<Vec<_>, _>>()
            .unwrap();
        assert!(
            !names.iter().any(|name| name == "plaintext_sha256"),
            "encrypted blob store must not persist plaintext-derived hashes"
        );
    }

    #[test]
    fn tenant_key_implements_zeroize_on_drop() {
        fn assert_zod<T: ZeroizeOnDrop>() {}
        assert_zod::<TenantKey>();
    }

    #[test]
    fn tenant_key_debug_redacts_material() {
        let key = TenantKey::from_bytes([7; 32]);
        assert_eq!(std::format!("{key:?}"), "TenantKey(<redacted>)");
    }
}