plonky2-crypto-gadgets/src/native_aes.rs
src/native_aes.rsRaw
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
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
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]
157 fn test_flatten_state() {
158 let s: [[u8; 4]; 4] = array::from_fn(|i| array::from_fn(|j| (i * 4 + j) as u8));
159
160 // the array::from_fn version is what we use later in the circuit
161 assert_eq!(flatten_state(s), array::from_fn(|i| s[i % 4][i / 4]));
162 }
163
164 /// test against test vectors from Appendix A
165 #[test]
166 fn test_key_expansion() {
167 // AES-128
168 let key = [
169 0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
170 0x4f, 0x3c,
171 ];
172 let expected: [u32; 4] = [0x2b7e1516, 0x28aed2a6, 0xabf71588, 0x09cf4f3c];
173 let w = key_expansion::</*Nk,Nb,Nr*/ 4, 4, 10>(&key);
174 for i in 0..4 {
175 assert_eq!(expected[i], u32::from_be_bytes(w[i]));
176 }
177
178 // AES-196
179 let key = [
180 0x8e, 0x73, 0xb0, 0xf7, 0xda, 0x0e, 0x64, 0x52, 0xc8, 0x10, 0xf3, 0x2b, 0x80, 0x90,
181 0x79, 0xe5, 0x62, 0xf8, 0xea, 0xd2, 0x52, 0x2c, 0x6b, 0x7b,
182 ];
183 let expected: [u32; 6] = [
184 0x8e73b0f7, 0xda0e6452, 0xc810f32b, 0x809079e5, 0x62f8ead2, 0x522c6b7b,
185 ];
186 let w = key_expansion::</*Nk,Nb,Nr*/ 6, 4, 12>(&key);
187 for i in 0..6 {
188 assert_eq!(expected[i], u32::from_be_bytes(w[i]));
189 }
190
191 // AES-256
192 let key = [
193 0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe, 0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d,
194 0x77, 0x81, 0x1f, 0x35, 0x2c, 0x07, 0x3b, 0x61, 0x08, 0xd7, 0x2d, 0x98, 0x10, 0xa3,
195 0x09, 0x14, 0xdf, 0xf4,
196 ];
197 let expected: [u32; 8] = [
198 0x603deb10, 0x15ca71be, 0x2b73aef0, 0x857d7781, 0x1f352c07, 0x3b6108d7, 0x2d9810a3,
199 0x0914dff4,
200 ];
201 let w = key_expansion::</*Nk,Nb,Nr*/ 8, 4, 14>(&key);
202 for i in 0..8 {
203 assert_eq!(expected[i], u32::from_be_bytes(w[i]));
204 }
205 }
206
207 /// test against paper test vector (appendix B)
208 #[test]
209 fn test_encrypt_block_test_vector() {
210 let input: [u8; 16] = [
211 0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
212 0x07, 0x34,
213 ];
214 let key: [u8; 16] = [
215 0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
216 0x4f, 0x3c,
217 ];
218 let expected: [u8; 16] = [
219 0x39, 0x25, 0x84, 0x1d, 0x02, 0xdc, 0x09, 0xfb, 0xdc, 0x11, 0x85, 0x97, 0x19, 0x6a,
220 0x0b, 0x32,
221 ];
222 let expanded_key = key_expansion::<4, 4, 10>(&key);
223 let c = encrypt_block::<10>(&input, &expanded_key);
224 assert_eq!(flatten_state(c), expected);
225 }
226
227 use aes::{
228 cipher::{generic_array::GenericArray, BlockDecrypt, BlockEncrypt, BlockSizeUser, KeyInit},
229 {Aes128, Aes192, Aes256},
230 };
231 use rand::RngExt;
232 /// test with loop of random-values against a rust well-tested library
233 #[test]
234 fn test_encrypt_block_external_lib() {
235 let mut rng = rand::rng();
236
237 // AES-128
238 {
239 const NK: usize = 4;
240 const NB: usize = 4;
241 const NR: usize = 10;
242
243 for _ in 0..10 {
244 let key: [u8; NK * NB] = array::from_fn(|_| rng.random());
245 let block: [u8; 16] = array::from_fn(|_| rng.random());
246 let block = GenericArray::from(block);
247 test_encrypt_block_external_lib_op::<NK, NB, NR, Aes128>(key, block);
248 }
249 }
250 // AES-192
251 {
252 const NK: usize = 6;
253 const NB: usize = 4;
254 const NR: usize = 12;
255
256 for _ in 0..10 {
257 let key: [u8; NK * NB] = array::from_fn(|_| rng.random());
258 let block: [u8; 16] = array::from_fn(|_| rng.random());
259 let block = GenericArray::from(block);
260 test_encrypt_block_external_lib_op::<NK, NB, NR, Aes192>(key, block);
261 }
262 }
263 // AES-256
264 {
265 const NK: usize = 8;
266 const NB: usize = 4;
267 const NR: usize = 14;
268
269 for _ in 0..10 {
270 let key: [u8; NK * NB] = array::from_fn(|_| rng.random());
271 let block: [u8; 16] = array::from_fn(|_| rng.random());
272 let block = GenericArray::from(block);
273 test_encrypt_block_external_lib_op::<NK, NB, NR, Aes256>(key, block);
274 }
275 }
276 }
277 fn test_encrypt_block_external_lib_op<const NK: usize, const NB: usize, const NR: usize, C>(
278 key_arr: [u8; NK * NB],
279 block: GenericArray<u8, C::BlockSize>,
280 ) where
281 C: BlockEncrypt + BlockDecrypt + BlockSizeUser + KeyInit,
282 [(); 4 * (NR + 1)]:,
283 {
284 let cipher = C::new_from_slice(&key_arr).unwrap();
285
286 let mut block_arr: [u8; 16] = [0; 16];
287 block_arr.clone_from_slice(block.as_slice());
288 let mut block = block.clone();
289 let block_copy = block.clone();
290
291 // encrypt block in-place
292 cipher.encrypt_block(&mut block);
293 let binding = block.clone();
294 let expected_cipher = binding.as_slice();
295
296 // and decrypt it back
297 cipher.decrypt_block(&mut block);
298 assert_eq!(block, block_copy);
299
300 // now use our native implementation
301 let expanded_key = key_expansion::<NK, NB, NR>(&key_arr);
302 let c = encrypt_block::<NR>(&block_arr, &expanded_key);
303 // check native implementation vs external lib
304 assert_eq!(flatten_state(c), expected_cipher);
305 }
306}