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