krypteia-arcana 0.1.0

//! AES-CCM AEAD (NIST SP 800-38C, RFC 3610).
//!
//! CCM = Counter with CBC-MAC. It is the second AES AEAD shipped by
//! `arcana` alongside AES-GCM, and the older of the two.
//! CCM is the AEAD used by Bluetooth Low Energy, ZigBee, IPsec
//! ESP-CCM, TLS 1.2 ciphersuites of the form `*_CCM*`, 802.15.4
//! mesh networks, and IETF protocols that need a small / portable
//! AEAD with no GHASH dependency.
//!
//! Compared to GCM:
//!
//! * **Pro**: only needs `encrypt_block`. Decryption never invokes
//!   AES decrypt -- it XORs the same CTR keystream and recomputes
//!   the MAC over the plaintext. Smaller code, smaller embedded
//!   footprint, no GHASH state.
//! * **Pro**: no soft-failure modes -- the construction is a
//!   straightforward MAC-then-encrypt with explicit length fields.
//! * **Con**: two-pass over the data (one for MAC, one for CTR), so
//!   it's slower than GCM on hardware that has carry-less multiply.
//! * **Con**: tag length and the unusual `(N, L)` parameter trade
//!   are easy to misconfigure -- see the parameter validation below.
//!
//! # Parameters
//!
//! CCM is parameterised by `(M, L)` where:
//!
//! * **`M`** is the **tag length in bytes**. Allowed: `{4, 6, 8, 10, 12, 14, 16}`.
//!   Smaller M = faster but lower forgery resistance. Most protocols
//!   use M = 8 (Bluetooth, ZigBee) or M = 16 (TLS, ESP).
//!
//! * **`L`** is the **length-of-length field in bytes**. Allowed:
//!   `{2, 3, 4, 5, 6, 7, 8}`. The maximum payload length is
//!   `2^(8*L) - 1` bytes; the nonce length is `15 - L` bytes.
//!   `L = 2` gives a 13-byte nonce and a max payload of ~64 KB,
//!   which is the choice in **RFC 3610** and almost every modern
//!   protocol that uses CCM. The convenience constructors below
//!   pin `L = 2`.
//!
//! # Construction (NIST SP 800-38C / RFC 3610)
//!
//! 1. **Format the first MAC block** `B0`:
//!
//!    ```text
//!     +-----+----------+----------+
//!     | flg | nonce N  | length Q |
//!     +-----+----------+----------+
//!       1B    15 - L      L
//!    ```
//!
//!    where `flg = (Adata << 6) | ((M-2)/2 << 3) | (L-1)`,
//!    `Adata = 1` iff AAD is non-empty, and `Q` is the big-endian
//!    payload length encoded in `L` bytes.
//!
//! 2. **Format the AAD blocks**: prepend `length(AAD)` encoded as
//!    2 bytes (BE) for `0 < len < 2^16 - 2^8`, or 6 bytes for the
//!    larger ranges (we only support the 2-byte form: caller AAD
//!    must fit in the canonical IETF range), then the AAD bytes,
//!    then zero-pad to a multiple of 16.
//!
//! 3. **Format the payload blocks**: just the plaintext bytes, zero-
//!    padded to a multiple of 16.
//!
//! 4. **CBC-MAC** all of the above (B0 || formatted_aad || formatted_payload)
//!    with AES under the key. The CBC-MAC output is the unencrypted
//!    tag T (taking the first M bytes of the final 16-byte CBC state).
//!
//! 5. **Counter blocks** A_i: format
//!
//!    ```text
//!     +-----+----------+--------+
//!     | flg | nonce N  | i (BE) |
//!     +-----+----------+--------+
//!       1B    15 - L      L
//!    ```
//!
//!    where `flg = L - 1`. `S_i = AES_K(A_i)` is the keystream block.
//!
//! 6. **Encrypt T** by XOR with the first M bytes of `S_0`. The result
//!    is the on-the-wire tag.
//!
//! 7. **Encrypt the payload** by XOR with `S_1, S_2, ...` (counter
//!    starts at 1).
//!
//! Decryption reverses steps 6-7 to recover the plaintext, then
//! recomputes T via steps 1-4 and compares in constant time.

use super::aes::Aes;
use crate::BlockCipher;

// ============================================================================
// Parameter validation
// ============================================================================

/// Validate CCM `(M, L)` parameters per NIST SP 800-38C / RFC 3610.
///
/// Returns an error string for diagnostic use; the public API
/// surfaces this as `None` from `encrypt`/`decrypt`. We use a string
/// here only as a debug aid for the test suite.
fn check_params(m: usize, l: usize) -> Result<(), &'static str> {
    if !matches!(m, 4 | 6 | 8 | 10 | 12 | 14 | 16) {
        return Err("CCM: tag length M must be in {4,6,8,10,12,14,16}");
    }
    if !(2..=8).contains(&l) {
        return Err("CCM: length-of-length L must be in {2..=8}");
    }
    Ok(())
}

// ============================================================================
// Generic AES-CCM core (any L, any M, any AES variant via Aes)
// ============================================================================

/// AES-CCM encrypt with the generic `(M, L)` parameters.
///
/// `nonce.len()` must be exactly `15 - L`. `plaintext.len()` must
/// fit in `2^(8*L)` bytes (no overflow check is enforced for `L >=
/// 8` since usize is bounded). AAD must currently be < `2^16 - 2^8`
/// bytes (the canonical 2-byte length encoding).
///
/// Returns `(ciphertext, tag)` where `ciphertext.len() ==
/// plaintext.len()` and `tag.len() == M`.
pub fn ccm_encrypt(
    aes: &Aes,
    m: usize,
    l: usize,
    nonce: &[u8],
    aad: &[u8],
    plaintext: &[u8],
) -> Option<(Vec<u8>, Vec<u8>)> {
    check_params(m, l).ok()?;
    if nonce.len() != 15 - l {
        return None;
    }
    if l < 8 {
        // Plaintext length must fit in L bytes.
        let max_pt: u128 = 1u128 << (8 * l);
        if (plaintext.len() as u128) >= max_pt {
            return None;
        }
    }
    if aad.len() >= (1usize << 16) - (1usize << 8) {
        // AAD too long for the 2-byte length encoding we support.
        return None;
    }

    // Step 1-4: compute CBC-MAC over (B0 || formatted_aad ||
    // formatted_payload). The result is the unencrypted tag T.
    let t = cbc_mac(aes, m, l, nonce, aad, plaintext);

    // Step 5-6: compute the keystream block A_0 = AES_K(format(0))
    // and XOR its first M bytes with T to produce the encrypted tag.
    let mut a0 = ctr_block(l, nonce, 0);
    aes.encrypt_block(&mut a0);
    let mut tag = vec![0u8; m];
    for i in 0..m {
        tag[i] = t[i] ^ a0[i];
    }

    // Step 7: CTR encrypt the payload starting at i = 1.
    let mut ct = plaintext.to_vec();
    let mut counter: u64 = 1;
    let mut pos = 0;
    while pos < ct.len() {
        let mut block = ctr_block(l, nonce, counter);
        aes.encrypt_block(&mut block);
        let take = (16).min(ct.len() - pos);
        for i in 0..take {
            ct[pos + i] ^= block[i];
        }
        pos += 16;
        counter += 1;
    }

    Some((ct, tag))
}

/// AES-CCM decrypt with the generic `(M, L)` parameters.
///
/// Returns `Some(plaintext)` only if the recomputed tag matches
/// (constant-time compare). Returns `None` for any malformed input
/// (wrong nonce length, wrong tag length, parameter out of range,
/// AAD too long, payload too long) **and** for tag mismatch.
///
/// Callers MUST treat `None` as a hard authentication failure and
/// MUST NOT use the (intermediate) decrypted bytes for any purpose
/// even if they were exposed by an aggressive optimiser -- the
/// function does not leak them.
pub fn ccm_decrypt(
    aes: &Aes,
    m: usize,
    l: usize,
    nonce: &[u8],
    aad: &[u8],
    ciphertext: &[u8],
    tag: &[u8],
) -> Option<Vec<u8>> {
    check_params(m, l).ok()?;
    if nonce.len() != 15 - l {
        return None;
    }
    if tag.len() != m {
        return None;
    }
    if l < 8 {
        let max_pt: u128 = 1u128 << (8 * l);
        if (ciphertext.len() as u128) >= max_pt {
            return None;
        }
    }
    if aad.len() >= (1usize << 16) - (1usize << 8) {
        return None;
    }

    // CTR-decrypt the payload (= same XOR as encrypt).
    let mut pt = ciphertext.to_vec();
    let mut counter: u64 = 1;
    let mut pos = 0;
    while pos < pt.len() {
        let mut block = ctr_block(l, nonce, counter);
        aes.encrypt_block(&mut block);
        let take = (16).min(pt.len() - pos);
        for i in 0..take {
            pt[pos + i] ^= block[i];
        }
        pos += 16;
        counter += 1;
    }

    // Recompute the tag over the recovered plaintext.
    let t = cbc_mac(aes, m, l, nonce, aad, &pt);
    let mut a0 = ctr_block(l, nonce, 0);
    aes.encrypt_block(&mut a0);
    let mut expected = vec![0u8; m];
    for i in 0..m {
        expected[i] = t[i] ^ a0[i];
    }

    // Constant-time compare the tag.
    let mut diff = 0u8;
    for i in 0..m {
        diff |= expected[i] ^ tag[i];
    }
    if diff != 0 {
        return None;
    }

    Some(pt)
}

// ============================================================================
// Internal helpers
// ============================================================================

/// Build the CCM CTR-block format (NIST SP 800-38C §6.2):
///
/// ```text
///   flg = L - 1
///   block = flg || nonce || counter (BE on L bytes)
/// ```
fn ctr_block(l: usize, nonce: &[u8], counter: u64) -> [u8; 16] {
    debug_assert_eq!(nonce.len(), 15 - l);
    let mut block = [0u8; 16];
    block[0] = (l - 1) as u8;
    block[1..1 + nonce.len()].copy_from_slice(nonce);
    // BE counter in the last L bytes.
    let ctr_be = counter.to_be_bytes();
    // Pick the trailing L bytes of ctr_be (a u64 has 8 bytes, so
    // for L <= 8 we always have enough).
    let l_used = l.min(8);
    block[16 - l_used..].copy_from_slice(&ctr_be[8 - l_used..]);
    block
}

/// Build the CCM B0 first MAC block (NIST SP 800-38C §A.2.1):
///
/// ```text
///   flg = (Adata << 6) | (((M - 2) / 2) << 3) | (L - 1)
///   B0  = flg || nonce || Q (BE on L bytes)
/// ```
fn b0_block(m: usize, l: usize, nonce: &[u8], aad_len: usize, payload_len: usize) -> [u8; 16] {
    debug_assert_eq!(nonce.len(), 15 - l);
    let adata: u8 = if aad_len > 0 { 1 } else { 0 };
    let m_field: u8 = (((m as u8) - 2) / 2) << 3;
    let l_field: u8 = (l as u8) - 1;
    let flg: u8 = (adata << 6) | m_field | l_field;

    let mut b0 = [0u8; 16];
    b0[0] = flg;
    b0[1..1 + nonce.len()].copy_from_slice(nonce);

    // Q = payload_len encoded BE in L bytes.
    let q_be = (payload_len as u64).to_be_bytes();
    let l_used = l.min(8);
    b0[16 - l_used..].copy_from_slice(&q_be[8 - l_used..]);
    b0
}

/// Run CBC-MAC over (B0 || formatted_aad || formatted_payload) and
/// return the final 16-byte block. The caller will truncate to M
/// bytes.
fn cbc_mac(aes: &Aes, m: usize, l: usize, nonce: &[u8], aad: &[u8], payload: &[u8]) -> [u8; 16] {
    let mut state = b0_block(m, l, nonce, aad.len(), payload.len());
    aes.encrypt_block(&mut state);

    // Format AAD: 2-byte BE length || aad || zero pad to 16-byte multiple.
    if !aad.is_empty() {
        // Build the prefix block: 2 bytes of BE length followed by
        // up to 14 bytes of AAD (the rest spills into subsequent
        // blocks).
        let len_bytes = (aad.len() as u16).to_be_bytes();
        let mut prefix = [0u8; 16];
        prefix[0] = len_bytes[0];
        prefix[1] = len_bytes[1];
        let take = (14).min(aad.len());
        prefix[2..2 + take].copy_from_slice(&aad[..take]);
        for i in 0..16 {
            state[i] ^= prefix[i];
        }
        aes.encrypt_block(&mut state);

        let mut pos = take;
        while pos < aad.len() {
            let mut block = [0u8; 16];
            let chunk = (16).min(aad.len() - pos);
            block[..chunk].copy_from_slice(&aad[pos..pos + chunk]);
            for i in 0..16 {
                state[i] ^= block[i];
            }
            aes.encrypt_block(&mut state);
            pos += chunk;
        }
    }

    // Format payload: just the bytes, zero-padded to a 16-byte multiple.
    let mut pos = 0;
    while pos < payload.len() {
        let mut block = [0u8; 16];
        let chunk = (16).min(payload.len() - pos);
        block[..chunk].copy_from_slice(&payload[pos..pos + chunk]);
        for i in 0..16 {
            state[i] ^= block[i];
        }
        aes.encrypt_block(&mut state);
        pos += chunk;
    }

    state
}

// ============================================================================
// Public convenience: AES-128-CCM with L = 2 (the RFC 3610 / IETF default)
// ============================================================================

/// AES-CCM with the **RFC 3610 default parameters**: `L = 2`, so
/// the nonce is 13 bytes and the maximum payload is `2^16 - 1`
/// bytes. The tag length `M` is configurable per call.
///
/// This is the wrapper most callers want. The generic `(M, L)`
/// form is exposed via [`ccm_encrypt`] / [`ccm_decrypt`] for
/// protocols that need a different `L`.
pub struct AesCcm {
    aes: Aes,
    m: usize,
}

impl AesCcm {
    /// Initialise AES-CCM with `M` ∈ `{4, 6, 8, 10, 12, 14, 16}`.
    /// Returns `None` for invalid `M` or invalid AES key length.
    pub fn new(key: &[u8], m: usize) -> Option<Self> {
        if !matches!(m, 4 | 6 | 8 | 10 | 12 | 14 | 16) {
            return None;
        }
        if !matches!(key.len(), 16 | 24 | 32) {
            return None;
        }
        Some(Self {
            aes: <Aes as BlockCipher>::new(key),
            m,
        })
    }

    /// Encrypt and authenticate. `nonce.len() == 13`.
    /// Returns `(ciphertext, tag)` where `tag.len() == M`.
    pub fn encrypt(&self, nonce: &[u8; 13], aad: &[u8], plaintext: &[u8]) -> Option<(Vec<u8>, Vec<u8>)> {
        ccm_encrypt(&self.aes, self.m, 2, nonce, aad, plaintext)
    }

    /// Decrypt and verify. `nonce.len() == 13`, `tag.len() == M`.
    /// Returns `None` on tag mismatch or any malformed input.
    pub fn decrypt(&self, nonce: &[u8; 13], aad: &[u8], ciphertext: &[u8], tag: &[u8]) -> Option<Vec<u8>> {
        ccm_decrypt(&self.aes, self.m, 2, nonce, aad, ciphertext, tag)
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    fn hex(s: &str) -> Vec<u8> {
        let s: String = s.chars().filter(|c| !c.is_whitespace()).collect();
        assert!(s.len() % 2 == 0);
        (0..s.len())
            .step_by(2)
            .map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
            .collect()
    }

    /// **RFC 3610 §8 Packet Vector #1** -- the canonical AES-CCM
    /// test vector. Pinned byte-exact against:
    ///
    /// ```text
    /// Key:    C0 C1 C2 C3 C4 C5 C6 C7  C8 C9 CA CB CC CD CE CF
    /// Nonce:  00 00 00 03 02 01 00 A0  A1 A2 A3 A4 A5
    /// AAD:    00 01 02 03 04 05 06 07
    /// Plain:  08 09 0A 0B 0C 0D 0E 0F  10 11 12 13 14 15 16 17
    ///         18 19 1A 1B 1C 1D 1E
    /// CCM:    58 8C 97 9A 61 C6 63 D2  F0 66 D0 C2 C0 F9 89 80
    ///         6D 5F 6B 61 DA C3 84
    /// Tag:    17 E8 D1 2C FD F9 26 E0   <-- last 8 bytes "17e8d12cfdf926e0"
    /// ```
    ///
    /// Note: in the RFC layout, the "CCM" output is the concatenation
    /// of (ciphertext || tag). We verify both halves separately.
    #[test]
    fn rfc3610_packet_vector_1() {
        let key = hex("c0c1c2c3c4c5c6c7c8c9cacbcccdcecf");
        let nonce: [u8; 13] = {
            let v = hex("00000003020100a0a1a2a3a4a5");
            v.try_into().unwrap()
        };
        let aad = hex("0001020304050607");
        let plaintext = hex("08090a0b0c0d0e0f101112131415161718191a1b1c1d1e");

        // RFC 3610 packet vector #1 uses M = 8 (8-byte tag).
        let ccm = AesCcm::new(&key, 8).unwrap();
        let (ct, tag) = ccm.encrypt(&nonce, &aad, &plaintext).unwrap();

        let expected_ct = hex("588c979a61c663d2f066d0c2c0f989806d5f6b61dac384");
        let expected_tag = hex("17e8d12cfdf926e0");
        assert_eq!(ct, expected_ct);
        assert_eq!(tag, expected_tag);

        // Round-trip: decrypt with the produced (ct, tag).
        let pt = ccm.decrypt(&nonce, &aad, &ct, &tag).unwrap();
        assert_eq!(pt, plaintext);
    }

    /// **RFC 3610 §8 Packet Vector #2** -- a second pinned vector
    /// for cross-checking. Same key, different nonce + payload.
    #[test]
    fn rfc3610_packet_vector_2() {
        let key = hex("c0c1c2c3c4c5c6c7c8c9cacbcccdcecf");
        let nonce: [u8; 13] = hex("00000004030201a0a1a2a3a4a5").try_into().unwrap();
        let aad = hex("0001020304050607");
        let plaintext = hex("08090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f");

        let ccm = AesCcm::new(&key, 8).unwrap();
        let (ct, tag) = ccm.encrypt(&nonce, &aad, &plaintext).unwrap();

        let expected_ct = hex("72c91a36e135f8cf291ca894085c87e3cc15c439c9e43a3b");
        let expected_tag = hex("a091d56e10400916");
        assert_eq!(ct, expected_ct);
        assert_eq!(tag, expected_tag);

        // Round-trip.
        let pt = ccm.decrypt(&nonce, &aad, &ct, &tag).unwrap();
        assert_eq!(pt, plaintext);
    }

    /// Round-trip on an arbitrary message with M = 16 (TLS / ESP profile).
    #[test]
    fn ccm_aes128_m16_roundtrip() {
        let key = [0x42u8; 16];
        let nonce = [0xa5u8; 13];
        let aad = b"some context";
        let pt = b"hello world; this is a test of moderate length to span more than one AES block.";

        let ccm = AesCcm::new(&key, 16).unwrap();
        let (ct, tag) = ccm.encrypt(&nonce, aad, pt).unwrap();
        assert_eq!(tag.len(), 16);
        assert_ne!(ct.as_slice(), pt.as_slice());

        let back = ccm.decrypt(&nonce, aad, &ct, &tag).unwrap();
        assert_eq!(back.as_slice(), pt.as_slice());
    }

    /// AES-256-CCM round-trip (longer key, same construction).
    #[test]
    fn ccm_aes256_m16_roundtrip() {
        let key = [0x77u8; 32];
        let nonce = [0x11u8; 13];
        let aad = b"";
        let pt = b"AES-256-CCM message";

        let ccm = AesCcm::new(&key, 16).unwrap();
        let (ct, tag) = ccm.encrypt(&nonce, aad, pt).unwrap();
        let back = ccm.decrypt(&nonce, aad, &ct, &tag).unwrap();
        assert_eq!(back.as_slice(), pt.as_slice());
    }

    /// Decrypt rejects a tampered ciphertext byte.
    #[test]
    fn ccm_rejects_tampered_ciphertext() {
        let key = [0x01u8; 16];
        let nonce = [0x02u8; 13];
        let pt = b"do not modify";

        let ccm = AesCcm::new(&key, 8).unwrap();
        let (mut ct, tag) = ccm.encrypt(&nonce, b"", pt).unwrap();
        ct[0] ^= 0x01;
        assert!(ccm.decrypt(&nonce, b"", &ct, &tag).is_none());
    }

    /// Decrypt rejects a tampered tag byte.
    #[test]
    fn ccm_rejects_tampered_tag() {
        let key = [0x01u8; 16];
        let nonce = [0x02u8; 13];
        let pt = b"do not modify";

        let ccm = AesCcm::new(&key, 8).unwrap();
        let (ct, mut tag) = ccm.encrypt(&nonce, b"", pt).unwrap();
        tag[0] ^= 0x01;
        assert!(ccm.decrypt(&nonce, b"", &ct, &tag).is_none());
    }

    /// Decrypt rejects modified AAD (proves AAD is in the MAC input).
    #[test]
    fn ccm_rejects_modified_aad() {
        let key = [0xffu8; 16];
        let nonce = [0x10u8; 13];
        let aad = b"context-A";
        let pt = b"shared payload";

        let ccm = AesCcm::new(&key, 8).unwrap();
        let (ct, tag) = ccm.encrypt(&nonce, aad, pt).unwrap();
        assert!(ccm.decrypt(&nonce, b"context-B", &ct, &tag).is_none());
    }

    /// Decrypt rejects when the wrong key is used.
    #[test]
    fn ccm_rejects_wrong_key() {
        let mut key1 = [0x33u8; 16];
        let mut key2 = key1;
        key2[0] ^= 0x01;
        let nonce = [0x44u8; 13];
        let pt = b"sensitive";

        let ccm1 = AesCcm::new(&key1, 8).unwrap();
        let (ct, tag) = ccm1.encrypt(&nonce, b"", pt).unwrap();
        let ccm2 = AesCcm::new(&key2, 8).unwrap();
        assert!(ccm2.decrypt(&nonce, b"", &ct, &tag).is_none());
        // Touch key1 so the compiler doesn't warn that it's only read once.
        key1[0] = 0;
    }

    /// Empty plaintext is allowed: ct.len() == 0 but the tag still
    /// authenticates the AAD and the length fields.
    #[test]
    fn ccm_empty_plaintext() {
        let key = [0x55u8; 16];
        let nonce = [0x66u8; 13];
        let aad = b"only-context";

        let ccm = AesCcm::new(&key, 8).unwrap();
        let (ct, tag) = ccm.encrypt(&nonce, aad, b"").unwrap();
        assert!(ct.is_empty());
        let back = ccm.decrypt(&nonce, aad, &ct, &tag).unwrap();
        assert!(back.is_empty());

        // Modifying the AAD must still be detected on empty plaintext.
        assert!(ccm.decrypt(&nonce, b"other", &ct, &tag).is_none());
    }

    /// Parameter validation: invalid M values are rejected.
    #[test]
    fn ccm_rejects_invalid_m() {
        let key = [0u8; 16];
        for bad_m in [0, 1, 2, 3, 5, 7, 9, 11, 13, 15, 17, 18, 32] {
            assert!(AesCcm::new(&key, bad_m).is_none(), "M={} should be rejected", bad_m);
        }
        for good_m in [4, 6, 8, 10, 12, 14, 16] {
            assert!(AesCcm::new(&key, good_m).is_some(), "M={} should be accepted", good_m);
        }
    }
}