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