plonky2-crypto-gadgets/aes-gcm/src/circuit_aes.rs
aes-gcm/src/circuit_aes.rsRaw
1//! For LICENSE check out https://github.com/0xPARC/plonky2-aes/blob/main/LICENSE
2
3use std::{array, sync::Arc};
4
5use plonky2::{
6 field::{extension::Extendable, goldilocks_field::GoldilocksField as F, types::Field},
7 hash::hash_types::RichField,
8 iop::{
9 target::{BoolTarget, Target},
10 witness::{PartialWitness, WitnessWrite},
11 },
12 plonk::circuit_builder::CircuitBuilder,
13};
14
15use crate::{
16 constants::{RCON, SBOX},
17 native_aes::{rot_word, shift_rows, State},
18 D,
19};
20
21#[derive(Debug, Copy, Clone)]
22pub struct StateTarget(pub [[ByteArrayTarget; 4]; 4]);
23
24// TODO: Normalise terminology!
25pub type ByteArrayTarget = Target;
26
27pub trait CircuitBuilderAESState<F: RichField + Extendable<D>, const D: usize> {
28 /// Adds state target.
29 fn add_virtual_state(&mut self) -> StateTarget;
30
31 /// AES cipher as in spec.
32 fn encrypt_block<const NR: usize>(
33 &mut self,
34 xor_lut_idx: usize,
35 sbox_lut_idx: usize,
36 mix_matrix: [[ByteArrayTarget; 4]; 4],
37 s: StateTarget,
38
39 w: [[ByteArrayTarget; 4]; 4 * (NR + 1)],
40 ) -> StateTarget;
41
42 /// Applies sub_bytes routine to state.
43 fn state_sub_bytes(&mut self, sbox_lut_idx: usize, s: StateTarget) -> StateTarget;
44
45 /// SubWord
46 fn state_sub_word(
47 &mut self,
48 sbox_lut_idx: usize,
49 word: [ByteArrayTarget; 4],
50 ) -> [ByteArrayTarget; 4];
51
52 /// MixColumns
53 fn state_mix_columns(
54 &mut self,
55 xor_lut_idx: usize,
56 mix_matrix: [[ByteArrayTarget; 4]; 4],
57 s: StateTarget,
58 ) -> StateTarget;
59
60 /// AddRoundKey
61 fn state_add_round_key(
62 &mut self,
63 xor_lut_idx: usize,
64 round_key: &[[ByteArrayTarget; 4]],
65 s: StateTarget,
66 ) -> StateTarget;
67
68 /// KeyExpansion
69 fn key_expansion<const NK: usize, const NB: usize, const NR: usize>(
70 &mut self,
71 xor_lut_idx: usize,
72 sbox_lut_idx: usize,
73 key: [ByteArrayTarget; NK * NB],
74 ) -> [[ByteArrayTarget; 4]; 4 * (NR + 1)];
75
76 /// GF(2^8) addition
77 fn gf_2_8_add(&mut self, xor_lut_idx: usize, x: Target, y: Target) -> Target;
78
79 /// GF(2^8) multiplication
80 fn gf_2_8_mul(&mut self, xor_lut_idx: usize, x: Target, y: Target) -> Target;
81
82 /// Bytearray inner product.
83 fn bytearray_ip_bits<const N: usize>(
84 &mut self,
85 xor_lut_idx: usize,
86 x: [ByteArrayTarget; N],
87 y: [ByteArrayTarget; N],
88 ) -> ByteArrayTarget;
89
90 /// Bytearray matrix application.
91 fn bytearray_matrix_apply_bits<const M: usize, const N: usize>(
92 &mut self,
93 xor_lut_idx: usize,
94 a: [[ByteArrayTarget; N]; M],
95 x: [ByteArrayTarget; N],
96 ) -> [ByteArrayTarget; M];
97}
98
99impl CircuitBuilderAESState<F, D> for CircuitBuilder<F, D> {
100 fn add_virtual_state(&mut self) -> StateTarget {
101 StateTarget(array::from_fn(|_| {
102 array::from_fn(|_| self.add_virtual_target())
103 }))
104 }
105
106 fn encrypt_block<const NR: usize>(
107 &mut self,
108 xor_lut_idx: usize,
109 sbox_lut_idx: usize,
110 mix_matrix: [[ByteArrayTarget; 4]; 4], // constant mix_matrix
111 s: StateTarget,
112 w: [[ByteArrayTarget; 4]; 4 * (NR + 1)], // expanded key
113 ) -> StateTarget {
114 let mut s = s;
115 s = self.state_add_round_key(xor_lut_idx, &w[0..4], s);
116 (1..NR).for_each(|i| {
117 s = self.state_sub_bytes(sbox_lut_idx, s);
118 s = StateTarget(shift_rows(s.0));
119 s = self.state_mix_columns(xor_lut_idx, mix_matrix, s);
120 s = self.state_add_round_key(xor_lut_idx, &w[4 * i..4 * (i + 1)], s);
121 });
122 s = self.state_sub_bytes(sbox_lut_idx, s);
123 s = StateTarget(shift_rows(s.0));
124 self.state_add_round_key(xor_lut_idx, &w[4 * NR..4 * (NR + 1)], s)
125 }
126
127 fn state_sub_bytes(&mut self, sbox_lut_idx: usize, s: StateTarget) -> StateTarget {
128 StateTarget(array::from_fn(|i| {
129 self.state_sub_word(sbox_lut_idx, s.0[i])
130 }))
131 }
132
133 fn state_sub_word(
134 &mut self,
135 sbox_lut_idx: usize,
136 word: [ByteArrayTarget; 4],
137 ) -> [ByteArrayTarget; 4] {
138 array::from_fn(|i| {
139 // let byte_target = target_from_bitarray(self, &word[i]);
140 let byte_target = word[i];
141 let out_target = self.add_lookup_from_index(byte_target, sbox_lut_idx);
142 out_target
143 })
144 }
145
146 fn state_mix_columns(
147 &mut self,
148 xor_lut_idx: usize,
149 mix_matrix: [[ByteArrayTarget; 4]; 4],
150 s: StateTarget,
151 ) -> StateTarget {
152 let cols: [_; 4] = array::from_fn(|i| array::from_fn(|j| s.0[j][i]));
153 let out_cols: [_; 4] =
154 array::from_fn(|i| self.bytearray_matrix_apply_bits(xor_lut_idx, mix_matrix, cols[i]));
155 StateTarget(array::from_fn(|i| array::from_fn(|j| out_cols[j][i])))
156 }
157
158 fn state_add_round_key(
159 &mut self,
160 xor_lut_idx: usize,
161 round_key: &[[ByteArrayTarget; 4]],
162 s: StateTarget,
163 ) -> StateTarget {
164 StateTarget(array::from_fn(|i| {
165 array::from_fn(|j| self.gf_2_8_add(xor_lut_idx, s.0[i][j], round_key[j][i]))
166 }))
167 }
168
169 fn key_expansion<const NK: usize, const NB: usize, const NR: usize>(
170 &mut self,
171 xor_lut_idx: usize,
172 sbox_lut_idx: usize,
173 key: [ByteArrayTarget; NK * NB],
174 ) -> [[ByteArrayTarget; 4]; 4 * (NR + 1)] {
175 let rcon: [ByteArrayTarget; 11] =
176 array::from_fn(|i| self.constant(F::from_canonical_u8(RCON[i])));
177
178 let key: [[_; 4]; NK] = array::from_fn(|i| array::from_fn(|j| key[4 * i + j]));
179
180 let key_additions = (NK..4 * (NR + 1))
181 .scan(key.to_vec(), |st, i| {
182 let offset = if i % NK == 0 {
183 let term = self.state_sub_word(sbox_lut_idx, rot_word(st[i - 1]));
184 array::from_fn(|j| {
185 if j == 0 {
186 self.gf_2_8_add(xor_lut_idx, term[j], rcon[i / NK])
187 } else {
188 term[j]
189 }
190 })
191 } else if (NK > 6) && (i % NK == 4) {
192 self.state_sub_word(sbox_lut_idx, st[i - 1])
193 } else {
194 st[i - 1]
195 };
196
197 let cur_val =
198 array::from_fn(|j| self.gf_2_8_add(xor_lut_idx, st[i - NK][j], offset[j]));
199 st.push(cur_val);
200 Some(cur_val)
201 })
202 .collect::<Vec<_>>();
203 std::array::from_fn(|i| {
204 if i < NK {
205 key[i]
206 } else {
207 key_additions[i - NK]
208 }
209 })
210 }
211
212 fn gf_2_8_add(&mut self, xor_lut_idx: usize, x: Target, y: Target) -> Target {
213 byte_xor(self, xor_lut_idx, x, y)
214 }
215
216 fn gf_2_8_mul(&mut self, xor_lut_idx: usize, x: Target, y: Target) -> Target {
217 let zero = self.zero();
218
219 let powers_times_x = (0..8)
220 .scan(x, |st, i| {
221 if i > 0 {
222 *st = byte_x_times(self, xor_lut_idx, *st);
223 }
224 Some(st.clone())
225 })
226 .collect::<Vec<_>>();
227
228 let y_bits = bitarray_from_bytetarget(self, y);
229 let prod_bits = std::iter::zip(y_bits, powers_times_x)
230 .map(|(c, o)| {
231 let o = bitarray_from_bytetarget(self, o);
232 o.into_iter().map(|b| self.and(c, b)).collect::<Vec<_>>()
233 })
234 .collect::<Vec<_>>()
235 .into_iter()
236 .fold(zero, |acc, term| {
237 let term = target_from_bitarray(self, &term);
238 byte_xor(self, xor_lut_idx, acc, term)
239 });
240
241 prod_bits
242 }
243
244 fn bytearray_ip_bits<const N: usize>(
245 &mut self,
246 xor_lut_idx: usize,
247 x: [ByteArrayTarget; N],
248 y: [ByteArrayTarget; N],
249 ) -> ByteArrayTarget {
250 let zero = self.zero();
251
252 std::iter::zip(x, y).fold(zero, |acc, (a, b)| {
253 let prod = self.gf_2_8_mul(xor_lut_idx, a, b);
254 self.gf_2_8_add(xor_lut_idx, acc, prod)
255 })
256 }
257
258 fn bytearray_matrix_apply_bits<const M: usize, const N: usize>(
259 &mut self,
260 xor_lut_idx: usize,
261 a: [[ByteArrayTarget; N]; M],
262 x: [ByteArrayTarget; N],
263 ) -> [ByteArrayTarget; M] {
264 std::array::from_fn(|i| self.bytearray_ip_bits(xor_lut_idx, a[i], x))
265 }
266}
267
268pub trait PartialWitnessByteArray {
269 fn set_byte_array_target(&mut self, target: ByteArrayTarget, value: u8) -> anyhow::Result<()>;
270}
271
272impl<F: Field> PartialWitnessByteArray for PartialWitness<F> {
273 fn set_byte_array_target(&mut self, target: ByteArrayTarget, value: u8) -> anyhow::Result<()> {
274 self.set_target(target, F::from_canonical_u8(value))
275 }
276}
277
278pub trait PartialWitnessAESState {
279 fn set_target_state(&mut self, target: StateTarget, value: State) -> anyhow::Result<()>;
280}
281
282impl<F: Field> PartialWitnessAESState for PartialWitness<F> {
283 fn set_target_state(&mut self, target: StateTarget, value: State) -> anyhow::Result<()> {
284 std::iter::zip(target.0, value).try_for_each(|(t, v)| {
285 std::iter::zip(t, v).try_for_each(|(t, v)| self.set_byte_array_target(t, v))
286 })
287 }
288}
289
290pub fn sbox_lut(builder: &mut CircuitBuilder<F, D>) -> usize {
291 builder.add_lookup_table_from_pairs(Arc::new(
292 SBOX.into_iter()
293 .enumerate()
294 .map(|(i, o)| (i as u16, o as u16))
295 .collect::<Vec<_>>(),
296 ))
297}
298
299pub fn byte_xor_lut(builder: &mut CircuitBuilder<F, D>) -> usize {
300 let xor_table: Vec<(u16, u16)> = (0..=u8::MAX as usize)
301 .flat_map(|x| {
302 (0..=u8::MAX as usize)
303 .map(|y| (((x as u16) << 8) + y as u16, (x as u16) ^ (y as u16)))
304 .collect::<Vec<_>>()
305 })
306 .collect();
307 builder.add_lookup_table_from_pairs(Arc::new(xor_table))
308}
309
310pub fn byte_and_lut(builder: &mut CircuitBuilder<F, D>) -> usize {
311 let and_table: Vec<(u16, u16)> = (0..=u8::MAX as usize)
312 .flat_map(|x| {
313 (0..=u8::MAX as usize)
314 .map(|y| (((x as u16) << 8) + y as u16, (x as u16) & (y as u16)))
315 .collect::<Vec<_>>()
316 })
317 .collect();
318 builder.add_lookup_table_from_pairs(Arc::new(and_table))
319}
320
321pub fn byte_x_times_lut(builder: &mut CircuitBuilder<F, D>) -> usize {
322 let x_times_table: Vec<(u16, u16)> = (0..=u8::MAX as usize)
323 .map(|y| {
324 (
325 y as u16,
326 if y < (1 << 7) {
327 (y << 1) as u16
328 } else {
329 (((y & ((1 << 7) - 1)) << 7) ^ (0x1b)) as u16
330 },
331 )
332 })
333 .collect();
334 builder.add_lookup_table_from_pairs(Arc::new(x_times_table))
335}
336
337pub fn state_mix_matrix_bits(builder: &mut CircuitBuilder<F, D>) -> [[ByteArrayTarget; 4]; 4] {
338 let one = builder.one();
339 let two = builder.two();
340 let three = builder.constant(F::from_canonical_u64(3));
341
342 let matrix = [
343 [two, three, one, one],
344 [one, two, three, one],
345 [one, one, two, three],
346 [three, one, one, two],
347 ];
348
349 matrix
350}
351
352pub fn xor(builder: &mut CircuitBuilder<F, D>, x: BoolTarget, y: BoolTarget) -> BoolTarget {
353 let x_or_y = builder.or(x, y);
354 BoolTarget::new_unsafe(builder.arithmetic(
355 F::NEG_ONE,
356 F::ONE,
357 x.target,
358 y.target,
359 x_or_y.target,
360 ))
361}
362
363pub fn byte_xor(
364 builder: &mut CircuitBuilder<F, D>,
365 xor_lut_idx: usize,
366 x: ByteArrayTarget,
367 y: ByteArrayTarget,
368) -> ByteArrayTarget {
369 let lookup_idx = builder.mul_const_add(F::from_canonical_u64(1 << 8), x, y);
370 builder.add_lookup_from_index(lookup_idx, xor_lut_idx)
371}
372
373pub fn byte_x_times(
374 builder: &mut CircuitBuilder<F, D>,
375 xor_lut_idx: usize,
376 y: Target,
377) -> ByteArrayTarget {
378 let y = bitarray_from_bytetarget(builder, y);
379 let zero = builder._false();
380 let one = builder._true();
381
382 let potential_offset = [one, one, zero, one, one, zero, zero, zero];
383 let y_shifted: [_; 8] = array::from_fn(|i| if i == 0 { zero } else { y[i - 1] });
384
385 let offset: [_; 8] = array::from_fn(|i| builder.and(y[7], potential_offset[i]));
386
387 let y_shifted = target_from_bitarray(builder, &y_shifted);
388 let offset = target_from_bitarray(builder, &offset);
389 byte_xor(builder, xor_lut_idx, y_shifted, offset)
390}
391
392// LE bit array conversions
393pub fn target_from_bitarray(builder: &mut CircuitBuilder<F, D>, x: &[BoolTarget]) -> Target {
394 let zero = builder.zero();
395 let two = builder.two();
396 x.iter()
397 .rev()
398 .fold(zero, |acc, b| builder.mul_add(two, acc, b.target))
399}
400
401pub fn bitarray_from_bytetarget(builder: &mut CircuitBuilder<F, D>, x: Target) -> [BoolTarget; 8] {
402 let x_bits = builder.split_le(x, 8);
403 array::from_fn(|i| x_bits[i])
404}
405
406pub fn le_bits_from_byte(v: u8) -> [bool; 8] {
407 let v_bits = (0..8)
408 .scan(v, |st, _| {
409 let cur_bit = (*st % 2) != 0;
410 *st >>= 1;
411 Some(cur_bit)
412 })
413 .collect::<Vec<_>>();
414 array::from_fn(|i| v_bits[i])
415}
416
417#[cfg(test)]
418mod tests {
419 use std::array;
420
421 use anyhow::Result;
422 use plonky2::{
423 field::{goldilocks_field::GoldilocksField as F, types::Field},
424 iop::witness::{PartialWitness, WitnessWrite},
425 plonk::{
426 circuit_builder::CircuitBuilder, circuit_data::CircuitConfig,
427 config::PoseidonGoldilocksConfig,
428 },
429 };
430 use rand::RngExt;
431
432 use super::*;
433 use crate::native_aes::{encrypt_block, key_expansion, mix_columns, sub_bytes, State};
434
435 #[test]
436 fn test_sub_bytes() -> Result<()> {
437 // Circuit declaration
438 let config = CircuitConfig::standard_recursion_config();
439 let mut builder = CircuitBuilder::<F, D>::new(config);
440
441 let state_target = builder.add_virtual_state();
442 let sbox_lut = sbox_lut(&mut builder);
443 let out_state_target = builder.state_sub_bytes(sbox_lut, state_target);
444
445 let data = builder.build::<PoseidonGoldilocksConfig>();
446
447 // set values to circuit
448 let mut rng = rand::rng();
449 let test_states: [State; 10] = array::from_fn(|_| array::from_fn(|_| rng.random()));
450
451 test_states.into_iter().try_for_each(|s| {
452 let expected_result = sub_bytes(s);
453
454 let mut pw = PartialWitness::<F>::new();
455 pw.set_target_state(state_target, s)?;
456 pw.set_target_state(out_state_target, expected_result)?;
457
458 let proof = data.prove(pw)?;
459 data.verify(proof)
460 })
461 }
462
463 #[test]
464 fn test_mix_columns() -> Result<()> {
465 // Circuit declaration
466 let config = CircuitConfig::standard_recursion_config();
467 let mut builder = CircuitBuilder::<F, D>::new(config);
468
469 let xor_lut_idx = byte_xor_lut(&mut builder);
470 let state_target = builder.add_virtual_state();
471 let mix_matrix = state_mix_matrix_bits(&mut builder);
472 let out_state_target = builder.state_mix_columns(xor_lut_idx, mix_matrix, state_target);
473
474 let data = builder.build::<PoseidonGoldilocksConfig>();
475
476 // set values to circuit
477 let mut rng = rand::rng();
478 let test_states: [State; 10] = array::from_fn(|_| array::from_fn(|_| rng.random()));
479
480 test_states.into_iter().try_for_each(|s| {
481 let expected_result = mix_columns(s);
482
483 let mut pw = PartialWitness::<F>::new();
484 pw.set_target_state(state_target, s)?;
485 pw.set_target_state(out_state_target, expected_result)?;
486
487 let proof = data.prove(pw)?;
488 data.verify(proof)
489 })
490 }
491
492 #[test]
493 fn test_gf_2_8_mul() -> Result<()> {
494 // Circuit declaration
495 let config = CircuitConfig::standard_recursion_config();
496 let mut builder = CircuitBuilder::<F, D>::new(config);
497
498 let xor_lut_idx = byte_xor_lut(&mut builder);
499
500 let a = builder.add_virtual_target();
501 let b = builder.add_virtual_target();
502
503 let a_times_b = builder.gf_2_8_mul(xor_lut_idx, a, b);
504
505 let data = builder.build::<PoseidonGoldilocksConfig>();
506
507 // set values to circuit
508 let test_values = [
509 (0x57, 0x01, 0x57),
510 (0x57, 0x02, 0xae),
511 (0x57, 0x04, 0x47),
512 (0x57, 0x08, 0x8e),
513 (0x57, 0x10, 0x07),
514 (0x57, 0x20, 0x0e),
515 (0x57, 0x40, 0x1c),
516 (0x57, 0x80, 0x38),
517 (0x57, 0x13, 0xfe),
518 ];
519 test_values
520 .into_iter()
521 .map(|(a, b, c)| {
522 (
523 F::from_canonical_u8(a),
524 F::from_canonical_u8(b),
525 F::from_canonical_u8(c),
526 )
527 })
528 .try_for_each(|(a_val, b_val, expected_result)| {
529 let mut pw = PartialWitness::<F>::new();
530 pw.set_target(a, a_val)?;
531 pw.set_target(b, b_val)?;
532 pw.set_target(a_times_b, expected_result)?;
533
534 let proof = data.prove(pw)?;
535 data.verify(proof)
536 })
537 }
538
539 #[test]
540 fn test_gf_2_8_add() -> Result<()> {
541 // Circuit declaration
542 let config = CircuitConfig::standard_recursion_config();
543 let mut builder = CircuitBuilder::<F, D>::new(config);
544
545 let xor_lut_idx = byte_xor_lut(&mut builder);
546
547 let a = builder.add_virtual_target();
548 let b = builder.add_virtual_target();
549
550 let a_plus_b = builder.gf_2_8_add(xor_lut_idx, a, b);
551
552 let data = builder.build::<PoseidonGoldilocksConfig>();
553
554 // set values to circuit
555 let mut rng = rand::rng();
556 let test_values: [(u8, u8, u8); 20] = std::array::from_fn(|_| {
557 let x = rng.random();
558 let y = rng.random();
559 (x, y, x ^ y)
560 });
561 test_values
562 .into_iter()
563 .map(|(a, b, c)| {
564 (
565 F::from_canonical_u8(a),
566 F::from_canonical_u8(b),
567 F::from_canonical_u8(c),
568 )
569 })
570 .try_for_each(|(a_val, b_val, expected_result)| {
571 let mut pw = PartialWitness::<F>::new();
572 pw.set_target(a, a_val)?;
573 pw.set_target(b, b_val)?;
574 pw.set_target(a_plus_b, expected_result)?;
575
576 let proof = data.prove(pw)?;
577 data.verify(proof)
578 })
579 }
580
581 #[test]
582 fn test_key_expansion() -> Result<()> {
583 // AES-128
584 let key = [
585 0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
586 0x4f, 0x3c,
587 ];
588 test_key_expansion_op::</*Nk,Nb,Nr*/ 4, 4, 10>(key)?;
589
590 // AES-196
591 let key = [
592 0x8e, 0x73, 0xb0, 0xf7, 0xda, 0x0e, 0x64, 0x52, 0xc8, 0x10, 0xf3, 0x2b, 0x80, 0x90,
593 0x79, 0xe5, 0x62, 0xf8, 0xea, 0xd2, 0x52, 0x2c, 0x6b, 0x7b,
594 ];
595 test_key_expansion_op::</*Nk,Nb,Nr*/ 6, 4, 12>(key)?;
596
597 // AES-256
598 let key = [
599 0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe, 0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d,
600 0x77, 0x81, 0x1f, 0x35, 0x2c, 0x07, 0x3b, 0x61, 0x08, 0xd7, 0x2d, 0x98, 0x10, 0xa3,
601 0x09, 0x14, 0xdf, 0xf4,
602 ];
603 test_key_expansion_op::</*Nk,Nb,Nr*/ 8, 4, 14>(key)?;
604
605 Ok(())
606 }
607
608 /// Analogous tests to `native_aes::tests::test_key_expansion`.
609 fn test_key_expansion_op<const NK: usize, const NB: usize, const NR: usize>(
610 key: [u8; NK * NB],
611 ) -> Result<()>
612 where
613 [(); 4 * (NR + 1)]:,
614 [(); 4 * NK]:,
615 {
616 // Circuit declaration
617 let config = CircuitConfig::standard_recursion_config();
618 let mut builder = CircuitBuilder::<F, D>::new(config);
619
620 let key_target: [ByteArrayTarget; NK * NB] =
621 array::from_fn(|_| builder.add_virtual_target());
622 let xor_lut_idx = byte_xor_lut(&mut builder);
623 let sbox_lut_idx = sbox_lut(&mut builder);
624 let expanded_key_target =
625 builder.key_expansion::<NK, NB, NR>(xor_lut_idx, sbox_lut_idx, key_target);
626
627 println!(
628 "key_expansion(NK:{}, NB:{}, NR:{}) num_gates: {}",
629 NK,
630 NB,
631 NR,
632 builder.num_gates()
633 );
634 let data = builder.build::<PoseidonGoldilocksConfig>();
635
636 // set values to circuit
637 let w = key_expansion::<NK, NB, NR>(&key);
638
639 let mut pw = PartialWitness::<F>::new();
640
641 std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
642
643 std::iter::zip(expanded_key_target, w).try_for_each(|(t, v)| {
644 std::iter::zip(t, v).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))
645 })?;
646
647 let proof = data.prove(pw)?;
648 data.verify(proof)
649 }
650
651 /// test against native version
652 #[test]
653 fn test_encrypt_block_test_vector() -> Result<()> {
654 // AES-128
655 let input_state: [u8; 16] = [
656 0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
657 0x07, 0x34,
658 ];
659 let key = [
660 0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
661 0x4f, 0x3c,
662 ];
663 test_encrypt_block_test_vector_op::<4, 4, 10>(input_state, key)?;
664
665 // AES-196
666 let input_state: [u8; 16] = [
667 0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
668 0x07, 0x34,
669 ];
670 let key = [
671 0x8e, 0x73, 0xb0, 0xf7, 0xda, 0x0e, 0x64, 0x52, 0xc8, 0x10, 0xf3, 0x2b, 0x80, 0x90,
672 0x79, 0xe5, 0x62, 0xf8, 0xea, 0xd2, 0x52, 0x2c, 0x6b, 0x7b,
673 ];
674 test_encrypt_block_test_vector_op::<6, 4, 12>(input_state, key)?;
675
676 // AES-256
677 let input_state: [u8; 16] = [
678 0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
679 0x07, 0x34,
680 ];
681 let key = [
682 0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe, 0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d,
683 0x77, 0x81, 0x1f, 0x35, 0x2c, 0x07, 0x3b, 0x61, 0x08, 0xd7, 0x2d, 0x98, 0x10, 0xa3,
684 0x09, 0x14, 0xdf, 0xf4,
685 ];
686 test_encrypt_block_test_vector_op::<8, 4, 14>(input_state, key)?;
687
688 Ok(())
689 }
690
691 fn test_encrypt_block_test_vector_op<const NK: usize, const NB: usize, const NR: usize>(
692 input_state: [u8; 16],
693 key: [u8; NK * NB],
694 ) -> Result<()>
695 where
696 [(); 4 * (NR + 1)]:,
697 [(); 4 * NK]:,
698 // [(); NK * NB]:,
699 {
700 // Circuit declaration
701 let config = CircuitConfig::standard_recursion_config();
702 let mut builder = CircuitBuilder::<F, D>::new(config);
703
704 let key_target: [ByteArrayTarget; NK * NB] =
705 array::from_fn(|_| builder.add_virtual_target());
706 let xor_lut_idx = byte_xor_lut(&mut builder);
707 let sbox_lut_idx = sbox_lut(&mut builder);
708 let mix_matrix = state_mix_matrix_bits(&mut builder);
709 let expanded_key_target: [[ByteArrayTarget; 4]; 4 * (NR + 1)] =
710 builder.key_expansion::<NK, NB, NR>(xor_lut_idx, sbox_lut_idx, key_target);
711
712 let input_state_target = builder.add_virtual_state();
713
714 let output = builder.encrypt_block(
715 xor_lut_idx,
716 sbox_lut_idx,
717 mix_matrix,
718 input_state_target,
719 expanded_key_target,
720 );
721
722 println!(
723 "encrypt_block (NK:{}, NB:{}, NR:{}) num_gates: {}",
724 NK,
725 NB,
726 NR,
727 builder.num_gates()
728 );
729 let data = builder.build::<PoseidonGoldilocksConfig>();
730
731 // set values to circuit
732 let mut input_state_matrix: State = [[0; 4]; 4];
733 for i in 0..4 {
734 for j in 0..4 {
735 input_state_matrix[i][j] = input_state[i + 4 * j];
736 }
737 }
738
739 let expanded_key: [[u8; 4]; 4 * (NR + 1)] = key_expansion::<NK, NB, NR>(&key);
740 let native_ciphertext = encrypt_block::<NR>(&input_state, &expanded_key);
741
742 let mut pw = PartialWitness::<F>::new();
743
744 std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
745
746 pw.set_target_state(input_state_target, input_state_matrix)?;
747
748 std::iter::zip(expanded_key_target, expanded_key).try_for_each(|(t, v)| {
749 std::iter::zip(t, v).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))
750 })?;
751
752 std::iter::zip(output.0, native_ciphertext).try_for_each(|(o, e)| {
753 std::iter::zip(o, e).try_for_each(|(o, e)| pw.set_byte_array_target(o, e))
754 })?;
755
756 let proof = data.prove(pw)?;
757 data.verify(proof)
758 }
759}