tsafe-core 1.0.12

//! Low-level cryptography primitives for tsafe.
//!
//! Key derivation: Argon2id with tunable cost parameters.  Encryption:
//! XChaCha20-Poly1305 by default, with optional AES-256-GCM support behind the
//! `fips` feature flag.  All secret material is handled via
//! [`Zeroizing`](zeroize::Zeroizing) wrappers so keys are wiped from heap
//! memory on drop.

#[cfg(feature = "fips")]
use aes_gcm::{Aes256Gcm, Nonce};
use argon2::{Algorithm, Argon2, Params, Version};
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
use chacha20poly1305::{
    aead::{Aead, KeyInit},
    XChaCha20Poly1305, XNonce,
};
use hkdf::Hkdf;
use rand::RngCore;
use sha2::Sha256;
use zeroize::{Zeroize, ZeroizeOnDrop};

use tracing::instrument;

use crate::errors::{SafeError, SafeResult};

// KDF defaults — validated against OWASP recommendations for Argon2id.
pub const VAULT_KDF_M_COST: u32 = 65536; // 64 MiB
pub const VAULT_KDF_T_COST: u32 = 3;
pub const VAULT_KDF_P_COST: u32 = 4;
pub const SALT_LEN: usize = 32;
pub const KEY_LEN: usize = 32;
pub const XCHACHA20POLY1305_NONCE_LEN: usize = 24;
pub const NONCE_LEN: usize = XCHACHA20POLY1305_NONCE_LEN; // snap + legacy helpers
#[cfg(feature = "fips")]
pub const AES256GCM_NONCE_LEN: usize = 12;

/// A 256-bit key that is zeroed from memory on drop.
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct VaultKey([u8; KEY_LEN]);

impl VaultKey {
    pub fn as_bytes(&self) -> &[u8; KEY_LEN] {
        &self.0
    }

    /// Construct a VaultKey from raw bytes (e.g. a randomly generated DEK).
    pub fn from_bytes(bytes: [u8; KEY_LEN]) -> Self {
        Self(bytes)
    }
}

/// The on-disk vault format historically used the root key directly. Newer
/// vaults keep the same file format but scope encryption through HKDF-derived
/// purpose keys so challenge/data material do not share a key.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KeySchedule {
    LegacyDirect,
    HkdfSha256V1,
}

/// Cipher selection for vault/team ciphertext on disk.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CipherKind {
    XChaCha20Poly1305,
    #[cfg(feature = "fips")]
    Aes256Gcm,
}

impl CipherKind {
    pub fn as_str(self) -> &'static str {
        match self {
            Self::XChaCha20Poly1305 => "xchacha20poly1305",
            #[cfg(feature = "fips")]
            Self::Aes256Gcm => "aes256gcm",
        }
    }

    pub fn nonce_len(self) -> usize {
        match self {
            Self::XChaCha20Poly1305 => XCHACHA20POLY1305_NONCE_LEN,
            #[cfg(feature = "fips")]
            Self::Aes256Gcm => AES256GCM_NONCE_LEN,
        }
    }
}

/// Purpose labels for HKDF expansion. Adding new purposes must use a distinct
/// label so old ciphertext domains remain isolated.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KeyPurpose {
    SecretData,
    VaultChallenge,
    AuditLog,
    Snapshot,
}

impl KeyPurpose {
    fn label(self) -> &'static str {
        match self {
            Self::SecretData => "tsafe/vault/secret-data/v1",
            Self::VaultChallenge => "tsafe/vault/challenge/v1",
            Self::AuditLog => "tsafe/vault/audit-log/v1",
            Self::Snapshot => "tsafe/vault/snapshot/v1",
        }
    }
}

pub fn default_vault_cipher() -> CipherKind {
    #[cfg(feature = "fips")]
    {
        CipherKind::Aes256Gcm
    }
    #[cfg(not(feature = "fips"))]
    {
        CipherKind::XChaCha20Poly1305
    }
}

pub fn parse_cipher_kind(label: &str) -> SafeResult<CipherKind> {
    match label {
        "xchacha20poly1305" => Ok(CipherKind::XChaCha20Poly1305),
        #[cfg(feature = "fips")]
        "aes256gcm" => Ok(CipherKind::Aes256Gcm),
        #[cfg(not(feature = "fips"))]
        "aes256gcm" => Err(SafeError::InvalidVault {
            reason: "cipher 'aes256gcm' requires a build with the 'fips' feature enabled".into(),
        }),
        other => Err(SafeError::InvalidVault {
            reason: format!("unsupported cipher: '{other}'"),
        }),
    }
}

pub fn random_salt() -> [u8; SALT_LEN] {
    let mut buf = [0u8; SALT_LEN];
    rand::rngs::OsRng.fill_bytes(&mut buf);
    buf
}

pub fn random_nonce() -> [u8; NONCE_LEN] {
    let mut buf = [0u8; NONCE_LEN];
    rand::rngs::OsRng.fill_bytes(&mut buf);
    buf
}

#[cfg(feature = "fips")]
fn random_aes_nonce() -> [u8; AES256GCM_NONCE_LEN] {
    let mut buf = [0u8; AES256GCM_NONCE_LEN];
    rand::rngs::OsRng.fill_bytes(&mut buf);
    buf
}

/// Derive a 256-bit key from a password + salt using Argon2id.
#[instrument(skip(password, salt), fields(m_cost, t_cost, p_cost))]
pub fn derive_key(
    password: &[u8],
    salt: &[u8],
    m_cost: u32,
    t_cost: u32,
    p_cost: u32,
) -> SafeResult<VaultKey> {
    let params =
        Params::new(m_cost, t_cost, p_cost, Some(KEY_LEN)).map_err(|e| SafeError::Crypto {
            context: e.to_string(),
        })?;
    let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, params);
    let mut key_bytes = [0u8; KEY_LEN];
    argon2
        .hash_password_into(password, salt, &mut key_bytes)
        .map_err(|e| SafeError::Crypto {
            context: e.to_string(),
        })?;
    Ok(VaultKey(key_bytes))
}

/// Derive a purpose-scoped 256-bit subkey from a root key using HKDF-SHA256.
pub fn derive_subkey(root_key: &VaultKey, purpose: KeyPurpose) -> SafeResult<VaultKey> {
    derive_labeled_subkey(root_key, purpose.label())
}

/// Derive a 256-bit subkey from a root key using an explicit HKDF label.
pub fn derive_labeled_subkey(root_key: &VaultKey, label: &str) -> SafeResult<VaultKey> {
    let hkdf = Hkdf::<Sha256>::new(None, root_key.as_bytes());
    let mut subkey = [0u8; KEY_LEN];
    hkdf.expand(label.as_bytes(), &mut subkey)
        .map_err(|e| SafeError::Crypto {
            context: format!("hkdf expand for {label}: {e}"),
        })?;
    Ok(VaultKey(subkey))
}

pub fn encrypt_with_key_schedule(
    root_key: &VaultKey,
    schedule: KeySchedule,
    purpose: KeyPurpose,
    cipher: CipherKind,
    plaintext: &[u8],
) -> SafeResult<(Vec<u8>, Vec<u8>)> {
    match schedule {
        KeySchedule::LegacyDirect => encrypt_for_cipher(cipher, root_key, plaintext),
        KeySchedule::HkdfSha256V1 => {
            let subkey = derive_subkey(root_key, purpose)?;
            encrypt_for_cipher(cipher, &subkey, plaintext)
        }
    }
}

pub fn decrypt_with_key_schedule(
    root_key: &VaultKey,
    schedule: KeySchedule,
    purpose: KeyPurpose,
    cipher: CipherKind,
    nonce_bytes: &[u8],
    ciphertext: &[u8],
) -> SafeResult<Vec<u8>> {
    match schedule {
        KeySchedule::LegacyDirect => decrypt_for_cipher(cipher, root_key, nonce_bytes, ciphertext),
        KeySchedule::HkdfSha256V1 => {
            let subkey = derive_subkey(root_key, purpose)?;
            decrypt_for_cipher(cipher, &subkey, nonce_bytes, ciphertext)
        }
    }
}

/// Detect which key schedule was used for a known-plaintext record.
pub fn detect_key_schedule(
    root_key: &VaultKey,
    purpose: KeyPurpose,
    cipher: CipherKind,
    nonce_bytes: &[u8],
    ciphertext: &[u8],
    expected_plaintext: &[u8],
) -> SafeResult<KeySchedule> {
    for schedule in [KeySchedule::HkdfSha256V1, KeySchedule::LegacyDirect] {
        match decrypt_with_key_schedule(
            root_key,
            schedule,
            purpose,
            cipher,
            nonce_bytes,
            ciphertext,
        ) {
            Ok(plaintext) if plaintext.as_slice() == expected_plaintext => return Ok(schedule),
            Ok(_) | Err(SafeError::DecryptionFailed) => continue,
            Err(err) => return Err(err),
        }
    }
    Err(SafeError::DecryptionFailed)
}

pub fn encrypt_for_cipher(
    cipher: CipherKind,
    key: &VaultKey,
    plaintext: &[u8],
) -> SafeResult<(Vec<u8>, Vec<u8>)> {
    match cipher {
        CipherKind::XChaCha20Poly1305 => encrypt(key, plaintext),
        #[cfg(feature = "fips")]
        CipherKind::Aes256Gcm => encrypt_aes_gcm(key, plaintext),
    }
}

pub fn decrypt_for_cipher(
    cipher: CipherKind,
    key: &VaultKey,
    nonce_bytes: &[u8],
    ciphertext: &[u8],
) -> SafeResult<Vec<u8>> {
    match cipher {
        CipherKind::XChaCha20Poly1305 => decrypt(key, nonce_bytes, ciphertext),
        #[cfg(feature = "fips")]
        CipherKind::Aes256Gcm => decrypt_aes_gcm(key, nonce_bytes, ciphertext),
    }
}

/// Encrypt plaintext with XChaCha20-Poly1305. Returns (nonce, ciphertext).
#[instrument(skip_all, fields(plaintext_len = plaintext.len()))]
pub fn encrypt(key: &VaultKey, plaintext: &[u8]) -> SafeResult<(Vec<u8>, Vec<u8>)> {
    let nonce_bytes = random_nonce();
    let nonce = XNonce::from_slice(&nonce_bytes);
    let cipher =
        XChaCha20Poly1305::new_from_slice(key.as_bytes()).map_err(|_| SafeError::Crypto {
            context: "invalid key length".into(),
        })?;
    let ciphertext = cipher
        .encrypt(nonce, plaintext)
        .map_err(|_| SafeError::Crypto {
            context: "encryption failed".into(),
        })?;
    Ok((nonce_bytes.to_vec(), ciphertext))
}

/// Decrypt with XChaCha20-Poly1305. Authentication failure returns `DecryptionFailed`.
#[instrument(skip_all, fields(ciphertext_len = ciphertext.len()))]
pub fn decrypt(key: &VaultKey, nonce_bytes: &[u8], ciphertext: &[u8]) -> SafeResult<Vec<u8>> {
    if nonce_bytes.len() != NONCE_LEN {
        return Err(SafeError::InvalidVault {
            reason: format!(
                "invalid nonce length: expected {NONCE_LEN} bytes, got {}",
                nonce_bytes.len()
            ),
        });
    }
    let nonce = XNonce::from_slice(nonce_bytes);
    let cipher =
        XChaCha20Poly1305::new_from_slice(key.as_bytes()).map_err(|_| SafeError::Crypto {
            context: "invalid key length".into(),
        })?;
    cipher
        .decrypt(nonce, ciphertext)
        .map_err(|_| SafeError::DecryptionFailed)
}

#[cfg(feature = "fips")]
fn encrypt_aes_gcm(key: &VaultKey, plaintext: &[u8]) -> SafeResult<(Vec<u8>, Vec<u8>)> {
    let nonce_bytes = random_aes_nonce();
    let nonce = Nonce::from_slice(&nonce_bytes);
    let cipher = Aes256Gcm::new_from_slice(key.as_bytes()).map_err(|_| SafeError::Crypto {
        context: "invalid AES-256-GCM key length".into(),
    })?;
    let ciphertext = cipher
        .encrypt(nonce, plaintext)
        .map_err(|_| SafeError::Crypto {
            context: "AES-256-GCM encryption failed".into(),
        })?;
    Ok((nonce_bytes.to_vec(), ciphertext))
}

#[cfg(feature = "fips")]
fn decrypt_aes_gcm(key: &VaultKey, nonce_bytes: &[u8], ciphertext: &[u8]) -> SafeResult<Vec<u8>> {
    if nonce_bytes.len() != AES256GCM_NONCE_LEN {
        return Err(SafeError::InvalidVault {
            reason: format!(
                "invalid AES-256-GCM nonce length: expected {AES256GCM_NONCE_LEN} bytes, got {}",
                nonce_bytes.len()
            ),
        });
    }
    let nonce = Nonce::from_slice(nonce_bytes);
    let cipher = Aes256Gcm::new_from_slice(key.as_bytes()).map_err(|_| SafeError::Crypto {
        context: "invalid AES-256-GCM key length".into(),
    })?;
    cipher
        .decrypt(nonce, ciphertext)
        .map_err(|_| SafeError::DecryptionFailed)
}

pub fn encode_b64(data: &[u8]) -> String {
    URL_SAFE_NO_PAD.encode(data)
}

pub fn decode_b64(s: &str) -> SafeResult<Vec<u8>> {
    URL_SAFE_NO_PAD
        .decode(s)
        .map_err(|e| SafeError::InvalidVault {
            reason: format!("base64 decode: {e}"),
        })
}

// ── Zero-knowledge snap crypto ────────────────────────────────────────────────
// The snap server stores only ciphertext. The decryption key lives in the URL
// fragment (#key) and is never transmitted to the server.

/// Encrypt a plaintext string for one-time snap sharing.
///
/// Returns `(blob_b64url, key_b64url)`.
/// - `blob_b64url` — base64url-encoded `nonce(24) || ciphertext` — safe to send to the snap server.
/// - `key_b64url`  — base64url-encoded 32-byte random key — embed in the URL fragment only, **never POST this**.
pub fn snap_encrypt(plaintext: &str) -> SafeResult<(String, String)> {
    let mut key_bytes = [0u8; KEY_LEN];
    rand::rngs::OsRng.fill_bytes(&mut key_bytes);
    let snap_key = VaultKey(key_bytes);
    let (nonce, ciphertext) = encrypt(&snap_key, plaintext.as_bytes())?;
    let mut blob = nonce;
    blob.extend_from_slice(&ciphertext);
    let blob_b64 = encode_b64(&blob);
    let key_b64 = encode_b64(snap_key.as_bytes());
    Ok((blob_b64, key_b64))
}

/// Decrypt a snap blob received from the snap server.
///
/// - `blob_b64`  — the `ciphertext` field from the snap server response.
/// - `key_b64`   — the URL fragment value (after `#`).
pub fn snap_decrypt(blob_b64: &str, key_b64: &str) -> SafeResult<String> {
    let blob = decode_b64(blob_b64).map_err(|_| SafeError::InvalidVault {
        reason: "snap blob is not valid base64url".into(),
    })?;
    let key_bytes = decode_b64(key_b64).map_err(|_| SafeError::InvalidVault {
        reason: "snap key is not valid base64url".into(),
    })?;
    if key_bytes.len() != KEY_LEN {
        return Err(SafeError::InvalidVault {
            reason: format!("snap key must be {KEY_LEN} bytes, got {}", key_bytes.len()),
        });
    }
    if blob.len() < NONCE_LEN {
        return Err(SafeError::InvalidVault {
            reason: "snap blob too short — nonce is missing".into(),
        });
    }
    let key_arr: [u8; KEY_LEN] = key_bytes.try_into().unwrap();
    let snap_key = VaultKey(key_arr);
    let nonce = &blob[..NONCE_LEN];
    let ciphertext = &blob[NONCE_LEN..];
    let plaintext_bytes = decrypt(&snap_key, nonce, ciphertext)?;
    String::from_utf8(plaintext_bytes).map_err(|_| SafeError::InvalidVault {
        reason: "decrypted snap is not valid UTF-8".into(),
    })
}

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

    #[test]
    fn roundtrip_encrypt_decrypt() {
        let salt = random_salt();
        let key = derive_key(
            b"test-password",
            &salt,
            VAULT_KDF_M_COST,
            VAULT_KDF_T_COST,
            VAULT_KDF_P_COST,
        )
        .unwrap();
        let plaintext = b"super-secret-value";
        let (nonce, ct) = encrypt(&key, plaintext).unwrap();
        let pt = decrypt(&key, &nonce, &ct).unwrap();
        assert_eq!(pt, plaintext);
    }

    #[test]
    fn wrong_password_returns_decryption_failed() {
        let salt = random_salt();
        let k1 = derive_key(
            b"correct",
            &salt,
            VAULT_KDF_M_COST,
            VAULT_KDF_T_COST,
            VAULT_KDF_P_COST,
        )
        .unwrap();
        let k2 = derive_key(
            b"wrong",
            &salt,
            VAULT_KDF_M_COST,
            VAULT_KDF_T_COST,
            VAULT_KDF_P_COST,
        )
        .unwrap();
        let (nonce, ct) = encrypt(&k1, b"data").unwrap();
        let result = decrypt(&k2, &nonce, &ct);
        assert!(matches!(result, Err(SafeError::DecryptionFailed)));
    }

    #[test]
    fn nonces_are_unique() {
        let n1 = random_nonce();
        let n2 = random_nonce();
        assert_ne!(n1, n2);
    }

    #[test]
    fn b64_roundtrip() {
        let data = b"hello world \x00\xff";
        assert_eq!(decode_b64(&encode_b64(data)).unwrap(), data);
    }

    #[test]
    fn hkdf_subkeys_are_domain_separated_and_deterministic() {
        let salt = random_salt();
        let root = derive_key(
            b"test-password",
            &salt,
            VAULT_KDF_M_COST,
            VAULT_KDF_T_COST,
            VAULT_KDF_P_COST,
        )
        .unwrap();
        let enc_1 = derive_subkey(&root, KeyPurpose::SecretData).unwrap();
        let enc_2 = derive_subkey(&root, KeyPurpose::SecretData).unwrap();
        let challenge = derive_subkey(&root, KeyPurpose::VaultChallenge).unwrap();
        let audit = derive_subkey(&root, KeyPurpose::AuditLog).unwrap();

        assert_eq!(enc_1.as_bytes(), enc_2.as_bytes());
        assert_ne!(enc_1.as_bytes(), challenge.as_bytes());
        assert_ne!(enc_1.as_bytes(), audit.as_bytes());
        assert_ne!(challenge.as_bytes(), audit.as_bytes());
    }

    #[test]
    fn detect_key_schedule_recognizes_hkdf_and_legacy_ciphertexts() {
        let salt = random_salt();
        let root = derive_key(
            b"test-password",
            &salt,
            VAULT_KDF_M_COST,
            VAULT_KDF_T_COST,
            VAULT_KDF_P_COST,
        )
        .unwrap();
        let expected = b"known-plaintext";

        let (hkdf_nonce, hkdf_ct) = encrypt_with_key_schedule(
            &root,
            KeySchedule::HkdfSha256V1,
            KeyPurpose::VaultChallenge,
            CipherKind::XChaCha20Poly1305,
            expected,
        )
        .unwrap();
        assert_eq!(
            detect_key_schedule(
                &root,
                KeyPurpose::VaultChallenge,
                CipherKind::XChaCha20Poly1305,
                &hkdf_nonce,
                &hkdf_ct,
                expected
            )
            .unwrap(),
            KeySchedule::HkdfSha256V1
        );

        let (legacy_nonce, legacy_ct) = encrypt_with_key_schedule(
            &root,
            KeySchedule::LegacyDirect,
            KeyPurpose::VaultChallenge,
            CipherKind::XChaCha20Poly1305,
            expected,
        )
        .unwrap();
        assert_eq!(
            detect_key_schedule(
                &root,
                KeyPurpose::VaultChallenge,
                CipherKind::XChaCha20Poly1305,
                &legacy_nonce,
                &legacy_ct,
                expected
            )
            .unwrap(),
            KeySchedule::LegacyDirect
        );
    }

    #[test]
    fn parse_cipher_kind_accepts_legacy_cipher() {
        assert_eq!(
            parse_cipher_kind("xchacha20poly1305").unwrap(),
            CipherKind::XChaCha20Poly1305
        );
    }

    #[cfg(feature = "fips")]
    #[test]
    fn default_cipher_is_aes_gcm_in_fips_builds() {
        assert_eq!(default_vault_cipher(), CipherKind::Aes256Gcm);
        assert_eq!(CipherKind::Aes256Gcm.nonce_len(), AES256GCM_NONCE_LEN);
    }

    #[cfg(not(feature = "fips"))]
    #[test]
    fn default_cipher_is_xchacha_without_fips() {
        assert_eq!(default_vault_cipher(), CipherKind::XChaCha20Poly1305);
    }

    #[cfg(feature = "fips")]
    #[test]
    fn aes_cipher_roundtrip_works_when_fips_enabled() {
        let salt = random_salt();
        let key = derive_key(
            b"test-password",
            &salt,
            VAULT_KDF_M_COST,
            VAULT_KDF_T_COST,
            VAULT_KDF_P_COST,
        )
        .unwrap();
        let plaintext = b"fips-mode-value";
        let (nonce, ct) = encrypt_for_cipher(CipherKind::Aes256Gcm, &key, plaintext).unwrap();
        let pt = decrypt_for_cipher(CipherKind::Aes256Gcm, &key, &nonce, &ct).unwrap();
        assert_eq!(pt, plaintext);
    }

    // ── Adversarial / negative-path tests ─────────────────────────────────────

    fn test_key() -> VaultKey {
        let salt = random_salt();
        derive_key(b"test-pw", &salt, 8192, 1, 1).unwrap()
    }

    /// Flipping any byte in the auth tag must cause authentication failure.
    #[test]
    fn tampered_ciphertext_tag_returns_decryption_failed() {
        let key = test_key();
        let (nonce, mut ct) = encrypt(&key, b"sensitive").unwrap();
        let last = ct.len() - 1;
        ct[last] ^= 0xff;
        let result = decrypt(&key, &nonce, &ct);
        assert!(matches!(result, Err(SafeError::DecryptionFailed)));
    }

    /// Flipping a byte in the ciphertext body (not the tag) also fails.
    #[test]
    fn tampered_ciphertext_body_returns_decryption_failed() {
        let key = test_key();
        let (nonce, mut ct) = encrypt(&key, b"another-secret-value").unwrap();
        ct[0] ^= 0x01;
        let result = decrypt(&key, &nonce, &ct);
        assert!(matches!(result, Err(SafeError::DecryptionFailed)));
    }

    /// Empty ciphertext (no auth tag) is rejected cleanly.
    #[test]
    fn empty_ciphertext_returns_decryption_failed() {
        let key = test_key();
        let nonce = random_nonce();
        let result = decrypt(&key, &nonce, &[]);
        assert!(matches!(result, Err(SafeError::DecryptionFailed)));
    }

    /// A nonce of the wrong length is rejected with InvalidVault (not a panic).
    #[test]
    fn wrong_nonce_length_returns_invalid_vault() {
        let key = test_key();
        let (_, ct) = encrypt(&key, b"data").unwrap();
        let short_nonce = [0u8; 12]; // 12 bytes instead of the required 24
        let result = decrypt(&key, &short_nonce, &ct);
        assert!(matches!(result, Err(SafeError::InvalidVault { .. })));
    }

    /// Correct key but a mismatched nonce must fail AEAD authentication.
    #[test]
    fn correct_key_wrong_nonce_returns_decryption_failed() {
        let key = test_key();
        let (_, ct) = encrypt(&key, b"data").unwrap();
        let wrong_nonce = random_nonce();
        let result = decrypt(&key, &wrong_nonce, &ct);
        assert!(matches!(result, Err(SafeError::DecryptionFailed)));
    }

    /// KeyPurpose segregation: SecretData ciphertext cannot be decrypted with
    /// the VaultChallenge subkey derived from the same root key.
    #[test]
    fn keypurpose_segregation_prevents_cross_purpose_decryption() {
        let root = test_key();
        let plaintext = b"cross-purpose-test";

        let (nonce, ct) = encrypt_with_key_schedule(
            &root,
            KeySchedule::HkdfSha256V1,
            KeyPurpose::SecretData,
            CipherKind::XChaCha20Poly1305,
            plaintext,
        )
        .unwrap();

        let result = decrypt_with_key_schedule(
            &root,
            KeySchedule::HkdfSha256V1,
            KeyPurpose::VaultChallenge,
            CipherKind::XChaCha20Poly1305,
            &nonce,
            &ct,
        );
        assert!(matches!(result, Err(SafeError::DecryptionFailed)));
    }

    /// HkdfSha256V1 ciphertext is rejected when fed to LegacyDirect even with the same root.
    #[test]
    fn hkdf_schedule_ciphertext_rejected_by_legacy_direct() {
        let root = test_key();
        let plaintext = b"schedule-isolation";

        let (nonce, ct) = encrypt_with_key_schedule(
            &root,
            KeySchedule::HkdfSha256V1,
            KeyPurpose::SecretData,
            CipherKind::XChaCha20Poly1305,
            plaintext,
        )
        .unwrap();

        let result = decrypt_with_key_schedule(
            &root,
            KeySchedule::LegacyDirect,
            KeyPurpose::SecretData,
            CipherKind::XChaCha20Poly1305,
            &nonce,
            &ct,
        );
        assert!(matches!(result, Err(SafeError::DecryptionFailed)));
    }

    /// detect_key_schedule returns DecryptionFailed when neither schedule matches.
    #[test]
    fn detect_key_schedule_with_wrong_key_returns_decryption_failed() {
        let enc_key = test_key();
        let dec_key = test_key(); // entirely different key
        let plaintext = b"detection-test";

        let (nonce, ct) = encrypt_with_key_schedule(
            &enc_key,
            KeySchedule::HkdfSha256V1,
            KeyPurpose::SecretData,
            CipherKind::XChaCha20Poly1305,
            plaintext,
        )
        .unwrap();

        let result = detect_key_schedule(
            &dec_key,
            KeyPurpose::SecretData,
            CipherKind::XChaCha20Poly1305,
            &nonce,
            &ct,
            plaintext,
        );
        assert!(matches!(result, Err(SafeError::DecryptionFailed)));
    }

    /// All four KeyPurpose subkeys derived from the same root are mutually distinct.
    #[test]
    fn all_keypurpose_subkeys_are_distinct() {
        let root = test_key();
        let sd = derive_subkey(&root, KeyPurpose::SecretData).unwrap();
        let vc = derive_subkey(&root, KeyPurpose::VaultChallenge).unwrap();
        let al = derive_subkey(&root, KeyPurpose::AuditLog).unwrap();
        let sn = derive_subkey(&root, KeyPurpose::Snapshot).unwrap();

        let keys = [sd.as_bytes(), vc.as_bytes(), al.as_bytes(), sn.as_bytes()];
        for i in 0..keys.len() {
            for j in (i + 1)..keys.len() {
                assert_ne!(keys[i], keys[j], "subkeys[{i}] and subkeys[{j}] collided");
            }
        }
    }

    // ── snap_encrypt / snap_decrypt adversarial paths ─────────────────────────

    #[test]
    fn snap_roundtrip() {
        let plaintext = "my one-time secret";
        let (blob, key) = snap_encrypt(plaintext).unwrap();
        let recovered = snap_decrypt(&blob, &key).unwrap();
        assert_eq!(recovered, plaintext);
    }

    #[test]
    fn snap_decrypt_tampered_blob_returns_decryption_failed() {
        let (mut blob_b64, key_b64) = snap_encrypt("value").unwrap();
        let mut blob = decode_b64(&blob_b64).unwrap();
        let last = blob.len() - 1;
        blob[last] ^= 0xff;
        blob_b64 = encode_b64(&blob);
        let result = snap_decrypt(&blob_b64, &key_b64);
        assert!(matches!(result, Err(SafeError::DecryptionFailed)));
    }

    #[test]
    fn snap_decrypt_wrong_key_returns_decryption_failed() {
        let (blob_b64, _) = snap_encrypt("value").unwrap();
        let (_, wrong_key_b64) = snap_encrypt("other").unwrap();
        let result = snap_decrypt(&blob_b64, &wrong_key_b64);
        assert!(matches!(result, Err(SafeError::DecryptionFailed)));
    }

    #[test]
    fn snap_decrypt_truncated_blob_returns_invalid_vault() {
        let short_blob = encode_b64(&[0u8; 4]); // shorter than NONCE_LEN (24)
        let (_, key_b64) = snap_encrypt("value").unwrap();
        let result = snap_decrypt(&short_blob, &key_b64);
        assert!(matches!(result, Err(SafeError::InvalidVault { .. })));
    }

    #[test]
    fn snap_decrypt_invalid_base64_blob_returns_invalid_vault() {
        let (_, key_b64) = snap_encrypt("value").unwrap();
        let result = snap_decrypt("!!!not-base64!!!", &key_b64);
        assert!(matches!(result, Err(SafeError::InvalidVault { .. })));
    }

    #[test]
    fn snap_decrypt_invalid_base64_key_returns_invalid_vault() {
        let (blob_b64, _) = snap_encrypt("value").unwrap();
        let result = snap_decrypt(&blob_b64, "!!!not-base64!!!");
        assert!(matches!(result, Err(SafeError::InvalidVault { .. })));
    }

    #[test]
    fn snap_decrypt_short_key_returns_invalid_vault() {
        let (blob_b64, _) = snap_encrypt("value").unwrap();
        let short_key = encode_b64(&[0u8; 16]); // 16 bytes, expected 32
        let result = snap_decrypt(&blob_b64, &short_key);
        assert!(matches!(result, Err(SafeError::InvalidVault { .. })));
    }

    // ── Key schedule full roundtrip for every purpose ─────────────────────────

    #[test]
    fn hkdf_schedule_roundtrip_for_all_purposes() {
        let root = test_key();
        let plaintext = b"purpose-roundtrip";
        for purpose in [
            KeyPurpose::SecretData,
            KeyPurpose::VaultChallenge,
            KeyPurpose::AuditLog,
            KeyPurpose::Snapshot,
        ] {
            let (nonce, ct) = encrypt_with_key_schedule(
                &root,
                KeySchedule::HkdfSha256V1,
                purpose,
                CipherKind::XChaCha20Poly1305,
                plaintext,
            )
            .unwrap();
            let pt = decrypt_with_key_schedule(
                &root,
                KeySchedule::HkdfSha256V1,
                purpose,
                CipherKind::XChaCha20Poly1305,
                &nonce,
                &ct,
            )
            .unwrap();
            assert_eq!(pt, plaintext, "roundtrip failed for {purpose:?}");
        }
    }

    #[test]
    fn legacy_direct_schedule_roundtrip() {
        let root = test_key();
        let plaintext = b"legacy-data";
        let (nonce, ct) = encrypt_with_key_schedule(
            &root,
            KeySchedule::LegacyDirect,
            KeyPurpose::SecretData,
            CipherKind::XChaCha20Poly1305,
            plaintext,
        )
        .unwrap();
        let pt = decrypt_with_key_schedule(
            &root,
            KeySchedule::LegacyDirect,
            KeyPurpose::SecretData,
            CipherKind::XChaCha20Poly1305,
            &nonce,
            &ct,
        )
        .unwrap();
        assert_eq!(pt, plaintext);
    }

    /// parse_cipher_kind rejects unknown labels.
    #[test]
    fn parse_cipher_kind_rejects_unknown_label() {
        let result = parse_cipher_kind("chacha8");
        assert!(matches!(result, Err(SafeError::InvalidVault { .. })));
    }

    /// parse_cipher_kind rejects empty string.
    #[test]
    fn parse_cipher_kind_rejects_empty_string() {
        let result = parse_cipher_kind("");
        assert!(matches!(result, Err(SafeError::InvalidVault { .. })));
    }
}