engram-retired/engrams/engram-crypto/src/engine.rs

/// CryptoEngine — encrypt and decrypt node content with algorithm versioning.
///
/// # Security Model
///
/// - **Symmetric encryption**: AES-256-GCM (authenticated encryption, quantum-resistant)
/// - **Key derivation**: BLAKE3 KDF — stretches the master key into per-operation keys
/// - **Authentication**: BLAKE3 keyed MAC over (algorithm_id || nonce || ciphertext)
/// - **Nonces**: 96-bit random nonce per encryption (from OS CSPRNG via `rand`)
///
/// # AES-256-GCM and Quantum Resistance
///
/// AES-256 is considered quantum-resistant: Grover's algorithm provides at most
/// a quadratic speedup, reducing 256-bit security to ~128-bit effective security
/// against quantum adversaries. 128-bit quantum security is currently considered
/// sufficient. The algorithm_id in EncryptedContent ensures an upgrade to
/// ML-KEM/Kyber is a transparent drop-in when the crates stabilize.
use aes_gcm::{
    aead::{Aead, AeadCore, KeyInit, OsRng},
    Aes256Gcm, Key, Nonce,
};
use serde::{Deserialize, Serialize};

use crate::error::{CryptoError, CryptoResult};
use crate::registry::AlgorithmRegistry;

/// An encrypted content blob, self-describing with its algorithm version.
///
/// The `algorithm_id` field allows any version to decrypt any record,
/// even after algorithm rotation. This is the key to migration-free upgrades.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncryptedContent {
    /// Which algorithm version encrypted this record.
    /// Maps to an entry in `AlgorithmRegistry::versions`.
    pub algorithm_id: String,
    /// The AES-256-GCM ciphertext (includes the GCM auth tag).
    pub ciphertext: Vec<u8>,
    /// In a PQ scheme: the KEM-encapsulated symmetric key.
    /// In the current AES-direct scheme: empty (key is derived from master key + context).
    pub encapsulated_key: Vec<u8>,
    /// 96-bit random AES-GCM nonce.
    pub nonce: Vec<u8>,
    /// BLAKE3 MAC over (algorithm_id || nonce || ciphertext).
    /// In a PQ scheme: this would be a Dilithium/ML-DSA signature.
    pub signature: Vec<u8>,
}

impl EncryptedContent {
    /// Serialize to a compact JSON string for storage.
    pub fn to_bytes(&self) -> CryptoResult<Vec<u8>> {
        serde_json::to_vec(self).map_err(|e| CryptoError::Serialization(e.to_string()))
    }

    /// Deserialize from bytes (JSON).
    pub fn from_bytes(bytes: &[u8]) -> CryptoResult<Self> {
        serde_json::from_slice(bytes).map_err(|e| CryptoError::DecryptionFailed(e.to_string()))
    }
}

/// The encryption engine.
///
/// One engine instance per process (or per-request for stateless usage).
/// The engine holds the master key and the algorithm registry.
pub struct CryptoEngine {
    /// Master key bytes — 32 bytes for AES-256.
    /// In a PQ scheme: this would be a keypair (public/private).
    master_key: [u8; 32],
    /// Algorithm registry — tracks active and historical versions.
    pub registry: AlgorithmRegistry,
}

impl CryptoEngine {
    /// Create an engine from a 32-byte master key (AES-256 requires 256-bit key).
    pub fn from_key(key: &[u8]) -> CryptoResult<Self> {
        if key.len() < 32 {
            return Err(CryptoError::InvalidKeyLength {
                expected: 32,
                got: key.len(),
            });
        }
        let mut master_key = [0u8; 32];
        master_key.copy_from_slice(&key[..32]);
        Ok(Self {
            master_key,
            registry: AlgorithmRegistry::default_registry(),
        })
    }

    /// Create an engine from an environment variable `ENGRAM_ENCRYPTION_KEY`.
    ///
    /// Returns `None` if the variable is not set (dev mode — plaintext storage).
    /// Returns an error if the variable is set but the key is too short.
    pub fn from_env() -> CryptoResult<Option<Self>> {
        match std::env::var("ENGRAM_ENCRYPTION_KEY") {
            Ok(key_str) => {
                let key_bytes = key_str.as_bytes();
                // Derive a 32-byte key from whatever the user provided
                let derived = derive_key(key_bytes, b"engram-master-key");
                Ok(Some(Self::from_key(&derived)?))
            }
            Err(std::env::VarError::NotPresent) => Ok(None),
            Err(e) => Err(CryptoError::KeyDerivation(e.to_string())),
        }
    }

    // ── Encrypt ───────────────────────────────────────────────────────────────

    /// Encrypt `plaintext` using the active algorithm.
    ///
    /// Returns an `EncryptedContent` that is self-describing: it carries its
    /// `algorithm_id`, so decryption never needs out-of-band version tracking.
    pub fn encrypt(&self, plaintext: &[u8]) -> CryptoResult<EncryptedContent> {
        let algorithm_id = self.registry.active_id().to_string();

        // Derive a per-operation encryption key from the master key
        let enc_key_bytes = derive_key(&self.master_key, b"encrypt");
        let key = Key::<Aes256Gcm>::from_slice(&enc_key_bytes);
        let cipher = Aes256Gcm::new(key);

        // Random 96-bit nonce
        let nonce = Aes256Gcm::generate_nonce(&mut OsRng);

        // Encrypt
        let ciphertext = cipher
            .encrypt(&nonce, plaintext)
            .map_err(|e| CryptoError::EncryptionFailed(e.to_string()))?;

        // MAC: BLAKE3 keyed hash over (algorithm_id || nonce || ciphertext)
        let mac_key = derive_key(&self.master_key, b"mac");
        let signature = compute_mac(&mac_key, &algorithm_id, nonce.as_slice(), &ciphertext);

        Ok(EncryptedContent {
            algorithm_id,
            ciphertext,
            encapsulated_key: vec![], // unused in AES-direct mode
            nonce: nonce.to_vec(),
            signature,
        })
    }

    // ── Decrypt ───────────────────────────────────────────────────────────────

    /// Decrypt an `EncryptedContent`, dispatching to the correct algorithm version.
    pub fn decrypt(&self, content: &EncryptedContent) -> CryptoResult<Vec<u8>> {
        // Look up the algorithm version that produced this ciphertext
        let _version = self.registry.get_version(&content.algorithm_id)?;

        // Verify MAC before decrypting (fail-fast on tampering)
        let mac_key = derive_key(&self.master_key, b"mac");
        let expected = compute_mac(&mac_key, &content.algorithm_id, &content.nonce, &content.ciphertext);
        if expected != content.signature {
            return Err(CryptoError::SignatureInvalid);
        }

        // Decrypt
        let enc_key_bytes = derive_key(&self.master_key, b"encrypt");
        let key = Key::<Aes256Gcm>::from_slice(&enc_key_bytes);
        let cipher = Aes256Gcm::new(key);

        if content.nonce.len() != 12 {
            return Err(CryptoError::DecryptionFailed(format!(
                "invalid nonce length: {}",
                content.nonce.len()
            )));
        }
        let nonce = Nonce::from_slice(&content.nonce);

        let plaintext = cipher
            .decrypt(nonce, content.ciphertext.as_slice())
            .map_err(|e| CryptoError::DecryptionFailed(e.to_string()))?;

        Ok(plaintext)
    }

    // ── Signature verification ─────────────────────────────────────────────────

    /// Verify the MAC/signature on an encrypted content blob.
    pub fn verify_signature(&self, content: &EncryptedContent) -> CryptoResult<bool> {
        let mac_key = derive_key(&self.master_key, b"mac");
        let expected = compute_mac(&mac_key, &content.algorithm_id, &content.nonce, &content.ciphertext);
        Ok(expected == content.signature)
    }

    // ── Algorithm rotation ────────────────────────────────────────────────────

    /// Rotate to a new KEM algorithm.
    ///
    /// After rotation, new encryptions use the new algorithm.
    /// Old records retain their `algorithm_id` and decrypt via the historical registry.
    pub fn rotate_algorithm(&mut self, new_kem: crate::algorithm::KemAlgorithm) -> CryptoResult<()> {
        self.registry.rotate_kem(new_kem)
    }
}

// ── Key derivation ────────────────────────────────────────────────────────────

/// Derive a 32-byte sub-key from the master key and a context string.
/// Uses BLAKE3's keyed hash for domain separation.
fn derive_key(master: &[u8], context: &[u8]) -> [u8; 32] {
    // BLAKE3 derive_key: master is the key material, context is the domain
    let mut hasher = blake3::Hasher::new_keyed(
        &padded_32(master),
    );
    hasher.update(context);
    let hash = hasher.finalize();
    *hash.as_bytes()
}

/// Pad or truncate bytes to exactly 32 bytes.
fn padded_32(bytes: &[u8]) -> [u8; 32] {
    let mut out = [0u8; 32];
    let len = bytes.len().min(32);
    out[..len].copy_from_slice(&bytes[..len]);
    out
}

/// Compute a BLAKE3 keyed MAC over (algorithm_id || nonce || ciphertext).
fn compute_mac(mac_key: &[u8; 32], algorithm_id: &str, nonce: &[u8], ciphertext: &[u8]) -> Vec<u8> {
    let mut hasher = blake3::Hasher::new_keyed(mac_key);
    hasher.update(algorithm_id.as_bytes());
    hasher.update(nonce);
    hasher.update(ciphertext);
    hasher.finalize().as_bytes().to_vec()
}

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

    fn make_engine() -> CryptoEngine {
        let key = b"test-master-key-must-be-32-bytes";
        CryptoEngine::from_key(key).unwrap()
    }

    #[test]
    fn test_encrypt_decrypt_roundtrip() {
        let engine = make_engine();
        let plaintext = b"sensitive memory content - do not store unencrypted";
        let enc = engine.encrypt(plaintext).unwrap();
        let dec = engine.decrypt(&enc).unwrap();
        assert_eq!(dec, plaintext);
    }

    #[test]
    fn test_nonce_is_random() {
        let engine = make_engine();
        let enc1 = engine.encrypt(b"same plaintext").unwrap();
        let enc2 = engine.encrypt(b"same plaintext").unwrap();
        // Nonces must differ (probabilistic — collision probability 1/2^96)
        assert_ne!(enc1.nonce, enc2.nonce);
        // Ciphertexts must differ (different nonces → different ciphertexts)
        assert_ne!(enc1.ciphertext, enc2.ciphertext);
    }

    #[test]
    fn test_tampered_ciphertext_rejected() {
        let engine = make_engine();
        let mut enc = engine.encrypt(b"original").unwrap();
        // Flip a byte in the ciphertext
        if let Some(b) = enc.ciphertext.first_mut() {
            *b ^= 0xFF;
        }
        // Decryption should fail due to GCM auth tag verification
        let result = engine.decrypt(&enc);
        assert!(result.is_err());
    }

    #[test]
    fn test_tampered_mac_rejected() {
        let engine = make_engine();
        let mut enc = engine.encrypt(b"original").unwrap();
        // Flip a byte in the signature
        if let Some(b) = enc.signature.first_mut() {
            *b ^= 0xFF;
        }
        let result = engine.decrypt(&enc);
        assert!(result.is_err());
    }

    #[test]
    fn test_verify_signature() {
        let engine = make_engine();
        let enc = engine.encrypt(b"content").unwrap();
        assert!(engine.verify_signature(&enc).unwrap());

        let mut tampered = enc.clone();
        tampered.signature[0] ^= 0x01;
        assert!(!engine.verify_signature(&tampered).unwrap());
    }

    #[test]
    fn test_algorithm_id_stored() {
        let engine = make_engine();
        let enc = engine.encrypt(b"data").unwrap();
        assert_eq!(enc.algorithm_id, "aes256gcm-direct-v1");
    }

    #[test]
    fn test_serialization_roundtrip() {
        let engine = make_engine();
        let enc = engine.encrypt(b"serialize me").unwrap();
        let bytes = enc.to_bytes().unwrap();
        let restored = EncryptedContent::from_bytes(&bytes).unwrap();
        let dec = engine.decrypt(&restored).unwrap();
        assert_eq!(dec, b"serialize me");
    }

    #[test]
    fn test_short_key_rejected() {
        let short_key = b"too-short";
        let result = CryptoEngine::from_key(short_key);
        assert!(result.is_err());
    }

    #[test]
    fn test_empty_plaintext() {
        let engine = make_engine();
        let enc = engine.encrypt(b"").unwrap();
        let dec = engine.decrypt(&enc).unwrap();
        assert_eq!(dec, b"");
    }

    #[test]
    fn test_large_plaintext() {
        let engine = make_engine();
        let plaintext = vec![0xABu8; 1_000_000]; // 1 MB
        let enc = engine.encrypt(&plaintext).unwrap();
        let dec = engine.decrypt(&enc).unwrap();
        assert_eq!(dec, plaintext);
    }

    #[test]
    fn test_unknown_algorithm_rejected() {
        let engine = make_engine();
        let mut enc = engine.encrypt(b"data").unwrap();
        enc.algorithm_id = "unknown-algo-v99".to_string();
        let result = engine.decrypt(&enc);
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), CryptoError::UnknownAlgorithm(_)));
    }
}