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