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 std::array;
 10
 11use crate::constants::{RCON, SBOX};
 12
 13pub(crate) type State = [[u8; 4]; 4];
 14
 15fn 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            // TODO
 20            r[i * 4 + j] = s[i][j];
 21        }
 22    }
 23    r
 24}
 25
 26/// encrypts an AES block (16 bytes). NR determines the number of rounds,
 27/// where AES-128: NR=10, AES-196: NR=12, AES-256: NR=14.
 28pub(crate) fn encrypt_block<const NR: usize>(
 29    input: &[u8; 16],
 30    w: &[[u8; 4]; 4 * (NR + 1)],
 31) -> State {
 32    assert_eq!(4 * (NR + 1), w.len());
 33
 34    let mut s: State = [[0; 4]; 4];
 35    for i in 0..4 {
 36        for j in 0..4 {
 37            // TODO
 38            s[i][j] = input[i + 4 * j];
 39        }
 40    }
 41
 42    s = add_round_key(s, &w[0..4]);
 43
 44    for round in 1..NR {
 45        s = sub_bytes(s);
 46        s = shift_rows(s);
 47        s = mix_columns(s);
 48        s = add_round_key(s, &w[4 * round..4 * round + 4]);
 49    }
 50
 51    s = sub_bytes(s);
 52    s = shift_rows(s);
 53    s = add_round_key(s, &w[4 * NR..4 * NR + 4]);
 54
 55    s
 56}
 57
 58pub(crate) fn sub_bytes(s: State) -> State {
 59    let mut r = [[0u8; 4]; 4];
 60    for i in 0..4 {
 61        for j in 0..4 {
 62            r[i][j] = SBOX[s[i][j] as usize];
 63        }
 64    }
 65    r
 66}
 67
 68pub(crate) fn shift_rows<T: Copy>(s: [[T; 4]; 4]) -> [[T; 4]; 4] {
 69    array::from_fn(|i| array::from_fn(|j| s[i][(i + j) % 4]))
 70}
 71
 72pub(crate) fn mix_columns(s: State) -> State {
 73    let mut r = [[0u8; 4]; 4];
 74    for c in 0..4 {
 75        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];
 76        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];
 77        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]);
 78        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]);
 79    }
 80    r
 81}
 82
 83/// multiplication in GF(2^8), section 4.2
 84fn gf_2_8_mul(a_raw: u8, b_raw: u8) -> u8 {
 85    let mut r = 0u8;
 86    let mut a = a_raw;
 87    let mut b = b_raw;
 88    for _ in 0..8 {
 89        if b & 1 == 1 {
 90            r ^= a
 91        }
 92        let high_bit = a & 0x80;
 93        a <<= 1;
 94        if high_bit == 0x80 {
 95            a ^= 0x1b; // ^0x1b = reduce mod m(x) = x^8 + x^4 + x^3 + x + 1
 96        }
 97        b >>= 1;
 98    }
 99    r
100}
101
102fn add_round_key(s: State, key: &[[u8; 4]]) -> State {
103    assert_eq!(key.len(), 4); // ensures that key: [[u8;4]; 4]
104    let mut r: [[u8; 4]; 4] = [[0; 4]; 4];
105    for i in 0..4 {
106        for j in 0..4 {
107            r[i][j] = s[i][j] ^ key[j][i];
108        }
109    }
110    r
111}
112
113pub(crate) fn key_expansion<const NK: usize, const NB: usize, const NR: usize>(
114    key: &[u8; NK * NB],
115) -> [[u8; 4]; 4 * (NR + 1)]
116where
117    [(); 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
138pub(crate) fn rot_word<T: Copy>(w: [T; 4]) -> [T; 4] {
139    array::from_fn(|i| w[(i + 1) % 4])
140}
141fn sub_word(w: [u8; 4]) -> [u8; 4] {
142    let mut r = [0u8; 4];
143    for i in 0..4 {
144        r[i] = SBOX[w[i] as usize];
145    }
146    r
147}
148fn xor_words(w1: [u8; 4], w2: [u8; 4]) -> [u8; 4] {
149    let mut r = [0u8; 4];
150    for i in 0..4 {
151        r[i] = w1[i] ^ w2[i];
152    }
153    r
154}
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159
160    /// test against test vectors from Appendix A
161    #[test]
162    fn test_key_expansion() {
163        // AES-128
164        let key = [
165            0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
166            0x4f, 0x3c,
167        ];
168        let expected: [u32; 4] = [0x2b7e1516, 0x28aed2a6, 0xabf71588, 0x09cf4f3c];
169        let w = key_expansion::</*Nk,Nb,Nr*/ 4, 4, 10>(&key);
170        for i in 0..4 {
171            assert_eq!(expected[i], u32::from_be_bytes(w[i]));
172        }
173
174        // AES-196
175        let key = [
176            0x8e, 0x73, 0xb0, 0xf7, 0xda, 0x0e, 0x64, 0x52, 0xc8, 0x10, 0xf3, 0x2b, 0x80, 0x90,
177            0x79, 0xe5, 0x62, 0xf8, 0xea, 0xd2, 0x52, 0x2c, 0x6b, 0x7b,
178        ];
179        let expected: [u32; 6] = [
180            0x8e73b0f7, 0xda0e6452, 0xc810f32b, 0x809079e5, 0x62f8ead2, 0x522c6b7b,
181        ];
182        let w = key_expansion::</*Nk,Nb,Nr*/ 6, 4, 12>(&key);
183        for i in 0..6 {
184            assert_eq!(expected[i], u32::from_be_bytes(w[i]));
185        }
186
187        // AES-256
188        let key = [
189            0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe, 0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d,
190            0x77, 0x81, 0x1f, 0x35, 0x2c, 0x07, 0x3b, 0x61, 0x08, 0xd7, 0x2d, 0x98, 0x10, 0xa3,
191            0x09, 0x14, 0xdf, 0xf4,
192        ];
193        let expected: [u32; 8] = [
194            0x603deb10, 0x15ca71be, 0x2b73aef0, 0x857d7781, 0x1f352c07, 0x3b6108d7, 0x2d9810a3,
195            0x0914dff4,
196        ];
197        let w = key_expansion::</*Nk,Nb,Nr*/ 8, 4, 14>(&key);
198        for i in 0..8 {
199            assert_eq!(expected[i], u32::from_be_bytes(w[i]));
200        }
201    }
202
203    /// test against paper test vector (appendix B)
204    #[test]
205    fn test_encrypt_block_test_vector() {
206        let input: [u8; 16] = [
207            0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
208            0x07, 0x34,
209        ];
210        let key = [
211            0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
212            0x4f, 0x3c,
213        ];
214        let expanded_key = key_expansion::<4, 4, 10>(&key);
215
216        let c = encrypt_block::<10>(&input, &expanded_key);
217
218        let expected: [u8; 16] = [
219            0x39, 0x02, 0xdc, 0x19, 0x25, 0xdc, 0x11, 0x6a, 0x84, 0x09, 0x85, 0x0b, 0x1d, 0xfb,
220            0x97, 0x32,
221        ];
222
223        assert_eq!(flatten_state(c), expected);
224    }
225
226    /// test with loop of random-values against a rust well-tested library
227    #[test]
228    fn test_encrypt_block_external_lib() {
229        // TODO random-values loop against a rust well-tested library
230    }
231}