krypteia-arcana 0.1.0

//! AES (FIPS 197) block cipher — single-block API.
//!
//! Supports 128-bit, 192-bit, and 256-bit keys with 10, 12, and 14
//! rounds respectively. Modes (ECB / CBC / CTR / GCM / CCM / XTS)
//! live in [`super::modes`], [`super::ccm`], [`super::xts`]; the
//! streaming wrapper lives in [`super::ctx`].
//!
//! # ⚠ Side-channel posture (evaluation-critical gap)
//!
//! This module is the **single largest open SCA gap on the
//! classical side**. Roadmap entries
//! (`arcana/doc/sca/countermeasures/aes.rst`):
//!
//! | Threat                                 | Status     | Roadmap item                                                                  |
//! |----------------------------------------|------------|-------------------------------------------------------------------------------|
//! | SPA / SEMA on key schedule + S-box     | vulnerable | `T1-A` — port fixsliced AES (Adomnicai-Peyrin TCHES 2021/1)                   |
//! | **Cache-timing on shared L1 / L2**     | vulnerable | Same `T1-A`. AES-NI / VAES is host-only (`T5`)                                |
//! | DPA / CPA on round-1 SubBytes          | vulnerable | `T2-G` — first-order Boolean masking on top of fixsliced AES                  |
//! | Template attacks (incl. ML-DPA)        | vulnerable | `T2-G`. ANSSI's protected AES was broken end-to-end by ML-DPA in 2023         |
//! | DFA on last AES round                  | vulnerable | `T4-AES-A` — redundancy + infective countermeasure (deferred)                 |
//!
//! ## Cache-timing leak — concrete model
//!
//! The `SBOX` array below is a 256-byte LUT. The first round of
//! AES indexes 16 bytes of `state[i] ^ K[i]`; observing which
//! cache lines (4 lines × 64 B = 256 B) are accessed reveals the
//! high bits of each byte of `state[i] ^ K[i]`. Combined T-table
//! implementations (which fold ShiftRows + MixColumns into 4 KiB
//! of pre-computed tables) leak more. References:
//! `bernstein2005_aes_cache_timing`, `osvik2006cache_aes`.
//!
//! **Until `T1-A` lands, this implementation must not be used in
//! deployments where a co-resident or near-shared-cache attacker
//! is in scope** — shared hosting, multi-VM tenants, or any
//! bare-metal target with shared L1 between cryptographic and
//! untrusted code.

use crate::BlockCipher;

// ============================================================
// S-box and inverse S-box (FIPS 197, Section 5.1.1)
// ============================================================

/// AES forward S-box.
const SBOX: [u8; 256] = [
    0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5, 0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76, 0xca, 0x82, 0xc9,
    0x7d, 0xfa, 0x59, 0x47, 0xf0, 0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0, 0xb7, 0xfd, 0x93, 0x26, 0x36, 0x3f,
    0xf7, 0xcc, 0x34, 0xa5, 0xe5, 0xf1, 0x71, 0xd8, 0x31, 0x15, 0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05, 0x9a, 0x07,
    0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75, 0x09, 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0, 0x52, 0x3b, 0xd6, 0xb3,
    0x29, 0xe3, 0x2f, 0x84, 0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b, 0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58,
    0xcf, 0xd0, 0xef, 0xaa, 0xfb, 0x43, 0x4d, 0x33, 0x85, 0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c, 0x9f, 0xa8, 0x51, 0xa3,
    0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5, 0xbc, 0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2, 0xcd, 0x0c, 0x13, 0xec, 0x5f,
    0x97, 0x44, 0x17, 0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73, 0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88,
    0x46, 0xee, 0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb, 0xe0, 0x32, 0x3a, 0x0a, 0x49, 0x06, 0x24, 0x5c, 0xc2, 0xd3, 0xac,
    0x62, 0x91, 0x95, 0xe4, 0x79, 0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9, 0x6c, 0x56, 0xf4, 0xea, 0x65, 0x7a,
    0xae, 0x08, 0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6, 0xb4, 0xc6, 0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a, 0x70,
    0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e, 0x61, 0x35, 0x57, 0xb9, 0x86, 0xc1, 0x1d, 0x9e, 0xe1, 0xf8, 0x98, 0x11,
    0x69, 0xd9, 0x8e, 0x94, 0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf, 0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42,
    0x68, 0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16,
];

/// AES inverse S-box.
const INV_SBOX: [u8; 256] = [
    0x52, 0x09, 0x6a, 0xd5, 0x30, 0x36, 0xa5, 0x38, 0xbf, 0x40, 0xa3, 0x9e, 0x81, 0xf3, 0xd7, 0xfb, 0x7c, 0xe3, 0x39,
    0x82, 0x9b, 0x2f, 0xff, 0x87, 0x34, 0x8e, 0x43, 0x44, 0xc4, 0xde, 0xe9, 0xcb, 0x54, 0x7b, 0x94, 0x32, 0xa6, 0xc2,
    0x23, 0x3d, 0xee, 0x4c, 0x95, 0x0b, 0x42, 0xfa, 0xc3, 0x4e, 0x08, 0x2e, 0xa1, 0x66, 0x28, 0xd9, 0x24, 0xb2, 0x76,
    0x5b, 0xa2, 0x49, 0x6d, 0x8b, 0xd1, 0x25, 0x72, 0xf8, 0xf6, 0x64, 0x86, 0x68, 0x98, 0x16, 0xd4, 0xa4, 0x5c, 0xcc,
    0x5d, 0x65, 0xb6, 0x92, 0x6c, 0x70, 0x48, 0x50, 0xfd, 0xed, 0xb9, 0xda, 0x5e, 0x15, 0x46, 0x57, 0xa7, 0x8d, 0x9d,
    0x84, 0x90, 0xd8, 0xab, 0x00, 0x8c, 0xbc, 0xd3, 0x0a, 0xf7, 0xe4, 0x58, 0x05, 0xb8, 0xb3, 0x45, 0x06, 0xd0, 0x2c,
    0x1e, 0x8f, 0xca, 0x3f, 0x0f, 0x02, 0xc1, 0xaf, 0xbd, 0x03, 0x01, 0x13, 0x8a, 0x6b, 0x3a, 0x91, 0x11, 0x41, 0x4f,
    0x67, 0xdc, 0xea, 0x97, 0xf2, 0xcf, 0xce, 0xf0, 0xb4, 0xe6, 0x73, 0x96, 0xac, 0x74, 0x22, 0xe7, 0xad, 0x35, 0x85,
    0xe2, 0xf9, 0x37, 0xe8, 0x1c, 0x75, 0xdf, 0x6e, 0x47, 0xf1, 0x1a, 0x71, 0x1d, 0x29, 0xc5, 0x89, 0x6f, 0xb7, 0x62,
    0x0e, 0xaa, 0x18, 0xbe, 0x1b, 0xfc, 0x56, 0x3e, 0x4b, 0xc6, 0xd2, 0x79, 0x20, 0x9a, 0xdb, 0xc0, 0xfe, 0x78, 0xcd,
    0x5a, 0xf4, 0x1f, 0xdd, 0xa8, 0x33, 0x88, 0x07, 0xc7, 0x31, 0xb1, 0x12, 0x10, 0x59, 0x27, 0x80, 0xec, 0x5f, 0x60,
    0x51, 0x7f, 0xa9, 0x19, 0xb5, 0x4a, 0x0d, 0x2d, 0xe5, 0x7a, 0x9f, 0x93, 0xc9, 0x9c, 0xef, 0xa0, 0xe0, 0x3b, 0x4d,
    0xae, 0x2a, 0xf5, 0xb0, 0xc8, 0xeb, 0xbb, 0x3c, 0x83, 0x53, 0x99, 0x61, 0x17, 0x2b, 0x04, 0x7e, 0xba, 0x77, 0xd6,
    0x26, 0xe1, 0x69, 0x14, 0x63, 0x55, 0x21, 0x0c, 0x7d,
];

/// Round constants for key expansion.
const RCON: [u8; 10] = [0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36];

// ============================================================
// GF(2^8) helpers
// ============================================================

/// Multiply by 2 in GF(2^8) with the AES irreducible polynomial x^8+x^4+x^3+x+1.
#[inline]
fn xtime(a: u8) -> u8 {
    let shifted = (a as u16) << 1;
    (shifted ^ (if a & 0x80 != 0 { 0x1b } else { 0x00 })) as u8
}

/// Multiply two elements in GF(2^8).
#[inline]
fn gmul(mut a: u8, mut b: u8) -> u8 {
    let mut p: u8 = 0;
    for _ in 0..8 {
        if b & 1 != 0 {
            p ^= a;
        }
        let hi = a & 0x80;
        a <<= 1;
        if hi != 0 {
            a ^= 0x1b;
        }
        b >>= 1;
    }
    p
}

// ============================================================
// Aes core
// ============================================================

/// Maximum number of round keys: AES-256 has 14 rounds → 15 round keys × 4 words = 60 words.
const MAX_ROUND_KEYS: usize = 60;

/// AES block cipher supporting 128, 192, and 256-bit keys.
///
/// See module-level documentation for cache-timing caveats.
pub struct Aes {
    /// Expanded round key words (Nk * (Nr+1) 32-bit words).
    round_keys: [u32; MAX_ROUND_KEYS],
    /// Number of rounds (10, 12, or 14).
    nr: usize,
}

impl Aes {
    /// Number of rounds for a given key length in bytes.
    fn rounds_for_key(key_len: usize) -> usize {
        match key_len {
            16 => 10,
            24 => 12,
            32 => 14,
            _ => panic!("AES: invalid key length (must be 16, 24, or 32 bytes)"),
        }
    }

    /// Key expansion (FIPS 197, Section 5.2).
    fn key_expansion(key: &[u8]) -> ([u32; MAX_ROUND_KEYS], usize) {
        let nk = key.len() / 4; // number of 32-bit words in the key
        let nr = Self::rounds_for_key(key.len());
        let total_words = 4 * (nr + 1);

        let mut w = [0u32; MAX_ROUND_KEYS];

        // Copy the key into the first Nk words
        for i in 0..nk {
            w[i] = u32::from_be_bytes([key[4 * i], key[4 * i + 1], key[4 * i + 2], key[4 * i + 3]]);
        }

        for i in nk..total_words {
            let mut temp = w[i - 1];
            if i % nk == 0 {
                // RotWord + SubWord + Rcon
                temp = Self::sub_word(Self::rot_word(temp)) ^ ((RCON[i / nk - 1] as u32) << 24);
            } else if nk > 6 && i % nk == 4 {
                temp = Self::sub_word(temp);
            }
            w[i] = w[i - nk] ^ temp;
        }

        (w, nr)
    }

    #[inline]
    fn sub_word(w: u32) -> u32 {
        let b = w.to_be_bytes();
        u32::from_be_bytes([
            SBOX[b[0] as usize],
            SBOX[b[1] as usize],
            SBOX[b[2] as usize],
            SBOX[b[3] as usize],
        ])
    }

    #[inline]
    fn rot_word(w: u32) -> u32 {
        w.rotate_left(8)
    }

    // ---- Encryption transforms ----

    fn sub_bytes(state: &mut [u8; 16]) {
        for b in state.iter_mut() {
            *b = SBOX[*b as usize];
        }
    }

    fn shift_rows(state: &mut [u8; 16]) {
        // Row 1: shift left by 1
        let tmp = state[1];
        state[1] = state[5];
        state[5] = state[9];
        state[9] = state[13];
        state[13] = tmp;

        // Row 2: shift left by 2
        let tmp0 = state[2];
        let tmp1 = state[6];
        state[2] = state[10];
        state[6] = state[14];
        state[10] = tmp0;
        state[14] = tmp1;

        // Row 3: shift left by 3 (= shift right by 1)
        let tmp = state[15];
        state[15] = state[11];
        state[11] = state[7];
        state[7] = state[3];
        state[3] = tmp;
    }

    fn mix_columns(state: &mut [u8; 16]) {
        for c in 0..4 {
            let i = 4 * c;
            let s0 = state[i];
            let s1 = state[i + 1];
            let s2 = state[i + 2];
            let s3 = state[i + 3];

            state[i] = xtime(s0) ^ xtime(s1) ^ s1 ^ s2 ^ s3;
            state[i + 1] = s0 ^ xtime(s1) ^ xtime(s2) ^ s2 ^ s3;
            state[i + 2] = s0 ^ s1 ^ xtime(s2) ^ xtime(s3) ^ s3;
            state[i + 3] = xtime(s0) ^ s0 ^ s1 ^ s2 ^ xtime(s3);
        }
    }

    fn add_round_key(state: &mut [u8; 16], rk: &[u32]) {
        for c in 0..4 {
            let k = rk[c].to_be_bytes();
            state[4 * c] ^= k[0];
            state[4 * c + 1] ^= k[1];
            state[4 * c + 2] ^= k[2];
            state[4 * c + 3] ^= k[3];
        }
    }

    // ---- Decryption transforms ----

    fn inv_sub_bytes(state: &mut [u8; 16]) {
        for b in state.iter_mut() {
            *b = INV_SBOX[*b as usize];
        }
    }

    fn inv_shift_rows(state: &mut [u8; 16]) {
        // Row 1: shift right by 1
        let tmp = state[13];
        state[13] = state[9];
        state[9] = state[5];
        state[5] = state[1];
        state[1] = tmp;

        // Row 2: shift right by 2
        let tmp0 = state[2];
        let tmp1 = state[6];
        state[2] = state[10];
        state[6] = state[14];
        state[10] = tmp0;
        state[14] = tmp1;

        // Row 3: shift right by 3 (= shift left by 1)
        let tmp = state[3];
        state[3] = state[7];
        state[7] = state[11];
        state[11] = state[15];
        state[15] = tmp;
    }

    fn inv_mix_columns(state: &mut [u8; 16]) {
        for c in 0..4 {
            let i = 4 * c;
            let s0 = state[i];
            let s1 = state[i + 1];
            let s2 = state[i + 2];
            let s3 = state[i + 3];

            state[i] = gmul(s0, 0x0e) ^ gmul(s1, 0x0b) ^ gmul(s2, 0x0d) ^ gmul(s3, 0x09);
            state[i + 1] = gmul(s0, 0x09) ^ gmul(s1, 0x0e) ^ gmul(s2, 0x0b) ^ gmul(s3, 0x0d);
            state[i + 2] = gmul(s0, 0x0d) ^ gmul(s1, 0x09) ^ gmul(s2, 0x0e) ^ gmul(s3, 0x0b);
            state[i + 3] = gmul(s0, 0x0b) ^ gmul(s1, 0x0d) ^ gmul(s2, 0x09) ^ gmul(s3, 0x0e);
        }
    }
}

impl BlockCipher for Aes {
    const BLOCK_LEN: usize = 16;
    const KEY_LENS: &'static [usize] = &[16, 24, 32];

    fn new(key: &[u8]) -> Self {
        let (round_keys, nr) = Self::key_expansion(key);
        Aes { round_keys, nr }
    }

    fn encrypt_block(&self, block: &mut [u8]) {
        assert!(block.len() >= 16, "AES: block must be at least 16 bytes");
        let mut state = [0u8; 16];
        state.copy_from_slice(&block[..16]);

        // Initial round key addition
        Self::add_round_key(&mut state, &self.round_keys[0..4]);

        // Rounds 1..Nr-1
        for round in 1..self.nr {
            Self::sub_bytes(&mut state);
            Self::shift_rows(&mut state);
            Self::mix_columns(&mut state);
            Self::add_round_key(&mut state, &self.round_keys[round * 4..(round + 1) * 4]);
        }

        // Final round (no MixColumns)
        Self::sub_bytes(&mut state);
        Self::shift_rows(&mut state);
        Self::add_round_key(&mut state, &self.round_keys[self.nr * 4..(self.nr + 1) * 4]);

        block[..16].copy_from_slice(&state);
    }

    fn decrypt_block(&self, block: &mut [u8]) {
        assert!(block.len() >= 16, "AES: block must be at least 16 bytes");
        let mut state = [0u8; 16];
        state.copy_from_slice(&block[..16]);

        // Initial round key addition (last round key)
        Self::add_round_key(&mut state, &self.round_keys[self.nr * 4..(self.nr + 1) * 4]);

        // Rounds Nr-1..1
        for round in (1..self.nr).rev() {
            Self::inv_shift_rows(&mut state);
            Self::inv_sub_bytes(&mut state);
            Self::add_round_key(&mut state, &self.round_keys[round * 4..(round + 1) * 4]);
            Self::inv_mix_columns(&mut state);
        }

        // Final round (no InvMixColumns)
        Self::inv_shift_rows(&mut state);
        Self::inv_sub_bytes(&mut state);
        Self::add_round_key(&mut state, &self.round_keys[0..4]);

        block[..16].copy_from_slice(&state);
    }
}

// ============================================================
// Convenience wrappers
// ============================================================

/// AES-128 (10 rounds, 128-bit key).
pub struct Aes128 {
    inner: Aes,
}

impl BlockCipher for Aes128 {
    const BLOCK_LEN: usize = 16;
    const KEY_LENS: &'static [usize] = &[16];

    fn new(key: &[u8]) -> Self {
        assert_eq!(key.len(), 16, "AES-128 requires a 16-byte key");
        Aes128 { inner: Aes::new(key) }
    }

    fn encrypt_block(&self, block: &mut [u8]) {
        self.inner.encrypt_block(block);
    }

    fn decrypt_block(&self, block: &mut [u8]) {
        self.inner.decrypt_block(block);
    }
}

/// AES-192 (12 rounds, 192-bit key).
pub struct Aes192 {
    inner: Aes,
}

impl BlockCipher for Aes192 {
    const BLOCK_LEN: usize = 16;
    const KEY_LENS: &'static [usize] = &[24];

    fn new(key: &[u8]) -> Self {
        assert_eq!(key.len(), 24, "AES-192 requires a 24-byte key");
        Aes192 { inner: Aes::new(key) }
    }

    fn encrypt_block(&self, block: &mut [u8]) {
        self.inner.encrypt_block(block);
    }

    fn decrypt_block(&self, block: &mut [u8]) {
        self.inner.decrypt_block(block);
    }
}

/// AES-256 (14 rounds, 256-bit key).
pub struct Aes256 {
    inner: Aes,
}

impl BlockCipher for Aes256 {
    const BLOCK_LEN: usize = 16;
    const KEY_LENS: &'static [usize] = &[32];

    fn new(key: &[u8]) -> Self {
        assert_eq!(key.len(), 32, "AES-256 requires a 32-byte key");
        Aes256 { inner: Aes::new(key) }
    }

    fn encrypt_block(&self, block: &mut [u8]) {
        self.inner.encrypt_block(block);
    }

    fn decrypt_block(&self, block: &mut [u8]) {
        self.inner.decrypt_block(block);
    }
}

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

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

    fn hex_to_bytes(s: &str) -> Vec<u8> {
        (0..s.len())
            .step_by(2)
            .map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
            .collect()
    }

    /// FIPS 197 Appendix B test vector.
    #[test]
    fn aes128_fips197_appendix_b() {
        let key = hex_to_bytes("2b7e151628aed2a6abf7158809cf4f3c");
        let plaintext = hex_to_bytes("3243f6a8885a308d313198a2e0370734");
        let expected_ct = hex_to_bytes("3925841d02dc09fbdc118597196a0b32");

        let cipher = Aes128::new(&key);

        // Encrypt
        let mut block = plaintext.clone();
        cipher.encrypt_block(&mut block);
        assert_eq!(block, expected_ct, "AES-128 encrypt mismatch");

        // Decrypt
        cipher.decrypt_block(&mut block);
        assert_eq!(block, plaintext, "AES-128 decrypt mismatch");
    }

    /// NIST FIPS 197 Appendix C.1 — AES-128
    #[test]
    fn aes128_nist_c1() {
        let key = hex_to_bytes("000102030405060708090a0b0c0d0e0f");
        let pt = hex_to_bytes("00112233445566778899aabbccddeeff");
        let expected = hex_to_bytes("69c4e0d86a7b0430d8cdb78070b4c55a");

        let cipher = Aes128::new(&key);
        let mut block = pt.clone();
        cipher.encrypt_block(&mut block);
        assert_eq!(block, expected);

        cipher.decrypt_block(&mut block);
        assert_eq!(block, pt);
    }

    /// NIST FIPS 197 Appendix C.2 — AES-192
    #[test]
    fn aes192_nist_c2() {
        let key = hex_to_bytes("000102030405060708090a0b0c0d0e0f1011121314151617");
        let pt = hex_to_bytes("00112233445566778899aabbccddeeff");
        let expected = hex_to_bytes("dda97ca4864cdfe06eaf70a0ec0d7191");

        let cipher = Aes192::new(&key);
        let mut block = pt.clone();
        cipher.encrypt_block(&mut block);
        assert_eq!(block, expected);

        cipher.decrypt_block(&mut block);
        assert_eq!(block, pt);
    }

    /// NIST FIPS 197 Appendix C.3 — AES-256
    #[test]
    fn aes256_nist_c3() {
        let key = hex_to_bytes("000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f");
        let pt = hex_to_bytes("00112233445566778899aabbccddeeff");
        let expected = hex_to_bytes("8ea2b7ca516745bfeafc49904b496089");

        let cipher = Aes256::new(&key);
        let mut block = pt.clone();
        cipher.encrypt_block(&mut block);
        assert_eq!(block, expected);

        cipher.decrypt_block(&mut block);
        assert_eq!(block, pt);
    }
}