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