Skip to main content

crypto/aes/
aes.rs

1/// Pure-Rust AES block cipher (128-bit and 256-bit keys).
2///
3/// Design notes:
4/// - Favours speed over constant-time execution (table-driven S-box lookups).
5/// - Uses T-tables (Te0..Te3) combining SubBytes+ShiftRows+MixColumns into a
6///   single 32-bit lookup per byte per round, matching Go's software AES approach.
7#[cfg(test)]
8use super::ghash::gf128_mul;
9
10// ── AES constants ─────────────────────────────────────────────────────────────
11
12pub const GCM_MAX_LEN: u64 = (u32::MAX as u64 - 1) * 16;
13
14#[rustfmt::skip]
15pub(crate) const SBOX: [u8; 256] = [
16    0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5, 0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76,
17    0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0, 0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0,
18    0xb7, 0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc, 0x34, 0xa5, 0xe5, 0xf1, 0x71, 0xd8, 0x31, 0x15,
19    0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05, 0x9a, 0x07, 0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75,
20    0x09, 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0, 0x52, 0x3b, 0xd6, 0xb3, 0x29, 0xe3, 0x2f, 0x84,
21    0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b, 0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf,
22    0xd0, 0xef, 0xaa, 0xfb, 0x43, 0x4d, 0x33, 0x85, 0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c, 0x9f, 0xa8,
23    0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5, 0xbc, 0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2,
24    0xcd, 0x0c, 0x13, 0xec, 0x5f, 0x97, 0x44, 0x17, 0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73,
25    0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88, 0x46, 0xee, 0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb,
26    0xe0, 0x32, 0x3a, 0x0a, 0x49, 0x06, 0x24, 0x5c, 0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79,
27    0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9, 0x6c, 0x56, 0xf4, 0xea, 0x65, 0x7a, 0xae, 0x08,
28    0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6, 0xb4, 0xc6, 0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a,
29    0x70, 0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e, 0x61, 0x35, 0x57, 0xb9, 0x86, 0xc1, 0x1d, 0x9e,
30    0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e, 0x94, 0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf,
31    0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68, 0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16,
32];
33
34#[rustfmt::skip]
35const SBOX_INV: [u8; 256] = [
36    0x52, 0x09, 0x6a, 0xd5, 0x30, 0x36, 0xa5, 0x38, 0xbf, 0x40, 0xa3, 0x9e, 0x81, 0xf3, 0xd7, 0xfb,
37    0x7c, 0xe3, 0x39, 0x82, 0x9b, 0x2f, 0xff, 0x87, 0x34, 0x8e, 0x43, 0x44, 0xc4, 0xde, 0xe9, 0xcb,
38    0x54, 0x7b, 0x94, 0x32, 0xa6, 0xc2, 0x23, 0x3d, 0xee, 0x4c, 0x95, 0x0b, 0x42, 0xfa, 0xc3, 0x4e,
39    0x08, 0x2e, 0xa1, 0x66, 0x28, 0xd9, 0x24, 0xb2, 0x76, 0x5b, 0xa2, 0x49, 0x6d, 0x8b, 0xd1, 0x25,
40    0x72, 0xf8, 0xf6, 0x64, 0x86, 0x68, 0x98, 0x16, 0xd4, 0xa4, 0x5c, 0xcc, 0x5d, 0x65, 0xb6, 0x92,
41    0x6c, 0x70, 0x48, 0x50, 0xfd, 0xed, 0xb9, 0xda, 0x5e, 0x15, 0x46, 0x57, 0xa7, 0x8d, 0x9d, 0x84,
42    0x90, 0xd8, 0xab, 0x00, 0x8c, 0xbc, 0xd3, 0x0a, 0xf7, 0xe4, 0x58, 0x05, 0xb8, 0xb3, 0x45, 0x06,
43    0xd0, 0x2c, 0x1e, 0x8f, 0xca, 0x3f, 0x0f, 0x02, 0xc1, 0xaf, 0xbd, 0x03, 0x01, 0x13, 0x8a, 0x6b,
44    0x3a, 0x91, 0x11, 0x41, 0x4f, 0x67, 0xdc, 0xea, 0x97, 0xf2, 0xcf, 0xce, 0xf0, 0xb4, 0xe6, 0x73,
45    0x96, 0xac, 0x74, 0x22, 0xe7, 0xad, 0x35, 0x85, 0xe2, 0xf9, 0x37, 0xe8, 0x1c, 0x75, 0xdf, 0x6e,
46    0x47, 0xf1, 0x1a, 0x71, 0x1d, 0x29, 0xc5, 0x89, 0x6f, 0xb7, 0x62, 0x0e, 0xaa, 0x18, 0xbe, 0x1b,
47    0xfc, 0x56, 0x3e, 0x4b, 0xc6, 0xd2, 0x79, 0x20, 0x9a, 0xdb, 0xc0, 0xfe, 0x78, 0xcd, 0x5a, 0xf4,
48    0x1f, 0xdd, 0xa8, 0x33, 0x88, 0x07, 0xc7, 0x31, 0xb1, 0x12, 0x10, 0x59, 0x27, 0x80, 0xec, 0x5f,
49    0x60, 0x51, 0x7f, 0xa9, 0x19, 0xb5, 0x4a, 0x0d, 0x2d, 0xe5, 0x7a, 0x9f, 0x93, 0xc9, 0x9c, 0xef,
50    0xa0, 0xe0, 0x3b, 0x4d, 0xae, 0x2a, 0xf5, 0xb0, 0xc8, 0xeb, 0xbb, 0x3c, 0x83, 0x53, 0x99, 0x61,
51    0x17, 0x2b, 0x04, 0x7e, 0xba, 0x77, 0xd6, 0x26, 0xe1, 0x69, 0x14, 0x63, 0x55, 0x21, 0x0c, 0x7d,
52];
53
54/// Round constants for AES key expansion (RCON[1..10]).
55pub(crate) const RCON: [u8; 11] = [0x00, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36];
56
57// ── AES key schedule ───────────────────────────────────────────────────────────
58
59/// AES expanded key: `N` round keys x 16 bytes.
60/// N = 11 for AES-128 and 15 for AES-256
61#[derive(Clone)]
62#[cfg_attr(feature = "zeroize", derive(zeroize::Zeroize, zeroize::ZeroizeOnDrop))]
63pub(crate) enum RoundKeys<const N: usize> {
64    #[cfg(target_arch = "x86_64")]
65    /// The platform supports x86 AES-NI instructions
66    X86_64([core::arch::x86_64::__m128i; N]),
67
68    #[cfg(target_arch = "aarch64")]
69    /// The platform supports ARMv8 AES instructions
70    Armv8([core::arch::aarch64::uint8x16_t; N]),
71
72    /// The platform doesn't support AES hardware acceleration
73    Software([[u8; 16]; N]),
74}
75
76pub(crate) type RoundKeysSoftware<const N: usize> = [[u8; 16]; N];
77
78#[inline(always)]
79fn rot_word(w: [u8; 4]) -> [u8; 4] {
80    [w[1], w[2], w[3], w[0]]
81}
82
83#[inline(always)]
84fn sub_word(w: [u8; 4]) -> [u8; 4] {
85    [
86        SBOX[w[0] as usize],
87        SBOX[w[1] as usize],
88        SBOX[w[2] as usize],
89        SBOX[w[3] as usize],
90    ]
91}
92
93/// Expand a key into AES round keys (FIPS 197 §5.2).
94///
95/// - `N = 11` -> AES-128 (Nk=4, Nr=10, 44 words -> 11 round keys)
96/// - `N = 15` -> AES-256 (Nk=8, Nr=14, 60 words -> 15 round keys)
97pub fn expand_key<const N: usize>(key: &[u8]) -> RoundKeysSoftware<N> {
98    const {
99        assert!(N == 11 || N == 15);
100    }
101    let nk = key.len() / 4;
102    let words = N * 4;
103
104    // Max 60 words (AES-256). Use a fixed-size buffer for both key sizes.
105    let mut w = [[0u8; 4]; 60];
106    for i in 0..nk {
107        w[i] = [key[4 * i], key[4 * i + 1], key[4 * i + 2], key[4 * i + 3]];
108    }
109
110    for i in nk..words {
111        let mut temp = w[i - 1];
112        if i % nk == 0 {
113            temp = sub_word(rot_word(temp));
114            temp[0] ^= RCON[i / nk];
115        } else if N == 15 && i % nk == 4 {
116            // AES-256 only: SubWord every 4 words within the Nk=8 group
117            temp = sub_word(temp);
118        }
119        w[i] = [
120            w[i - nk][0] ^ temp[0],
121            w[i - nk][1] ^ temp[1],
122            w[i - nk][2] ^ temp[2],
123            w[i - nk][3] ^ temp[3],
124        ];
125    }
126
127    let mut rk = [[0u8; 16]; N];
128    for i in 0..N {
129        for j in 0..4 {
130            rk[i][4 * j..4 * j + 4].copy_from_slice(&w[4 * i + j]);
131        }
132    }
133    rk
134}
135
136// ── AES block cipher (encrypt / decrypt) ─────────────────────────────────────
137
138/// Const-time gf128 xtime
139#[inline(always)]
140const fn xtime(a: u8) -> u8 {
141    let hi = a & 0x80;
142    let b = a << 1;
143    if hi != 0 { b ^ 0x1b } else { b }
144}
145
146// ── T-tables: combine SubBytes + ShiftRows + MixColumns ─────────────────────
147//
148// Each Teᵢ[x] = S[x] multiplied by a column of the MixColumns matrix
149// (rotated left by i positions). A u32 holds row 0..row 3 in LE byte order.
150
151pub(crate) const TE0: [u32; 256] = {
152    let mut t = [0u32; 256];
153    let mut i = 0usize;
154    while i < 256 {
155        let s = SBOX[i] as u32;
156        let s2 = xtime(SBOX[i]) as u32;
157        let s3 = (s ^ s2) as u32;
158        // Column: [2·S, 1·S, 1·S, 3·S] in rows 0..3
159        t[i] = (s3 << 24) | (s << 16) | (s << 8) | s2;
160        i += 1;
161    }
162    t
163};
164
165pub(crate) const TE1: [u32; 256] = {
166    let mut t = [0u32; 256];
167    let mut i = 0usize;
168    while i < 256 {
169        let s = SBOX[i] as u32;
170        let s2 = xtime(SBOX[i]) as u32;
171        let s3 = (s ^ s2) as u32;
172        // Column: [3·S, 2·S, 1·S, 1·S]  -> LE bytes [3S,2S,1S,1S]
173        t[i] = (s << 24) | (s << 16) | (s2 << 8) | s3;
174        i += 1;
175    }
176    t
177};
178
179pub(crate) const TE2: [u32; 256] = {
180    let mut t = [0u32; 256];
181    let mut i = 0usize;
182    while i < 256 {
183        let s = SBOX[i] as u32;
184        let s2 = xtime(SBOX[i]) as u32;
185        let s3 = (s ^ s2) as u32;
186        // Column: [1·S, 3·S, 2·S, 1·S]
187        t[i] = (s << 24) | (s2 << 16) | (s3 << 8) | s;
188        i += 1;
189    }
190    t
191};
192
193pub(crate) const TE3: [u32; 256] = {
194    let mut t = [0u32; 256];
195    let mut i = 0usize;
196    while i < 256 {
197        let s = SBOX[i] as u32;
198        let s2 = xtime(SBOX[i]) as u32;
199        let s3 = (s ^ s2) as u32;
200        // Column: [1·S, 1·S, 3·S, 2·S]
201        t[i] = (s2 << 24) | (s3 << 16) | (s << 8) | s;
202        i += 1;
203    }
204    t
205};
206
207// Inverse T-tables for decryption (InvSubBytes + InvShiftRows + InvMixColumns).
208
209const TD0: [u32; 256] = {
210    let mut t = [0u32; 256];
211    let mut i = 0usize;
212    while i < 256 {
213        let s = SBOX_INV[i] as u32;
214        let s2 = xtime(s as u8) as u32;
215        let s4 = xtime(s2 as u8) as u32;
216        let s8 = xtime(s4 as u8) as u32;
217        let c0 = (s8 ^ s4 ^ s2) as u32; // 0e·S
218        let c1 = (s8 ^ s2 ^ s) as u32; // 0b·S
219        let c2 = (s8 ^ s4 ^ s) as u32; // 0d·S
220        let c3 = (s8 ^ s) as u32; // 09·S
221        // Column: [0e·S, 09·S, 0d·S, 0b·S]  -> LE bytes [c0,c3,c2,c1]
222        t[i] = (c1 << 24) | (c2 << 16) | (c3 << 8) | c0;
223        i += 1;
224    }
225    t
226};
227
228const TD1: [u32; 256] = {
229    let mut t = [0u32; 256];
230    let mut i = 0usize;
231    while i < 256 {
232        let s = SBOX_INV[i] as u32;
233        let s2 = xtime(s as u8) as u32;
234        let s4 = xtime(s2 as u8) as u32;
235        let s8 = xtime(s4 as u8) as u32;
236        let c0 = (s8 ^ s4 ^ s2) as u32;
237        let c1 = (s8 ^ s2 ^ s) as u32;
238        let c2 = (s8 ^ s4 ^ s) as u32;
239        let c3 = (s8 ^ s) as u32;
240        // Column: [0b·S, 0e·S, 09·S, 0d·S]
241        t[i] = (c2 << 24) | (c3 << 16) | (c0 << 8) | c1;
242        i += 1;
243    }
244    t
245};
246
247const TD2: [u32; 256] = {
248    let mut t = [0u32; 256];
249    let mut i = 0usize;
250    while i < 256 {
251        let s = SBOX_INV[i] as u32;
252        let s2 = xtime(s as u8) as u32;
253        let s4 = xtime(s2 as u8) as u32;
254        let s8 = xtime(s4 as u8) as u32;
255        let c0 = (s8 ^ s4 ^ s2) as u32;
256        let c1 = (s8 ^ s2 ^ s) as u32;
257        let c2 = (s8 ^ s4 ^ s) as u32;
258        let c3 = (s8 ^ s) as u32;
259        // Column: [0d·S, 0b·S, 0e·S, 09·S]
260        t[i] = (c3 << 24) | (c0 << 16) | (c1 << 8) | c2;
261        i += 1;
262    }
263    t
264};
265
266const TD3: [u32; 256] = {
267    let mut t = [0u32; 256];
268    let mut i = 0usize;
269    while i < 256 {
270        let s = SBOX_INV[i] as u32;
271        let s2 = xtime(s as u8) as u32;
272        let s4 = xtime(s2 as u8) as u32;
273        let s8 = xtime(s4 as u8) as u32;
274        let c0 = (s8 ^ s4 ^ s2) as u32;
275        let c1 = (s8 ^ s2 ^ s) as u32;
276        let c2 = (s8 ^ s4 ^ s) as u32;
277        let c3 = (s8 ^ s) as u32;
278        // Column: [09·S, 0d·S, 0b·S, 0e·S]
279        t[i] = (c0 << 24) | (c1 << 16) | (c2 << 8) | c3;
280        i += 1;
281    }
282    t
283};
284
285/// Encrypt one 16-byte block using T-table accelerated routine.
286///
287/// Combines SubBytes, ShiftRows, and MixColumns into four 32-bit table lookups
288/// per column, matching the approach used by Go and OpenSSL for software AES.
289///
290/// `N = 11` for AES-128 (10 rounds), `N = 15` for AES-256 (14 rounds).
291pub fn encrypt_block<const N: usize>(round_keys: &RoundKeysSoftware<N>, block: &[u8; 16]) -> [u8; 16] {
292    const {
293        assert!(N == 11 || N == 15);
294    }
295
296    let mut s = *block;
297
298    // Round 0: AddRoundKey
299    for i in 0..16 {
300        s[i] ^= round_keys[0][i];
301    }
302
303    // Rounds 1..N-2: SubBytes + ShiftRows + MixColumns + AddRoundKey via T-tables
304    let round_keys_pointer = round_keys.as_ptr();
305    for round in 1..N - 1 {
306        let round_keys = unsafe { &*round_keys_pointer.add(round) };
307        let t0 = TE0[s[0] as usize]
308            ^ TE1[s[5] as usize]
309            ^ TE2[s[10] as usize]
310            ^ TE3[s[15] as usize]
311            ^ u32::from_ne_bytes(round_keys[0..4].try_into().unwrap());
312        let t1 = TE0[s[4] as usize]
313            ^ TE1[s[9] as usize]
314            ^ TE2[s[14] as usize]
315            ^ TE3[s[3] as usize]
316            ^ u32::from_ne_bytes(round_keys[4..8].try_into().unwrap());
317        let t2 = TE0[s[8] as usize]
318            ^ TE1[s[13] as usize]
319            ^ TE2[s[2] as usize]
320            ^ TE3[s[7] as usize]
321            ^ u32::from_ne_bytes(round_keys[8..12].try_into().unwrap());
322        let t3 = TE0[s[12] as usize]
323            ^ TE1[s[1] as usize]
324            ^ TE2[s[6] as usize]
325            ^ TE3[s[11] as usize]
326            ^ u32::from_ne_bytes(round_keys[12..16].try_into().unwrap());
327
328        s[0..4].copy_from_slice(&t0.to_ne_bytes());
329        s[4..8].copy_from_slice(&t1.to_ne_bytes());
330        s[8..12].copy_from_slice(&t2.to_ne_bytes());
331        s[12..16].copy_from_slice(&t3.to_ne_bytes());
332    }
333
334    // Final round (N-1): SubBytes + ShiftRows + AddRoundKey (no MixColumns)
335    let rk_last = unsafe { &*round_keys_pointer.add(N - 1) };
336    s = [
337        SBOX[s[0] as usize] ^ rk_last[0],
338        SBOX[s[5] as usize] ^ rk_last[1],
339        SBOX[s[10] as usize] ^ rk_last[2],
340        SBOX[s[15] as usize] ^ rk_last[3],
341        SBOX[s[4] as usize] ^ rk_last[4],
342        SBOX[s[9] as usize] ^ rk_last[5],
343        SBOX[s[14] as usize] ^ rk_last[6],
344        SBOX[s[3] as usize] ^ rk_last[7],
345        SBOX[s[8] as usize] ^ rk_last[8],
346        SBOX[s[13] as usize] ^ rk_last[9],
347        SBOX[s[2] as usize] ^ rk_last[10],
348        SBOX[s[7] as usize] ^ rk_last[11],
349        SBOX[s[12] as usize] ^ rk_last[12],
350        SBOX[s[1] as usize] ^ rk_last[13],
351        SBOX[s[6] as usize] ^ rk_last[14],
352        SBOX[s[11] as usize] ^ rk_last[15],
353    ];
354    s
355}
356
357/// Decrypt one 16-byte block using inverse T-tables.
358///
359/// Note: the inverse T-tables (TD0..TD3) combine InvSubBytes + InvMixColumns.
360/// Because AddRoundKey falls between InvSubBytes and InvMixColumns in the
361/// decryption round, each round key rk[1]..rk[N-2] must have InvMixColumns
362/// applied to it before the XOR.
363///
364/// `N = 11` for AES-128 (10 rounds), `N = 15` for AES-256 (14 rounds).
365pub fn decrypt_block<const N: usize>(round_keys: &RoundKeysSoftware<N>, block: &[u8; 16]) -> [u8; 16] {
366    const {
367        assert!(N == 11 || N == 15);
368    }
369
370    let mut s = *block;
371
372    // Round N-1 reversed: AddRoundKey
373    let p_last = unsafe { &*round_keys.as_ptr().add(N - 1) };
374    for i in 0..16 {
375        s[i] ^= p_last[i];
376    }
377
378    // Rounds N-2..1: InvShiftRows + InvSubBytes + AddRoundKey + InvMixColumns via inverse T-tables
379    let round_keys_pointer = round_keys.as_ptr();
380    for round in (1..N - 1).rev() {
381        // Apply InvMixColumns to round key columns for correct T-table XOR
382        let mut rk_adj = unsafe { *round_keys_pointer.add(round) };
383        inv_mix_columns(&mut rk_adj);
384
385        let t0 = TD0[s[0] as usize]
386            ^ TD1[s[13] as usize]
387            ^ TD2[s[10] as usize]
388            ^ TD3[s[7] as usize]
389            ^ u32::from_ne_bytes(rk_adj[0..4].try_into().unwrap());
390        let t1 = TD0[s[4] as usize]
391            ^ TD1[s[1] as usize]
392            ^ TD2[s[14] as usize]
393            ^ TD3[s[11] as usize]
394            ^ u32::from_ne_bytes(rk_adj[4..8].try_into().unwrap());
395        let t2 = TD0[s[8] as usize]
396            ^ TD1[s[5] as usize]
397            ^ TD2[s[2] as usize]
398            ^ TD3[s[15] as usize]
399            ^ u32::from_ne_bytes(rk_adj[8..12].try_into().unwrap());
400        let t3 = TD0[s[12] as usize]
401            ^ TD1[s[9] as usize]
402            ^ TD2[s[6] as usize]
403            ^ TD3[s[3] as usize]
404            ^ u32::from_ne_bytes(rk_adj[12..16].try_into().unwrap());
405
406        s[0..4].copy_from_slice(&t0.to_ne_bytes());
407        s[4..8].copy_from_slice(&t1.to_ne_bytes());
408        s[8..12].copy_from_slice(&t2.to_ne_bytes());
409        s[12..16].copy_from_slice(&t3.to_ne_bytes());
410    }
411
412    // Round 0: InvShiftRows + InvSubBytes + AddRoundKey (no InvMixColumns)
413    s = [
414        SBOX_INV[s[0] as usize] ^ round_keys[0][0],
415        SBOX_INV[s[13] as usize] ^ round_keys[0][1],
416        SBOX_INV[s[10] as usize] ^ round_keys[0][2],
417        SBOX_INV[s[7] as usize] ^ round_keys[0][3],
418        SBOX_INV[s[4] as usize] ^ round_keys[0][4],
419        SBOX_INV[s[1] as usize] ^ round_keys[0][5],
420        SBOX_INV[s[14] as usize] ^ round_keys[0][6],
421        SBOX_INV[s[11] as usize] ^ round_keys[0][7],
422        SBOX_INV[s[8] as usize] ^ round_keys[0][8],
423        SBOX_INV[s[5] as usize] ^ round_keys[0][9],
424        SBOX_INV[s[2] as usize] ^ round_keys[0][10],
425        SBOX_INV[s[15] as usize] ^ round_keys[0][11],
426        SBOX_INV[s[12] as usize] ^ round_keys[0][12],
427        SBOX_INV[s[9] as usize] ^ round_keys[0][13],
428        SBOX_INV[s[6] as usize] ^ round_keys[0][14],
429        SBOX_INV[s[3] as usize] ^ round_keys[0][15],
430    ];
431    s
432}
433
434// ── helpers ────────────────────────────────────────────────────────
435
436/// Multiply in GF(2^8) mod 0x11b.
437#[inline(always)]
438fn gmul(mut a: u8, mut b: u8) -> u8 {
439    let mut p = 0u8;
440    for _ in 0..8 {
441        if b & 1 != 0 {
442            p ^= a;
443        }
444        let carry = a & 0x80;
445        a <<= 1;
446        if carry != 0 {
447            a ^= 0x1b;
448        }
449        b >>= 1;
450    }
451    p
452}
453
454#[inline(always)]
455fn inv_mix_col(s: &mut [u8; 16], col: usize) {
456    let i = col * 4;
457    let s0 = s[i];
458    let s1 = s[i + 1];
459    let s2 = s[i + 2];
460    let s3 = s[i + 3];
461    s[i] = gmul(0x0e, s0) ^ gmul(0x0b, s1) ^ gmul(0x0d, s2) ^ gmul(0x09, s3);
462    s[i + 1] = gmul(0x09, s0) ^ gmul(0x0e, s1) ^ gmul(0x0b, s2) ^ gmul(0x0d, s3);
463    s[i + 2] = gmul(0x0d, s0) ^ gmul(0x09, s1) ^ gmul(0x0e, s2) ^ gmul(0x0b, s3);
464    s[i + 3] = gmul(0x0b, s0) ^ gmul(0x0d, s1) ^ gmul(0x09, s2) ^ gmul(0x0e, s3);
465}
466
467#[inline(always)]
468fn inv_mix_columns(state: &mut [u8; 16]) {
469    inv_mix_col(state, 0);
470    inv_mix_col(state, 1);
471    inv_mix_col(state, 2);
472    inv_mix_col(state, 3);
473}
474
475#[cfg(test)]
476mod tests {
477    use super::*;
478
479    // TODO: Add tests:
480    // https://www.tuhs.org/cgi-bin/utree.pl?file=OpenBSD-4.6/regress/sys/crypto/aes/vectors/ecbnk48.txt
481    // https://android.googlesource.com/platform/libcore/+/1db6bf619611525020518a180f0ee82c8cd50af2/luni/src/test/resources/crypto/aes-cbc.csv
482
483    // ── AES-256 block cipher (FIPS 197 Appendix B + C) ────────────────────────
484
485    /// NIST FIPS 197 Appendix B – AES-128 vectors (re-confirmed in AES-256 test)
486    /// These come from FIPS 197 Appendix C.3 (AES-256).
487    #[test]
488    fn fips197_aes256_encrypt() {
489        // FIPS 197 Appendix C.3
490        let key: [u8; 32] =
491            hex::decode_array::<32>(b"000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f").unwrap();
492        let pt: [u8; 16] = hex::decode_array::<16>(b"00112233445566778899aabbccddeeff").unwrap();
493        let ct_expected: [u8; 16] = hex::decode_array::<16>(b"8ea2b7ca516745bfeafc49904b496089").unwrap();
494
495        let rk = expand_key::<15>(&key);
496        let ct = encrypt_block(&rk, &pt);
497        assert_eq!(ct, ct_expected);
498    }
499
500    #[test]
501    fn fips197_aes256_decrypt() {
502        let key: [u8; 32] =
503            hex::decode_array::<32>(b"000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f").unwrap();
504        let ct: [u8; 16] = hex::decode_array::<16>(b"8ea2b7ca516745bfeafc49904b496089").unwrap();
505        let pt_expected: [u8; 16] = hex::decode_array::<16>(b"00112233445566778899aabbccddeeff").unwrap();
506
507        let rk = expand_key::<15>(&key);
508        let pt = decrypt_block(&rk, &ct);
509        assert_eq!(pt, pt_expected);
510    }
511
512    /// NIST SP 800-38A ECB-AES256 vectors (F.1.5 / F.1.6).
513    #[test]
514    fn nist_sp800_38a_aes256_ecb() {
515        let key: [u8; 32] =
516            hex::decode_array::<32>(b"603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4").unwrap();
517
518        let blocks: &[([u8; 16], [u8; 16])] = &[
519            (
520                hex::decode_array::<16>(b"6bc1bee22e409f96e93d7e117393172a").unwrap(),
521                hex::decode_array::<16>(b"f3eed1bdb5d2a03c064b5a7e3db181f8").unwrap(),
522            ),
523            (
524                hex::decode_array::<16>(b"ae2d8a571e03ac9c9eb76fac45af8e51").unwrap(),
525                hex::decode_array::<16>(b"591ccb10d410ed26dc5ba74a31362870").unwrap(),
526            ),
527            (
528                hex::decode_array::<16>(b"30c81c46a35ce411e5fbc1191a0a52ef").unwrap(),
529                hex::decode_array::<16>(b"b6ed21b99ca6f4f9f153e7b1beafed1d").unwrap(),
530            ),
531            (
532                hex::decode_array::<16>(b"f69f2445df4f9b17ad2b417be66c3710").unwrap(),
533                hex::decode_array::<16>(b"23304b7a39f9f3ff067d8d8f9e24ecc7").unwrap(),
534            ),
535        ];
536
537        let rk = expand_key::<15>(&key);
538        for (pt, ct) in blocks {
539            assert_eq!(encrypt_block(&rk, pt), *ct);
540            assert_eq!(decrypt_block(&rk, ct), *pt);
541        }
542    }
543
544    /// NIST Known Answer Test (KAT) – a few AES-256 single-block KATs.
545    #[test]
546    fn aes256_kat_vectors() {
547        // key, plaintext, ciphertext
548        let vectors: &[([u8; 32], [u8; 16], [u8; 16])] = &[
549            // All-zero key and plaintext
550            (
551                [0u8; 32],
552                [0u8; 16],
553                hex::decode_array::<16>(b"dc95c078a2408989ad48a21492842087").unwrap(),
554            ),
555            // Key = 0x01..0x20, PT = 0
556            (
557                hex::decode_array::<32>(b"0101010101010101010101010101010101010101010101010101010101010101").unwrap(),
558                [0u8; 16],
559                hex::decode_array::<16>(b"7298caa565031eadc6ce23d23ea66378").unwrap(),
560            ),
561            // Key = 0xff..0xff
562            (
563                [0xff; 32],
564                [0u8; 16],
565                hex::decode_array::<16>(b"4bf85f1b5d54adbc307b0a048389adcb").unwrap(),
566            ),
567        ];
568
569        for (key, pt, ct_expected) in vectors {
570            let rk = expand_key::<15>(key);
571            let ct = encrypt_block(&rk, pt);
572            assert_eq!(ct, *ct_expected, "key={}", hex::encode(key));
573            let pt2 = decrypt_block(&rk, &ct);
574            assert_eq!(pt2, *pt, "round-trip failed");
575        }
576    }
577
578    #[test]
579    fn encrypt_decrypt_roundtrip_random() {
580        let key: [u8; 32] =
581            hex::decode_array::<32>(b"deadbeefcafebabedeadbeefcafebabe0011223344556677deadbeefcafebabe").unwrap();
582        let rk = expand_key::<15>(&key);
583        for seed in 0u8..=255 {
584            let pt = [seed; 16];
585            let ct = encrypt_block(&rk, &pt);
586            let pt2 = decrypt_block(&rk, &ct);
587            assert_eq!(pt2, pt);
588        }
589    }
590
591    // ── GF(2^128) multiplication ───────────────────────────────────────────────
592
593    #[test]
594    fn gf128_mul_zero() {
595        let h = [
596            0x66, 0xe9, 0x4b, 0xd4, 0xef, 0x8a, 0x2c, 0x3b, 0x88, 0x4c, 0xfa, 0x59, 0xca, 0x34, 0x2b, 0x2e,
597        ];
598        let zero = [0u8; 16];
599        assert_eq!(gf128_mul(&zero, &h), zero);
600        assert_eq!(gf128_mul(&h, &zero), zero);
601    }
602
603    #[test]
604    fn gf128_mul_commutativity() {
605        let a: [u8; 16] = hex::decode_array::<16>(b"66e94bd4ef8a2c3b884cfa59ca342b2e").unwrap();
606        let b: [u8; 16] = hex::decode_array::<16>(b"feedfacedeadbeeffeedfacedeadbeef").unwrap();
607        assert_eq!(gf128_mul(&a, &b), gf128_mul(&b, &a));
608    }
609
610    // H and X from GCM Test Case 2 (NIST SP 800-38D Appendix B).
611    #[test]
612    fn gf128_mul_nist_tv2() {
613        // H = AES_K(0) for K = all-zeros 128-bit key -> irrelevant for 256-bit here,
614        // but we test the raw GF multiplication with known values from NIST test vectors.
615        // From TC2: H = 66e94bd4ef8a2c3b884cfa59ca342b2e
616        //           X (first GHASH input) = feedfacedeadbeeffeedfacedeadbeef
617        // Expected product from the NIST spec reference implementation.
618        let h: [u8; 16] = hex::decode_array::<16>(b"66e94bd4ef8a2c3b884cfa59ca342b2e").unwrap();
619        let x: [u8; 16] = hex::decode_array::<16>(b"feedfacedeadbeeffeedfacedeadbeef").unwrap();
620        // Computed offline with a reference implementation.
621        let expected: [u8; 16] = hex::decode_array::<16>(b"88eddca9968dec8b9c952d6ae0290a82").unwrap();
622        assert_eq!(gf128_mul(&x, &h), expected);
623    }
624}