1//! For LICENSE check out https://github.com/0xPARC/plonky2-aes/blob/main/LICENSE
  2//!
  3//! Rust native implementation of
  4//! [AES](https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.197-upd1.pdf) having
  5//! in mind the approach that will be done in-circuit; this means that it is not
  6//! written in the most rust idiomatic way, nor memory-efficient, but simulating
  7//! the behavior that we will do later inside the circuit.
  8
  9use std::array;
 10
 11use crate::constants::{RCON, SBOX};
 12
 13pub(crate) type State = [[u8; 4]; 4];
 14
 15pub(crate) fn flatten_state(s: State) -> [u8; 16] {
 16    let mut r = [0u8; 16];
 17    for i in 0..4 {
 18        for j in 0..4 {
 19            r[i * 4 + j] = s[j][i];
 20        }
 21    }
 22    r
 23}
 24
 25/// encrypts an AES block (16 bytes). NR determines the number of rounds,
 26/// where AES-128: NR=10, AES-196: NR=12, AES-256: NR=14.
 27pub(crate) fn encrypt_block<const NR: usize>(
 28    input: &[u8; 16],
 29    w: &[[u8; 4]; 4 * (NR + 1)],
 30) -> State {
 31    assert_eq!(4 * (NR + 1), w.len());
 32
 33    let mut s: State = [[0; 4]; 4];
 34    for i in 0..4 {
 35        for j in 0..4 {
 36            s[i][j] = input[i + 4 * j];
 37        }
 38    }
 39
 40    s = add_round_key(s, &w[0..4]);
 41
 42    for round in 1..NR {
 43        s = sub_bytes(s);
 44        s = shift_rows(s);
 45        s = mix_columns(s);
 46        s = add_round_key(s, &w[4 * round..4 * round + 4]);
 47    }
 48
 49    s = sub_bytes(s);
 50    s = shift_rows(s);
 51    s = add_round_key(s, &w[4 * NR..4 * NR + 4]);
 52
 53    s
 54}
 55
 56pub(crate) fn sub_bytes(s: State) -> State {
 57    let mut r = [[0u8; 4]; 4];
 58    for i in 0..4 {
 59        for j in 0..4 {
 60            r[i][j] = SBOX[s[i][j] as usize];
 61        }
 62    }
 63    r
 64}
 65
 66pub(crate) fn shift_rows<T: Copy>(s: [[T; 4]; 4]) -> [[T; 4]; 4] {
 67    array::from_fn(|i| array::from_fn(|j| s[i][(i + j) % 4]))
 68}
 69
 70pub(crate) fn mix_columns(s: State) -> State {
 71    let mut r = [[0u8; 4]; 4];
 72    for c in 0..4 {
 73        r[0][c] = gf_2_8_mul(0x02, s[0][c]) ^ gf_2_8_mul(0x03, s[1][c]) ^ s[2][c] ^ s[3][c];
 74        r[1][c] = s[0][c] ^ gf_2_8_mul(0x02, s[1][c]) ^ gf_2_8_mul(0x03, s[2][c]) ^ s[3][c];
 75        r[2][c] = s[0][c] ^ s[1][c] ^ gf_2_8_mul(0x02, s[2][c]) ^ gf_2_8_mul(0x03, s[3][c]);
 76        r[3][c] = gf_2_8_mul(0x03, s[0][c]) ^ s[1][c] ^ s[2][c] ^ gf_2_8_mul(0x02, s[3][c]);
 77    }
 78    r
 79}
 80
 81/// multiplication in GF(2^8), section 4.2
 82fn gf_2_8_mul(a_raw: u8, b_raw: u8) -> u8 {
 83    let mut r = 0u8;
 84    let mut a = a_raw;
 85    let mut b = b_raw;
 86    for _ in 0..8 {
 87        if b & 1 == 1 {
 88            r ^= a
 89        }
 90        let high_bit = a & 0x80;
 91        a <<= 1;
 92        if high_bit == 0x80 {
 93            a ^= 0x1b; // ^0x1b = reduce mod m(x) = x^8 + x^4 + x^3 + x + 1
 94        }
 95        b >>= 1;
 96    }
 97    r
 98}
 99
100fn add_round_key(s: State, key: &[[u8; 4]]) -> State {
101    assert_eq!(key.len(), 4); // ensures that key: [[u8;4]; 4]
102    let mut r: [[u8; 4]; 4] = [[0; 4]; 4];
103    for i in 0..4 {
104        for j in 0..4 {
105            r[i][j] = s[i][j] ^ key[j][i];
106        }
107    }
108    r
109}
110
111pub(crate) fn key_expansion<const NK: usize, const NB: usize, const NR: usize>(
112    key: &[u8; NK * NB],
113) -> [[u8; 4]; 4 * (NR + 1)] {
114    let mut w = [[0u8; 4]; 4 * (NR + 1)]; // expanded key
115
116    for i in 0..NK {
117        w[i].clone_from_slice(&key[i * 4..i * 4 + 4]);
118    }
119
120    for i in NK..4 * (NR + 1) {
121        let mut temp = w[i - 1];
122        if i % NK == 0 {
123            let rcon: [u8; 4] = [RCON[i / NK], 0, 0, 0];
124            temp = xor_words(sub_word(rot_word(temp)), rcon);
125        } else if NK > 6 && i % NK == 4 {
126            temp = sub_word(temp);
127        }
128        w[i] = xor_words(w[i - NK], temp);
129    }
130    w
131}
132
133pub(crate) fn rot_word<T: Copy>(w: [T; 4]) -> [T; 4] {
134    array::from_fn(|i| w[(i + 1) % 4])
135}
136fn sub_word(w: [u8; 4]) -> [u8; 4] {
137    let mut r = [0u8; 4];
138    for i in 0..4 {
139        r[i] = SBOX[w[i] as usize];
140    }
141    r
142}
143fn xor_words(w1: [u8; 4], w2: [u8; 4]) -> [u8; 4] {
144    let mut r = [0u8; 4];
145    for i in 0..4 {
146        r[i] = w1[i] ^ w2[i];
147    }
148    r
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    #[test]
156    fn test_flatten_state() {
157        let s: [[u8; 4]; 4] = array::from_fn(|i| array::from_fn(|j| (i * 4 + j) as u8));
158
159        // the array::from_fn version is what we use later in the circuit
160        assert_eq!(flatten_state(s), array::from_fn(|i| s[i % 4][i / 4]));
161    }
162
163    /// test against test vectors from Appendix A
164    #[test]
165    fn test_key_expansion() {
166        // AES-128
167        let key = [
168            0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
169            0x4f, 0x3c,
170        ];
171        let expected: [u32; 4] = [0x2b7e1516, 0x28aed2a6, 0xabf71588, 0x09cf4f3c];
172        let w = key_expansion::</*Nk,Nb,Nr*/ 4, 4, 10>(&key);
173        for i in 0..4 {
174            assert_eq!(expected[i], u32::from_be_bytes(w[i]));
175        }
176
177        // AES-196
178        let key = [
179            0x8e, 0x73, 0xb0, 0xf7, 0xda, 0x0e, 0x64, 0x52, 0xc8, 0x10, 0xf3, 0x2b, 0x80, 0x90,
180            0x79, 0xe5, 0x62, 0xf8, 0xea, 0xd2, 0x52, 0x2c, 0x6b, 0x7b,
181        ];
182        let expected: [u32; 6] = [
183            0x8e73b0f7, 0xda0e6452, 0xc810f32b, 0x809079e5, 0x62f8ead2, 0x522c6b7b,
184        ];
185        let w = key_expansion::</*Nk,Nb,Nr*/ 6, 4, 12>(&key);
186        for i in 0..6 {
187            assert_eq!(expected[i], u32::from_be_bytes(w[i]));
188        }
189
190        // AES-256
191        let key = [
192            0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe, 0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d,
193            0x77, 0x81, 0x1f, 0x35, 0x2c, 0x07, 0x3b, 0x61, 0x08, 0xd7, 0x2d, 0x98, 0x10, 0xa3,
194            0x09, 0x14, 0xdf, 0xf4,
195        ];
196        let expected: [u32; 8] = [
197            0x603deb10, 0x15ca71be, 0x2b73aef0, 0x857d7781, 0x1f352c07, 0x3b6108d7, 0x2d9810a3,
198            0x0914dff4,
199        ];
200        let w = key_expansion::</*Nk,Nb,Nr*/ 8, 4, 14>(&key);
201        for i in 0..8 {
202            assert_eq!(expected[i], u32::from_be_bytes(w[i]));
203        }
204    }
205
206    /// test against paper test vector (appendix B)
207    #[test]
208    fn test_encrypt_block_test_vector() {
209        let input: [u8; 16] = [
210            0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
211            0x07, 0x34,
212        ];
213        let key: [u8; 16] = [
214            0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
215            0x4f, 0x3c,
216        ];
217        let expected: [u8; 16] = [
218            0x39, 0x25, 0x84, 0x1d, 0x02, 0xdc, 0x09, 0xfb, 0xdc, 0x11, 0x85, 0x97, 0x19, 0x6a,
219            0x0b, 0x32,
220        ];
221        let expanded_key = key_expansion::<4, 4, 10>(&key);
222        let c = encrypt_block::<10>(&input, &expanded_key);
223        assert_eq!(flatten_state(c), expected);
224    }
225
226    use aes::{
227        cipher::{BlockDecrypt, BlockEncrypt, BlockSizeUser, KeyInit, generic_array::GenericArray},
228        {Aes128, Aes192, Aes256},
229    };
230    use rand::RngExt;
231    /// test with loop of random-values against a rust well-tested library
232    #[test]
233    fn test_encrypt_block_external_lib() {
234        let mut rng = rand::rng();
235
236        // AES-128
237        {
238            const NK: usize = 4;
239            const NB: usize = 4;
240            const NR: usize = 10;
241
242            for _ in 0..10 {
243                let key: [u8; NK * NB] = array::from_fn(|_| rng.random());
244                let block: [u8; 16] = array::from_fn(|_| rng.random());
245                let block = GenericArray::from(block);
246                test_encrypt_block_external_lib_op::<NK, NB, NR, Aes128>(key, block);
247            }
248        }
249        // AES-192
250        {
251            const NK: usize = 6;
252            const NB: usize = 4;
253            const NR: usize = 12;
254
255            for _ in 0..10 {
256                let key: [u8; NK * NB] = array::from_fn(|_| rng.random());
257                let block: [u8; 16] = array::from_fn(|_| rng.random());
258                let block = GenericArray::from(block);
259                test_encrypt_block_external_lib_op::<NK, NB, NR, Aes192>(key, block);
260            }
261        }
262        // AES-256
263        {
264            const NK: usize = 8;
265            const NB: usize = 4;
266            const NR: usize = 14;
267
268            for _ in 0..10 {
269                let key: [u8; NK * NB] = array::from_fn(|_| rng.random());
270                let block: [u8; 16] = array::from_fn(|_| rng.random());
271                let block = GenericArray::from(block);
272                test_encrypt_block_external_lib_op::<NK, NB, NR, Aes256>(key, block);
273            }
274        }
275    }
276    fn test_encrypt_block_external_lib_op<const NK: usize, const NB: usize, const NR: usize, C>(
277        key_arr: [u8; NK * NB],
278        block: GenericArray<u8, C::BlockSize>,
279    ) where
280        C: BlockEncrypt + BlockDecrypt + BlockSizeUser + KeyInit,
281        [(); 4 * (NR + 1)]:,
282    {
283        let cipher = C::new_from_slice(&key_arr).unwrap();
284
285        let mut block_arr: [u8; 16] = [0; 16];
286        block_arr.clone_from_slice(block.as_slice());
287        let mut block = block.clone();
288        let block_copy = block.clone();
289
290        // encrypt block in-place
291        cipher.encrypt_block(&mut block);
292        let binding = block.clone();
293        let expected_cipher = binding.as_slice();
294
295        // and decrypt it back
296        cipher.decrypt_block(&mut block);
297        assert_eq!(block, block_copy);
298
299        // now use our native implementation
300        let expanded_key = key_expansion::<NK, NB, NR>(&key_arr);
301        let c = encrypt_block::<NR>(&block_arr, &expanded_key);
302        // check native implementation vs external lib
303        assert_eq!(flatten_state(c), expected_cipher);
304    }
305}