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    [(); 4 * 4 * (NR + 1)]:,
119{
120    let mut w = [[0u8; 4]; 4 * (NR + 1)]; // expanded key
121
122    for i in 0..NK {
123        w[i].clone_from_slice(&key[i * 4..i * 4 + 4]);
124    }
125
126    for i in NK..4 * (NR + 1) {
127        let mut temp = w[i - 1];
128        if i % NK == 0 {
129            let rcon: [u8; 4] = [RCON[i / NK], 0, 0, 0];
130            temp = xor_words(sub_word(rot_word(temp)), rcon);
131        } else if NK > 6 && i % NK == 4 {
132            temp = sub_word(temp);
133        }
134        w[i] = xor_words(w[i - NK], temp);
135    }
136    w
137}
138
139pub(crate) fn rot_word<T: Copy>(w: [T; 4]) -> [T; 4] {
140    array::from_fn(|i| w[(i + 1) % 4])
141}
142fn sub_word(w: [u8; 4]) -> [u8; 4] {
143    let mut r = [0u8; 4];
144    for i in 0..4 {
145        r[i] = SBOX[w[i] as usize];
146    }
147    r
148}
149fn xor_words(w1: [u8; 4], w2: [u8; 4]) -> [u8; 4] {
150    let mut r = [0u8; 4];
151    for i in 0..4 {
152        r[i] = w1[i] ^ w2[i];
153    }
154    r
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160
161    /// test against test vectors from Appendix A
162    #[test]
163    fn test_key_expansion() {
164        // AES-128
165        let key = [
166            0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
167            0x4f, 0x3c,
168        ];
169        let expected: [u32; 4] = [0x2b7e1516, 0x28aed2a6, 0xabf71588, 0x09cf4f3c];
170        let w = key_expansion::</*Nk,Nb,Nr*/ 4, 4, 10>(&key);
171        for i in 0..4 {
172            assert_eq!(expected[i], u32::from_be_bytes(w[i]));
173        }
174
175        // AES-196
176        let key = [
177            0x8e, 0x73, 0xb0, 0xf7, 0xda, 0x0e, 0x64, 0x52, 0xc8, 0x10, 0xf3, 0x2b, 0x80, 0x90,
178            0x79, 0xe5, 0x62, 0xf8, 0xea, 0xd2, 0x52, 0x2c, 0x6b, 0x7b,
179        ];
180        let expected: [u32; 6] = [
181            0x8e73b0f7, 0xda0e6452, 0xc810f32b, 0x809079e5, 0x62f8ead2, 0x522c6b7b,
182        ];
183        let w = key_expansion::</*Nk,Nb,Nr*/ 6, 4, 12>(&key);
184        for i in 0..6 {
185            assert_eq!(expected[i], u32::from_be_bytes(w[i]));
186        }
187
188        // AES-256
189        let key = [
190            0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe, 0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d,
191            0x77, 0x81, 0x1f, 0x35, 0x2c, 0x07, 0x3b, 0x61, 0x08, 0xd7, 0x2d, 0x98, 0x10, 0xa3,
192            0x09, 0x14, 0xdf, 0xf4,
193        ];
194        let expected: [u32; 8] = [
195            0x603deb10, 0x15ca71be, 0x2b73aef0, 0x857d7781, 0x1f352c07, 0x3b6108d7, 0x2d9810a3,
196            0x0914dff4,
197        ];
198        let w = key_expansion::</*Nk,Nb,Nr*/ 8, 4, 14>(&key);
199        for i in 0..8 {
200            assert_eq!(expected[i], u32::from_be_bytes(w[i]));
201        }
202    }
203
204    /// test against paper test vector (appendix B)
205    #[test]
206    fn test_encrypt_block_test_vector() {
207        let input: [u8; 16] = [
208            0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
209            0x07, 0x34,
210        ];
211        let key = [
212            0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
213            0x4f, 0x3c,
214        ];
215        let expanded_key = key_expansion::<4, 4, 10>(&key);
216
217        let c = encrypt_block::<10>(&input, &expanded_key);
218
219        let expected: [u8; 16] = [
220            0x39, 0x02, 0xdc, 0x19, 0x25, 0xdc, 0x11, 0x6a, 0x84, 0x09, 0x85, 0x0b, 0x1d, 0xfb,
221            0x97, 0x32,
222        ];
223
224        assert_eq!(flatten_state(c), expected);
225    }
226
227    /// test with loop of random-values against a rust well-tested library
228    #[test]
229    fn test_encrypt_block_external_lib() {
230        // TODO random-values loop against a rust well-tested library
231    }
232}