klieo-ops 0.5.0

//! Pluggable at-rest encryption for klieo-ops adapters that want to
//! crypto-shred per-tenant data (GDPR Art 17 erasure).
//!
//! `Crypter` is the abstract port; `TenantKeyedCrypter` is the software
//! default impl using AES-256-GCM. `KeyProvider` resolves tenant → key
//! bytes; `TenantKeyedCrypter` consults `klieo_ops::tenant::current_tenant()`
//! at encrypt/decrypt time so the same Crypter instance can serve
//! multiple tenants.
//!
//! ## Wire format (v1)
//!
//! ```text
//! version_u8 || nonce[12] || aad_len_u16_le || aad[aad_len] || ct_and_tag
//! ```
//!
//! AAD = `tenant_bytes || version_byte` — bound to the ciphertext via the
//! AEAD tag so key-version mismatches and cross-tenant decryption attempts
//! are detected at the tag-check stage (defence-in-depth).
//!
//! Closes spec § 5.4 OQ-5. Integration with `EpisodicMemory` is the
//! adapter author's responsibility per spec § 1.9 — see the chapter
//! at `docs/book/src/ops/encryption.md` for the recommended wrapping
//! pattern.

use crate::tenant::current_tenant;
use crate::types::TenantId;
use aes_gcm::aead::{Aead, KeyInit, OsRng};
use aes_gcm::{AeadCore, Aes256Gcm, Nonce};
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
use thiserror::Error;

/// Errors raised by [`Crypter`] / [`KeyProvider`] impls.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum CrypterError {
    /// Tenant key not available (crypto-shred or unconfigured tenant).
    #[error("key destroyed or not configured for tenant {tenant:?}")]
    KeyDestroyed {
        /// Tenant tag that was looked up.
        tenant: Option<TenantId>,
    },
    /// Underlying AEAD operation failed (e.g. authentication tag mismatch).
    #[error("aead failed: {0}")]
    Aead(String),
    /// AAD embedded in the ciphertext does not match the current context
    /// (wrong tenant scope or stale key version). Defence-in-depth: the
    /// AEAD tag already fails; this error gives a more actionable signal.
    #[error("aad mismatch: stored aad does not match current context")]
    AadMismatch,
    /// Ciphertext header is too short to contain the framing fields.
    #[error("ciphertext too short: {0} bytes")]
    Truncated(usize),
    /// Other internal error.
    #[error("internal: {0}")]
    Internal(String),
}

/// Key material returned by a versioned lookup.
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct KeyMaterial {
    /// 32-byte AES-256 key.
    pub key: [u8; 32],
    /// Monotonic version counter used in AAD binding. `0` for
    /// non-versioned providers.
    pub version: u8,
}

impl KeyMaterial {
    /// Construct key material from raw bytes and a version counter.
    #[must_use]
    pub fn new(key: [u8; 32], version: u8) -> Self {
        Self { key, version }
    }
}

/// Resolve a tenant tag to an encryption key. Implementations include
/// [`StaticKeyProvider`] (in-memory map for dev) and operator HSM-backed
/// adapters.
///
/// Returning `None` from `lookup` is the crypto-shred signal: past
/// ciphertext written under this tenant is now unreadable.
pub trait KeyProvider: Send + Sync {
    /// Resolve the tenant to a 32-byte AES-256 key. `None` means
    /// "key destroyed or unconfigured" — [`TenantKeyedCrypter`] raises
    /// [`CrypterError::KeyDestroyed`].
    fn lookup(&self, tenant: &Option<TenantId>) -> Option<[u8; 32]>;

    /// Versioned lookup. Returns key bytes **and** a monotonic version
    /// counter that will be embedded in the ciphertext AAD.
    ///
    /// The default implementation delegates to [`lookup`](Self::lookup)
    /// and reports version `0`. Version-aware providers override this
    /// method to supply the current key version.
    fn lookup_versioned(&self, tenant: &Option<TenantId>) -> Option<KeyMaterial> {
        self.lookup(tenant)
            .map(|key| KeyMaterial { key, version: 0 })
    }
}

/// In-memory [`KeyProvider`] for dev / non-regulated deployments.
#[non_exhaustive]
pub struct StaticKeyProvider {
    keys: HashMap<Option<TenantId>, [u8; 32]>,
}

impl StaticKeyProvider {
    /// Build from a map. The `None` entry serves as the default key for
    /// unscoped operations (e.g. global audit events with no tenant).
    #[must_use]
    pub fn from_map(keys: HashMap<Option<TenantId>, [u8; 32]>) -> Self {
        Self { keys }
    }
}

impl KeyProvider for StaticKeyProvider {
    fn lookup(&self, tenant: &Option<TenantId>) -> Option<[u8; 32]> {
        self.keys.get(tenant).copied()
    }
}

/// Abstract at-rest encryption port. Implementations are pure (no I/O).
///
/// Ciphertext format is implementation-defined; callers must feed the
/// output of `encrypt` directly into `decrypt` on the same impl.
#[async_trait]
pub trait Crypter: Send + Sync {
    /// Encrypt `plaintext`. The returned bytes carry whatever framing
    /// the impl needs for later decryption (e.g. nonce prefix).
    async fn encrypt(&self, plaintext: &[u8]) -> Result<Vec<u8>, CrypterError>;

    /// Decrypt `ciphertext` previously produced by [`Crypter::encrypt`].
    async fn decrypt(&self, ciphertext: &[u8]) -> Result<Vec<u8>, CrypterError>;
}

/// Tenant-keyed AES-256-GCM [`Crypter`]. Looks up the current tenant's
/// key via [`KeyProvider`] and encrypts / decrypts under it.
///
/// ## Wire format
///
/// ```text
/// version_u8 || nonce[12] || aad_len_u16_le || aad[aad_len] || ct_and_tag
/// ```
///
/// AAD = `tenant_bytes || version_byte`. The AEAD tag covers the AAD so
/// cross-tenant and wrong-key-version decryptions fail at tag verification.
#[non_exhaustive]
pub struct TenantKeyedCrypter {
    keys: Arc<dyn KeyProvider>,
}

impl TenantKeyedCrypter {
    /// Build with the provided [`KeyProvider`].
    #[must_use]
    pub fn new(keys: Arc<dyn KeyProvider>) -> Self {
        Self { keys }
    }

    fn resolve_key(&self, tenant: &Option<TenantId>) -> Result<KeyMaterial, CrypterError> {
        self.keys
            .lookup_versioned(tenant)
            .ok_or_else(|| CrypterError::KeyDestroyed {
                tenant: tenant.clone(),
            })
    }

    /// Canonical AAD: tenant bytes (empty for `None`) concatenated with the
    /// single key-version byte.
    fn build_aad(tenant: &Option<TenantId>, version: u8) -> Vec<u8> {
        let mut aad = match tenant {
            Some(t) => t.0.as_bytes().to_vec(),
            None => Vec::new(),
        };
        aad.push(version);
        aad
    }
}

// Header byte counts used for framing.
const VERSION_LEN: usize = 1;
const NONCE_LEN: usize = 12;
const AAD_LEN_FIELD: usize = 2; // u16 little-endian
const MIN_HEADER: usize = VERSION_LEN + NONCE_LEN + AAD_LEN_FIELD;

#[async_trait]
impl Crypter for TenantKeyedCrypter {
    async fn encrypt(&self, plaintext: &[u8]) -> Result<Vec<u8>, CrypterError> {
        let tenant = current_tenant();
        let km = self.resolve_key(&tenant)?;
        let cipher = Aes256Gcm::new((&km.key).into());
        let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
        let aad = Self::build_aad(&tenant, km.version);
        let aad_len = aad.len() as u16;

        let body = cipher
            .encrypt(
                &nonce,
                aes_gcm::aead::Payload {
                    msg: plaintext,
                    aad: &aad,
                },
            )
            .map_err(|e| CrypterError::Aead(format!("encrypt: {e}")))?;

        let mut out =
            Vec::with_capacity(VERSION_LEN + NONCE_LEN + AAD_LEN_FIELD + aad.len() + body.len());
        out.push(km.version);
        out.extend_from_slice(nonce.as_slice());
        out.extend_from_slice(&aad_len.to_le_bytes());
        out.extend_from_slice(&aad);
        out.extend_from_slice(&body);
        Ok(out)
    }

    async fn decrypt(&self, ciphertext: &[u8]) -> Result<Vec<u8>, CrypterError> {
        if ciphertext.len() < MIN_HEADER {
            return Err(CrypterError::Truncated(ciphertext.len()));
        }

        // version byte is part of the framing but its value is validated
        // implicitly via the AAD comparison below; we skip a named binding.
        let nonce = Nonce::from_slice(&ciphertext[VERSION_LEN..VERSION_LEN + NONCE_LEN]);
        let aad_len = u16::from_le_bytes([
            ciphertext[VERSION_LEN + NONCE_LEN],
            ciphertext[VERSION_LEN + NONCE_LEN + 1],
        ]) as usize;

        let aad_start = MIN_HEADER;
        let body_start = aad_start + aad_len;
        if ciphertext.len() < body_start {
            return Err(CrypterError::Truncated(ciphertext.len()));
        }

        let stored_aad = &ciphertext[aad_start..body_start];
        let body = &ciphertext[body_start..];

        let tenant = current_tenant();
        let km = self.resolve_key(&tenant)?;

        // Verify the AAD stored in the ciphertext matches what the current
        // context and key version would produce. Using `km.version` (the
        // provider's current version) ensures that a key rotation — where
        // `km.version` advances — surfaces as `AadMismatch` before we even
        // attempt the AEAD tag check. Cross-tenant attempts also fail here
        // because the tenant bytes in `stored_aad` won't match the current
        // tenant's bytes in `expected_aad`.
        let expected_aad = Self::build_aad(&tenant, km.version);
        if stored_aad != expected_aad.as_slice() {
            return Err(CrypterError::AadMismatch);
        }

        let cipher = Aes256Gcm::new((&km.key).into());
        cipher
            .decrypt(
                nonce,
                aes_gcm::aead::Payload {
                    msg: body,
                    aad: stored_aad,
                },
            )
            .map_err(|e| CrypterError::Aead(format!("decrypt: {e}")))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tenant::scope_tenant;

    fn provider_with(seed: u8, tenant: &str) -> StaticKeyProvider {
        let mut keys = HashMap::new();
        keys.insert(Some(TenantId(tenant.to_string())), [seed; 32]);
        StaticKeyProvider::from_map(keys)
    }

    fn two_tenant_provider(
        seed_a: u8,
        tenant_a: &str,
        seed_b: u8,
        tenant_b: &str,
    ) -> StaticKeyProvider {
        let mut keys = HashMap::new();
        keys.insert(Some(TenantId(tenant_a.to_string())), [seed_a; 32]);
        keys.insert(Some(TenantId(tenant_b.to_string())), [seed_b; 32]);
        StaticKeyProvider::from_map(keys)
    }

    #[tokio::test]
    async fn encrypt_then_decrypt_round_trip() {
        let crypter = TenantKeyedCrypter::new(Arc::new(provider_with(1, "BL_auto")));
        let plaintext = b"sensitive";
        let result = scope_tenant(Some(TenantId("BL_auto".into())), async {
            let ct = crypter.encrypt(plaintext).await.expect("encrypt");
            crypter.decrypt(&ct).await.expect("decrypt")
        })
        .await;
        assert_eq!(result, plaintext);
    }

    #[tokio::test]
    async fn missing_tenant_key_returns_key_destroyed() {
        let crypter =
            TenantKeyedCrypter::new(Arc::new(StaticKeyProvider::from_map(HashMap::new())));
        let err = scope_tenant(Some(TenantId("BL_motor".into())), async {
            crypter.encrypt(b"x").await.unwrap_err()
        })
        .await;
        assert!(matches!(err, CrypterError::KeyDestroyed { .. }));
    }

    /// Cross-tenant decrypt must fail with AadMismatch even when both tenants
    /// happen to share the same key bytes. The AAD binding prevents the
    /// ciphertext from being interpreted under a different tenant context.
    #[tokio::test]
    async fn cross_tenant_decrypt_fails_with_aad_mismatch() {
        // Both tenants use the same underlying key bytes (worst-case collision).
        let crypter = TenantKeyedCrypter::new(Arc::new(two_tenant_provider(
            42, "tenant_a", 42, "tenant_b",
        )));
        let plaintext = b"cross-tenant-secret";

        // Encrypt under tenant_a.
        let ciphertext = scope_tenant(Some(TenantId("tenant_a".into())), async {
            crypter.encrypt(plaintext).await.expect("encrypt")
        })
        .await;

        // Attempt to decrypt under tenant_b — must fail.
        let err = scope_tenant(Some(TenantId("tenant_b".into())), async {
            crypter.decrypt(&ciphertext).await.unwrap_err()
        })
        .await;

        assert!(
            matches!(err, CrypterError::AadMismatch),
            "expected AadMismatch, got {err:?}"
        );
    }

    /// A versioned key provider stub that reports version 1 for the active key.
    struct VersionedProvider {
        tenant: TenantId,
        key: [u8; 32],
        current_version: u8,
    }

    impl KeyProvider for VersionedProvider {
        fn lookup(&self, tenant: &Option<TenantId>) -> Option<[u8; 32]> {
            if tenant.as_ref() == Some(&self.tenant) {
                Some(self.key)
            } else {
                None
            }
        }

        fn lookup_versioned(&self, tenant: &Option<TenantId>) -> Option<KeyMaterial> {
            self.lookup(tenant).map(|key| KeyMaterial {
                key,
                version: self.current_version,
            })
        }
    }

    /// Rotating to a new key version makes old ciphertexts unreadable.
    /// The AEAD tag or AAD mismatch catches the stale version cleanly.
    #[tokio::test]
    async fn old_key_version_decrypt_fails_clean() {
        let tenant = TenantId("tenant_rotation".into());

        // Encrypt with version 0 key.
        let provider_v0 = Arc::new(VersionedProvider {
            tenant: tenant.clone(),
            key: [7u8; 32],
            current_version: 0,
        });
        let crypter_v0 = TenantKeyedCrypter::new(provider_v0);
        let ciphertext = scope_tenant(Some(tenant.clone()), async {
            crypter_v0.encrypt(b"old-secret").await.expect("encrypt v0")
        })
        .await;

        // Attempt decrypt with version 1 key (same key bytes, different version).
        let provider_v1 = Arc::new(VersionedProvider {
            tenant: tenant.clone(),
            key: [7u8; 32],
            current_version: 1,
        });
        let crypter_v1 = TenantKeyedCrypter::new(provider_v1);
        let err = scope_tenant(Some(tenant.clone()), async {
            crypter_v1.decrypt(&ciphertext).await.unwrap_err()
        })
        .await;

        assert!(
            matches!(err, CrypterError::AadMismatch | CrypterError::Aead(_)),
            "expected AadMismatch or Aead on version mismatch, got {err:?}"
        );
    }
}