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, but simulating the behavior that we
  7//! will do later inside the circuit.
  8
  9use crate::constants::{RCON, SBOX};
 10
 11type State = [[u8; 4]; 4];
 12
 13fn flatten_state(s: State) -> [u8; 16] {
 14    let mut r = [0u8; 16];
 15    for i in 0..4 {
 16        for j in 0..4 {
 17            r[i * 4 + j] = s[i][j];
 18        }
 19    }
 20    r
 21}
 22
 23/// encrypts an AES block (16 bytes). NR determines the number of rounds,
 24/// where AES-128: NR=10, AES-196: NR=12, AES-256: NR=14.
 25fn encrypt_block<const NR: usize>(input: &[u8; 16], w: &[[u8; 4]; 4 * (NR + 1)]) -> State {
 26    assert_eq!(4 * (NR + 1), w.len());
 27
 28    let mut s: State = [[0; 4]; 4];
 29    for i in 0..4 {
 30        for j in 0..4 {
 31            s[i][j] = input[i + 4 * j];
 32        }
 33    }
 34
 35    s = add_round_key(s, &w[0..4]);
 36
 37    for round in 1..NR {
 38        s = sub_bytes(s);
 39        s = shift_rows(s);
 40        s = mix_columns(s);
 41        s = add_round_key(s, &w[4 * round..4 * round + 4]);
 42    }
 43
 44    s = sub_bytes(s);
 45    s = shift_rows(s);
 46    s = add_round_key(s, &w[4 * NR..4 * NR + 4]);
 47
 48    s
 49}
 50
 51fn sub_bytes(s: State) -> State {
 52    let mut r = [[0u8; 4]; 4];
 53    for i in 0..4 {
 54        for j in 0..4 {
 55            r[i][j] = SBOX[s[i][j] as usize];
 56        }
 57    }
 58    r
 59}
 60
 61fn shift_rows(s: State) -> State {
 62    let mut r = [[0u8; 4]; 4];
 63    for i in 0..4 {
 64        for j in 0..4 {
 65            r[i][j] = s[i][(i + j) % 4];
 66        }
 67    }
 68    r
 69}
 70
 71fn 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
 83fn 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
112fn key_expansion<const NK: usize, const NB: usize, const NR: usize>(
113    key: &[u8; NK * NB],
114) -> [[u8; 4]; 4 * (NR + 1)]
115where
116    [(); 4 * (NR + 1)]:,
117    [(); 4 * 4 * (NR + 1)]:,
118{
119    let mut w = [[0u8; 4]; 4 * (NR + 1)]; // expanded key
120
121    for i in 0..NK {
122        w[i].clone_from_slice(&key[i * 4..i * 4 + 4]);
123    }
124
125    for i in NK..4 * (NR + 1) {
126        let mut temp = w[i - 1];
127        if i % NK == 0 {
128            let rcon: [u8; 4] = [RCON[i / NK], 0, 0, 0];
129            temp = xor_words(sub_word(rot_word(temp)), rcon);
130        } else if NK > 6 && i % NK == 4 {
131            temp = sub_word(temp);
132        }
133        w[i] = xor_words(w[i - NK], temp);
134    }
135    w
136}
137
138fn rot_word(w: [u8; 4]) -> [u8; 4] {
139    let mut r = [0u8; 4];
140    (r[0], r[1], r[2], r[3]) = (w[1], w[2], w[3], w[0]);
141    r
142}
143fn sub_word(w: [u8; 4]) -> [u8; 4] {
144    let mut r = [0u8; 4];
145    for i in 0..4 {
146        r[i] = SBOX[w[i] as usize];
147    }
148    r
149}
150fn xor_words(w1: [u8; 4], w2: [u8; 4]) -> [u8; 4] {
151    let mut r = [0u8; 4];
152    for i in 0..4 {
153        r[i] = w1[i] ^ w2[i];
154    }
155    r
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    /// test against test vectors from Appendix A
163    #[test]
164    fn test_key_expansion() {
165        // AES-128
166        let key = [
167            0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
168            0x4f, 0x3c,
169        ];
170        let expected: [u32; 4] = [0x2b7e1516, 0x28aed2a6, 0xabf71588, 0x09cf4f3c];
171        let w = key_expansion::</*Nk,Nb,Nr*/ 4, 4, 10>(&key);
172        for i in 0..4 {
173            assert_eq!(expected[i], u32::from_be_bytes(w[i]));
174        }
175
176        // AES-196
177        let key = [
178            0x8e, 0x73, 0xb0, 0xf7, 0xda, 0x0e, 0x64, 0x52, 0xc8, 0x10, 0xf3, 0x2b, 0x80, 0x90,
179            0x79, 0xe5, 0x62, 0xf8, 0xea, 0xd2, 0x52, 0x2c, 0x6b, 0x7b,
180        ];
181        let expected: [u32; 6] = [
182            0x8e73b0f7, 0xda0e6452, 0xc810f32b, 0x809079e5, 0x62f8ead2, 0x522c6b7b,
183        ];
184        let w = key_expansion::</*Nk,Nb,Nr*/ 6, 4, 12>(&key);
185        for i in 0..6 {
186            assert_eq!(expected[i], u32::from_be_bytes(w[i]));
187        }
188
189        // AES-256
190        let key = [
191            0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe, 0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d,
192            0x77, 0x81, 0x1f, 0x35, 0x2c, 0x07, 0x3b, 0x61, 0x08, 0xd7, 0x2d, 0x98, 0x10, 0xa3,
193            0x09, 0x14, 0xdf, 0xf4,
194        ];
195        let expected: [u32; 8] = [
196            0x603deb10, 0x15ca71be, 0x2b73aef0, 0x857d7781, 0x1f352c07, 0x3b6108d7, 0x2d9810a3,
197            0x0914dff4,
198        ];
199        let w = key_expansion::</*Nk,Nb,Nr*/ 8, 4, 14>(&key);
200        for i in 0..8 {
201            assert_eq!(expected[i], u32::from_be_bytes(w[i]));
202        }
203    }
204
205    /// test against paper test vector (appendix B)
206    #[test]
207    fn test_encrypt_block_test_vector() {
208        let input: [u8; 16] = [
209            0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
210            0x07, 0x34,
211        ];
212        let key = [
213            0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
214            0x4f, 0x3c,
215        ];
216        let expanded_key = key_expansion::<4, 4, 10>(&key);
217
218        let c = encrypt_block::<10>(&input, &expanded_key);
219
220        let expected: [u8; 16] = [
221            0x39, 0x02, 0xdc, 0x19, 0x25, 0xdc, 0x11, 0x6a, 0x84, 0x09, 0x85, 0x0b, 0x1d, 0xfb,
222            0x97, 0x32,
223        ];
224
225        assert_eq!(flatten_state(c), expected);
226    }
227
228    /// test with loop of random-values against a rust well-tested library
229    #[test]
230    fn test_encrypt_block_external_lib() {
231        // TODO random-values loop against a rust well-tested library
232    }
233}