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