seshcookie 0.1.0

//! ChaCha20-Poly1305 AEAD via [`ring::aead::LessSafeKey`].
//!
//! pattern: Functional Core
//!
//! This module exposes three crate-internal primitives: `DerivedKey::derive`
//! (HKDF-SHA256 domain-separated by version and AEAD choice), `seal`
//! (ChaCha20-Poly1305 with a fresh 12-byte nonce), and `try_decrypt`
//! (rotation-aware decrypt against an ordered key list).
//!
//! Nonces are drawn from a [`ring::rand::SystemRandom`] constructed once at
//! [`crate::SessionLayer::new`] time and primed with a one-byte `fill` so the OS
//! RNG is initialized before the first request arrives.
//!
//! The caller-supplied `rng` parameter in `seal` makes the randomness source
//! explicit and injected, keeping the module's logic pure data transformation
//! with no global state, no I/O, and no wall-clock reads. Higher layers compose
//! these primitives with the envelope codec to produce cookie values.
//!
//! ## Wire layout
//!
//! - Sealed payload: `nonce(12) || ciphertext || tag(16)` — fixed-length nonce
//!   prefix, then the ChaCha20-Poly1305 sealed bytes (ciphertext + 16-byte
//!   authentication tag).
//! - Minimum size of any valid sealed payload is therefore 28 bytes (the case where
//!   the plaintext was empty).
//!
//! ## Rotation
//!
//! `try_decrypt` walks an ordered slice of derived keys and returns the
//! first key that authenticates, paired with its index in the slice. Phase 3 uses
//! the index to decide whether to auto-migrate the cookie to the primary key on
//! response.

use ring::{aead, digest, hkdf, rand::SecureRandom};

use crate::keys::{DERIVED_KEY_LEN, HKDF_INFO, HKDF_SALT_CONTEXT};

/// ChaCha20-Poly1305 nonce length (96 bits).
const NONCE_LEN: usize = 12;

/// ChaCha20-Poly1305 authentication tag length (128 bits).
const TAG_LEN: usize = 16;

/// Minimum byte length for any sealed payload: a 12-byte nonce plus a 16-byte tag.
const MIN_SEALED_LEN: usize = NONCE_LEN + TAG_LEN;

/// A ChaCha20-Poly1305 sealing/opening key derived from caller IKM via HKDF-SHA256.
///
/// Constructed via [`DerivedKey::derive`]. The wrapped `LessSafeKey` exposes both
/// `seal_in_place_append_tag` and `open_in_place`; rotation requires both
/// directions, so `LessSafeKey` (rather than the direction-restricted
/// `SealingKey`/`OpeningKey`) is the correct choice.
pub(crate) struct DerivedKey {
    sealing: aead::LessSafeKey,
}

impl DerivedKey {
    /// Derive a ChaCha20-Poly1305 key from the supplied input keying material.
    ///
    /// Uses HKDF-SHA256 with a fixed salt (the SHA-256 of [`HKDF_SALT_CONTEXT`])
    /// and the version-bound [`HKDF_INFO`] string. The salt provides domain
    /// separation; the info parameter binds the derived key to this crate's
    /// wire-format version and AEAD choice.
    pub(crate) fn derive(ikm: &[u8]) -> Self {
        let salt_digest = digest::digest(&digest::SHA256, HKDF_SALT_CONTEXT);
        let salt = hkdf::Salt::new(hkdf::HKDF_SHA256, salt_digest.as_ref());
        let prk = salt.extract(ikm);
        let okm = prk
            .expand(&[HKDF_INFO], &aead::CHACHA20_POLY1305)
            .expect("HKDF-SHA256 expand to a static KeyType cannot fail");
        let mut raw = [0u8; DERIVED_KEY_LEN];
        okm.fill(&mut raw)
            .expect("OKM length matches DERIVED_KEY_LEN by construction");
        let unbound = aead::UnboundKey::new(&aead::CHACHA20_POLY1305, &raw)
            .expect("raw is exactly DERIVED_KEY_LEN bytes, matching the AEAD key length");
        DerivedKey {
            sealing: aead::LessSafeKey::new(unbound),
        }
    }
}

/// Seal `plaintext` under `key` using a fresh random 12-byte nonce drawn from `rng`.
///
/// Output layout: `nonce(12) || ciphertext || tag(16)`. The total length is always
/// exactly `plaintext.len() + 28`.
///
/// Panics only on RNG failure (an OS RNG initialization error, which `ring`
/// surfaces as `Unspecified`); ChaCha20-Poly1305 sealing itself cannot fail for
/// a well-formed key and a 12-byte nonce.
pub(crate) fn seal(key: &DerivedKey, rng: &dyn SecureRandom, plaintext: &[u8]) -> Vec<u8> {
    let mut nonce_bytes = [0u8; NONCE_LEN];
    rng.fill(&mut nonce_bytes)
        .expect("OS RNG must succeed; failure indicates an unrecoverable system fault");

    // Two allocations: one for the seal buffer (plaintext copy that grows by
    // the 16-byte tag), one for the final nonce||ciphertext||tag output.
    // A single-allocation form would prepend the nonce inside the seal buffer
    // and use an offset nonce-position trick with seal_in_place_append_tag,
    // but ring's API does not support sealing a sub-slice in-place. The cost
    // is one extra Vec per call; for the cookie-per-response use case this is
    // negligible, and keeping the code straightforward is preferable.
    let mut sealed_payload: Vec<u8> = plaintext.to_vec();
    let nonce = aead::Nonce::assume_unique_for_key(nonce_bytes);
    key.sealing
        .seal_in_place_append_tag(nonce, aead::Aad::empty(), &mut sealed_payload)
        .expect("ChaCha20-Poly1305 seal cannot fail for a valid key and 12-byte nonce");

    let mut out = Vec::with_capacity(NONCE_LEN + sealed_payload.len());
    out.extend_from_slice(&nonce_bytes);
    out.extend_from_slice(&sealed_payload);
    out
}

/// Attempt to decrypt `payload = nonce(12) || ciphertext || tag(16)` against an
/// ordered list of keys.
///
/// Returns `Some((plaintext, index))` for the first key that authenticates, where
/// `index` is the position of that key in the input slice. Returns `None` if:
///
/// - `payload.len() < 28` (no nonce/tag space — silent fallback, AC6.2)
/// - no key in `keys` authenticates the ciphertext (AC4.5, AC6.3)
///
/// Each key attempt allocates a fresh decryption buffer because `ring`'s
/// `open_in_place` mutates its input. Typical deployments carry 0–2 fallback
/// keys, so the bounded copy cost is acceptable.
pub(crate) fn try_decrypt(keys: &[DerivedKey], payload: &[u8]) -> Option<(Vec<u8>, usize)> {
    if payload.len() < MIN_SEALED_LEN {
        return None;
    }

    let (nonce_slice, ct_plus_tag) = payload.split_at(NONCE_LEN);
    let nonce_array: [u8; NONCE_LEN] = nonce_slice
        .try_into()
        .expect("split_at(NONCE_LEN) yields a slice of exactly NONCE_LEN bytes");

    for (idx, key) in keys.iter().enumerate() {
        let mut buf = ct_plus_tag.to_vec();
        let nonce = aead::Nonce::assume_unique_for_key(nonce_array);
        if let Ok(plain) = key
            .sealing
            .open_in_place(nonce, aead::Aad::empty(), &mut buf)
        {
            let plain_len = plain.len();
            buf.truncate(plain_len);
            return Some((buf, idx));
        }
    }
    None
}

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

    use ring::rand::SystemRandom;
    use std::collections::HashSet;

    /// Two distinct IKMs used as fixtures across most tests. Choosing different
    /// byte patterns ensures HKDF expand produces unrelated keys.
    const IKM_PRIMARY: &[u8; 32] = b"primary-ikm-fixed-bytes-32-len!!";
    const IKM_OTHER_1: &[u8; 32] = b"o1-fallback-ikm-fixed-bytes-32!!";
    const IKM_OTHER_2: &[u8; 32] = b"o2-fallback-ikm-fixed-bytes-32!!";
    const IKM_UNKNOWN: &[u8; 32] = b"x-unknown-ikm-not-in-keylist-32!";

    fn key_from(ikm: &[u8]) -> DerivedKey {
        DerivedKey::derive(ikm)
    }

    // --- AC1.4: bit-flip tamper detection across the entire sealed payload --------

    /// seshcookie-rs.AC1.4: a single-bit flip anywhere in the sealed payload
    /// (nonce, ciphertext, or tag) causes AEAD authentication to fail; no panic
    /// propagates and the handler-visible result is `None`.
    ///
    /// The test iterates every byte position and every bit position within that
    /// byte (8 bits per byte * payload_len bytes), flips exactly that bit, and
    /// confirms `try_decrypt` returns `None`.
    #[test]
    fn try_decrypt_rejects_every_single_bit_flip_ac1_4() {
        let rng = SystemRandom::new();
        let key = key_from(IKM_PRIMARY);
        // Use a non-empty plaintext so the ciphertext region is non-trivial. The
        // exact contents are immaterial; what matters is that we exercise nonce,
        // ciphertext, and tag bit-flips.
        let plaintext: Vec<u8> = (0..64u8).collect();
        let sealed = seal(&key, &rng, &plaintext);

        let keys = [key_from(IKM_PRIMARY)];

        for byte_index in 0..sealed.len() {
            for bit in 0..8u8 {
                let mut tampered = sealed.clone();
                tampered[byte_index] ^= 1 << bit;
                let result = try_decrypt(&keys, &tampered);
                assert!(
                    result.is_none(),
                    "bit flip at byte {byte_index} bit {bit} unexpectedly authenticated"
                );
            }
        }
    }

    // --- AC2.3: tampering inside the issued_at region inside the AEAD plaintext ---

    /// seshcookie-rs.AC2.3: `issued_at` lives inside the AEAD plaintext, so any
    /// tampering with the bytes corresponding to it corrupts the authentication
    /// tag and is rejected. We simulate the layer-1 plaintext shape
    /// (`version(1) || issued_at(8) || payload`) here directly so the test does
    /// not depend on the envelope module.
    #[test]
    fn try_decrypt_rejects_tamper_in_issued_at_region_ac2_3() {
        let rng = SystemRandom::new();
        let key = key_from(IKM_PRIMARY);

        // Synthetic layer-1 plaintext: version byte (1), 8 bytes of issued_at, then
        // arbitrary trailing payload. The exact issued_at value is unused; the test
        // only verifies that tampering inside that 8-byte region defeats the AEAD.
        let mut plaintext = Vec::new();
        plaintext.push(1u8);
        plaintext.extend_from_slice(&42i64.to_le_bytes());
        plaintext.extend_from_slice(b"trailing-payload-bytes");

        let sealed = seal(&key, &rng, &plaintext);

        // Bytes 1..9 of the *plaintext* correspond to bytes
        // (NONCE_LEN + 1) .. (NONCE_LEN + 9) of the *sealed* payload. ChaCha20 is
        // a stream cipher, so the ciphertext byte at position i decrypts to
        // plaintext byte at position i.
        for offset in 1..9 {
            let mut tampered = sealed.clone();
            tampered[NONCE_LEN + offset] ^= 0x01;
            let result = try_decrypt(&[key_from(IKM_PRIMARY)], &tampered);
            assert!(
                result.is_none(),
                "tamper inside issued_at byte {offset} unexpectedly authenticated"
            );
        }
    }

    // --- AC4.1: primary-key decrypt returns index 0 -------------------------------

    /// seshcookie-rs.AC4.1: a cookie sealed under the primary key decrypts via a
    /// keylist whose first entry is that primary, returning index `0`.
    #[test]
    fn try_decrypt_with_primary_returns_index_zero_ac4_1() {
        let rng = SystemRandom::new();
        let key = key_from(IKM_PRIMARY);
        let plaintext = b"hello, primary";
        let sealed = seal(&key, &rng, plaintext);

        let result = try_decrypt(&[key_from(IKM_PRIMARY)], &sealed);
        let (decrypted, idx) = result.expect("primary key must authenticate its own seal");
        assert_eq!(decrypted, plaintext);
        assert_eq!(idx, 0);
    }

    // --- AC4.2: rotation - sealed under fallback, decrypted via fallback index ----

    /// seshcookie-rs.AC4.2: a cookie sealed under an old key still decrypts when
    /// that key is present as a fallback; the returned index identifies which
    /// fallback authenticated.
    #[test]
    fn try_decrypt_falls_back_to_index_one_ac4_2() {
        let rng = SystemRandom::new();
        let old_key = key_from(IKM_OTHER_1);
        let plaintext = b"hello, fallback";
        let sealed = seal(&old_key, &rng, plaintext);

        let keys = [key_from(IKM_PRIMARY), key_from(IKM_OTHER_1)];
        let result = try_decrypt(&keys, &sealed);
        let (decrypted, idx) = result.expect("fallback key must authenticate its own seal");
        assert_eq!(decrypted, plaintext);
        assert_eq!(idx, 1);
    }

    // --- AC4.5: unknown key yields None, no panic ---------------------------------

    /// seshcookie-rs.AC4.5: a cookie sealed under a key not present in the trial
    /// list returns `None`. The test asserts the absence of panic by using
    /// `is_none()` rather than `unwrap`.
    #[test]
    fn try_decrypt_returns_none_for_unknown_key_ac4_5() {
        let rng = SystemRandom::new();
        let unknown_key = key_from(IKM_UNKNOWN);
        let plaintext = b"sealed under an unknown key";
        let sealed = seal(&unknown_key, &rng, plaintext);

        let keys = [key_from(IKM_PRIMARY), key_from(IKM_OTHER_1)];
        let result = try_decrypt(&keys, &sealed);
        assert!(result.is_none());
    }

    // --- AC4.6: position-3 fallback still matches ---------------------------------

    /// seshcookie-rs.AC4.6: with three keys [P, O1, O2] and a cookie sealed under
    /// O2, `try_decrypt` returns the matching plaintext and `index == 2`. The
    /// position of the matching key inside the fallbacks does not prevent a
    /// match — the iterator walks all keys.
    #[test]
    fn try_decrypt_matches_third_key_ac4_6() {
        let rng = SystemRandom::new();
        let o2_key = key_from(IKM_OTHER_2);
        let plaintext = b"sealed under second fallback";
        let sealed = seal(&o2_key, &rng, plaintext);

        let keys = [
            key_from(IKM_PRIMARY),
            key_from(IKM_OTHER_1),
            key_from(IKM_OTHER_2),
        ];
        let result = try_decrypt(&keys, &sealed);
        let (decrypted, idx) = result.expect("third key must authenticate its own seal");
        assert_eq!(decrypted, plaintext);
        assert_eq!(idx, 2);
    }

    // --- AC6.2: too-short inputs produce None -------------------------------------

    /// seshcookie-rs.AC6.2: payloads strictly shorter than 28 bytes (the nonce +
    /// tag minimum) yield `None`. Length 28 itself is also rejected because no
    /// well-formed authenticator can match an arbitrary all-zero buffer; we
    /// assert it as a sanity boundary on the AEAD path.
    #[test]
    fn try_decrypt_rejects_inputs_under_min_len_ac6_2() {
        let key = key_from(IKM_PRIMARY);
        let keys = [key];

        assert!(try_decrypt(&keys, &[]).is_none());
        assert!(try_decrypt(&keys, &[0u8; 27]).is_none());
        // 28 bytes is the minimum length; the all-zero buffer will not authenticate.
        assert!(try_decrypt(&keys, &[0u8; 28]).is_none());
    }

    // --- Sanity: seal output shape ------------------------------------------------

    /// `seal(key, rng, plaintext)` returns exactly `plaintext.len() + 28` bytes:
    /// 12-byte nonce prefix, ciphertext of the same length as plaintext, plus
    /// a 16-byte authentication tag.
    #[test]
    fn seal_output_length_is_plaintext_plus_28() {
        let rng = SystemRandom::new();
        let key = key_from(IKM_PRIMARY);

        for &size in &[0usize, 1, 100, 3000] {
            let plaintext = vec![0xABu8; size];
            let sealed = seal(&key, &rng, &plaintext);
            assert_eq!(
                sealed.len(),
                size + NONCE_LEN + TAG_LEN,
                "seal output length wrong for plaintext size {size}"
            );
        }
    }

    // --- Sanity: round-trip across plaintext sizes --------------------------------

    /// `try_decrypt(&[k], &seal(k, rng, pt))` recovers `pt` exactly for plaintext
    /// sizes spanning empty, small, and "large" (3 KB) inputs.
    #[test]
    fn seal_then_decrypt_round_trip() {
        let rng = SystemRandom::new();
        let key = key_from(IKM_PRIMARY);
        let keys = [key_from(IKM_PRIMARY)];

        for size in [0usize, 1, 16, 100, 1024, 3000] {
            let plaintext: Vec<u8> = (0..size).map(|i| (i % 251) as u8).collect();
            let sealed = seal(&key, &rng, &plaintext);
            let result = try_decrypt(&keys, &sealed);
            let (decrypted, idx) = result.expect("seal -> decrypt must round-trip");
            assert_eq!(decrypted, plaintext, "round-trip mismatch at size {size}");
            assert_eq!(idx, 0);
        }
    }

    // --- Sanity: 100 distinct nonces ----------------------------------------------

    /// 100 successive `seal` calls with the same key produce 100 pairwise-distinct
    /// 12-byte nonce prefixes. With a 96-bit random nonce the birthday-bound
    /// probability of a collision in a sample of 100 is on the order of
    /// 10**-24, so the test is both stable and meaningful for the real RNG.
    #[test]
    fn seal_produces_distinct_nonces() {
        let rng = SystemRandom::new();
        let key = key_from(IKM_PRIMARY);

        let plaintext = b"same plaintext every time";
        let mut nonces: HashSet<[u8; NONCE_LEN]> = HashSet::with_capacity(100);
        for _ in 0..100 {
            let sealed = seal(&key, &rng, plaintext);
            let nonce: [u8; NONCE_LEN] = sealed[..NONCE_LEN]
                .try_into()
                .expect("seal output starts with NONCE_LEN bytes");
            assert!(
                nonces.insert(nonce),
                "duplicate nonce produced across 100 seals"
            );
        }
        assert_eq!(nonces.len(), 100);
    }

    // --- Sanity: derive determinism -----------------------------------------------

    /// `DerivedKey::derive(ikm)` produces a key with deterministic sealing
    /// behavior: a key derived a second time from the same IKM authenticates the
    /// first key's seal output.
    #[test]
    fn derive_is_deterministic_across_invocations() {
        let rng = SystemRandom::new();
        let plaintext = b"determinism probe";

        let key1 = key_from(IKM_PRIMARY);
        let sealed = seal(&key1, &rng, plaintext);

        let key2 = key_from(IKM_PRIMARY);
        let result = try_decrypt(&[key2], &sealed);
        let (decrypted, idx) = result.expect("re-derived key must authenticate the same seal");
        assert_eq!(decrypted, plaintext);
        assert_eq!(idx, 0);
    }

    // --- Sanity: HKDF info domain separation --------------------------------------

    /// A key derived with the production [`HKDF_INFO`] cannot decrypt a payload
    /// sealed under a key derived from the same IKM but a *different* info
    /// string. This exercises the version-binding role of `HKDF_INFO`.
    #[test]
    fn derive_domain_separates_via_info_string() {
        let rng = SystemRandom::new();

        // Production key derivation.
        let production_key = key_from(IKM_PRIMARY);
        let plaintext = b"sealed under production info";
        let sealed = seal(&production_key, &rng, plaintext);

        // Re-derive with a different info string. We only build a probe key here;
        // we are not introducing a parallel public path. This stays in tests.
        let alt_info: &[u8] = b"alt-info-string-not-the-one-used";
        let alt_key = derive_with_alt_info(IKM_PRIMARY, alt_info);

        let result = try_decrypt(&[alt_key], &sealed);
        assert!(
            result.is_none(),
            "key derived with a different info must not authenticate"
        );

        // Sanity: the production key still decrypts its own seal.
        let ok = try_decrypt(&[key_from(IKM_PRIMARY)], &sealed);
        assert!(ok.is_some());
    }

    /// Test-only helper that mirrors `DerivedKey::derive` but takes the HKDF info
    /// string as a parameter, so domain-separation tests can construct a "wrong"
    /// key without exposing a parallel API on `DerivedKey`.
    fn derive_with_alt_info(ikm: &[u8], info: &[u8]) -> DerivedKey {
        let salt_digest = digest::digest(&digest::SHA256, HKDF_SALT_CONTEXT);
        let salt = hkdf::Salt::new(hkdf::HKDF_SHA256, salt_digest.as_ref());
        let prk = salt.extract(ikm);
        let info_parts = [info];
        let okm = prk.expand(&info_parts, &aead::CHACHA20_POLY1305).unwrap();
        let mut raw = [0u8; DERIVED_KEY_LEN];
        okm.fill(&mut raw).unwrap();
        let unbound = aead::UnboundKey::new(&aead::CHACHA20_POLY1305, &raw).unwrap();
        DerivedKey {
            sealing: aead::LessSafeKey::new(unbound),
        }
    }
}