arnaucube
arnaucube/plonky2-crypto-gadgetsPublic Self-hosted repository
plonky2-crypto-gadgets/aes-gcm/src/circuit_gcm.rs
aes-gcm/src/circuit_gcm.rsRaw
  1//! For LICENSE check out https://github.com/0xPARC/plonky2-crypto-gadgets/blob/main/LICENSE
  2//!
  3//! Plonky2 circuit implementation of
  4//! [AES-GCM (Galois Counter Mode)](https://nvlpubs.nist.gov/nistpubs/Legacy/SP/nistspecialpublication800-38d.pdf).
  5
  6use std::{array, sync::Arc};
  7
  8use anyhow::Result;
  9use plonky2::{
 10    field::{goldilocks_field::GoldilocksField as F, types::Field},
 11    iop::{target::BoolTarget, witness::PartialWitness},
 12    plonk::circuit_builder::CircuitBuilder,
 13};
 14
 15use crate::{
 16    D,
 17    circuit_aes::{
 18        ByteTarget, CircuitBuilderAESState, PartialWitnessByteArray, StateTarget, byte_xor,
 19        byte_xor_lut, gf_2_8_mul_lut, sbox_lut, state_mix_matrix,
 20    },
 21    constants::TAG_LEN,
 22};
 23
 24pub struct AesGcmTarget<
 25    const NK: usize,
 26    const NB: usize,
 27    const NR: usize,
 28    // L: max size of plaintext to cipher.
 29    const L: usize,
 30    const TAG: bool,
 31> where
 32    [(); NK * NB]:,
 33    [(); 4 * (NR + 1)]:,
 34{
 35    key: [ByteTarget; NK * NB],
 36    nonce: [ByteTarget; 12],
 37    pt: [ByteTarget; L],
 38
 39    ct: [ByteTarget; L],
 40    tag: [ByteTarget; TAG_LEN / 8],
 41}
 42
 43impl<const NK: usize, const NB: usize, const NR: usize, const L: usize, const TAG: bool>
 44    AesGcmTarget<NK, NB, NR, L, TAG>
 45where
 46    [(); NK * NB]:,
 47    [(); 4 * (NR + 1)]:,
 48{
 49    pub fn build(builder: &mut CircuitBuilder<F, D>) -> Self {
 50        // Declare LUTs
 51        let sbox_lut = sbox_lut(builder);
 52        let xor_lut = byte_xor_lut(builder);
 53        let gf_2_8_mul_lut = gf_2_8_mul_lut(builder);
 54
 55        // add targets
 56        let key: [ByteTarget; NK * NB] =
 57            array::from_fn(|_| builder.add_virtual_byte_target(sbox_lut));
 58        let nonce: [ByteTarget; 12] = array::from_fn(|_| builder.add_virtual_byte_target(sbox_lut));
 59        let pt: [ByteTarget; L] = array::from_fn(|_| builder.add_virtual_byte_target(sbox_lut));
 60
 61        let tag: [ByteTarget; TAG_LEN / 8] =
 62            array::from_fn(|_| builder.add_virtual_byte_target(sbox_lut));
 63
 64        let mix_matrix = state_mix_matrix(builder);
 65
 66        let a: &[ByteTarget] = &[]; // additional authenticated data
 67
 68        let expanded_key = builder.key_expansion::<NK, NB, NR>(xor_lut, sbox_lut, key);
 69
 70        let empty_state = builder.empty_state();
 71
 72        // 1. CIPH_K
 73        let h = builder
 74            .encrypt_block::<NR>(
 75                xor_lut,
 76                gf_2_8_mul_lut,
 77                sbox_lut,
 78                mix_matrix,
 79                empty_state,
 80                expanded_key,
 81            )
 82            .flatten();
 83
 84        // 2. J_0
 85        let zero_byte = builder.zero_byte();
 86        let one_byte = builder.byte_constant(1);
 87        let j0: [ByteTarget; 16] = if nonce.len() * 8 == 96 {
 88            // J_0 = IV || 0^31 || 1
 89            let mut out = [zero_byte; 16];
 90            out[..12].copy_from_slice(nonce.as_slice());
 91            out[12..16].copy_from_slice(&[zero_byte, zero_byte, zero_byte, one_byte]);
 92            out
 93        } else {
 94            panic!("unsupported at initial version; nonce.len()=12 (96 bits)");
 95        };
 96
 97        // 3. C=GCTR()
 98        let inc32_j0 = inc32_target(builder, j0);
 99        let ct = gctr_target::<NR, L>(
100            builder,
101            (xor_lut, gf_2_8_mul_lut, sbox_lut),
102            mix_matrix,
103            expanded_key,
104            &inc32_j0,
105            &pt,
106        );
107
108        // the rest of the logic is for the tag; for some use cases we might skip it
109        // (saving a notable amount of gates, about double)
110        if !TAG {
111            return Self {
112                key,
113                nonce,
114                pt,
115                ct,
116                tag,
117            };
118        }
119
120        // 4. u, v
121        let u: usize = 16 * (L as f64 / 16_f64).ceil() as usize - L; // L=x.len()
122        let v: usize = 16 * (a.len() as f64 / 16_f64).ceil() as usize - a.len();
123
124        // 5. S = GHASH()
125        let const_a_len_u8: [u8; 8] = (a.len() * 8).to_be_bytes(); // const
126        let const_c_len_u8: [u8; 8] = (ct.len() * 8).to_be_bytes(); // const
127        let a_len: [ByteTarget; 8] =
128            array::from_fn(|byte_i| builder.byte_constant(const_a_len_u8[byte_i]));
129        let c_len: [ByteTarget; 8] =
130            array::from_fn(|byte_i| builder.byte_constant(const_c_len_u8[byte_i]));
131        let ghash_input_vec = [
132            a.to_vec(),
133            vec![zero_byte; v],
134            ct.to_vec(),
135            vec![zero_byte; u],
136            a_len.to_vec(),
137            c_len.to_vec(),
138        ]
139        .concat();
140
141        let u8_unit_right_shift_lut_idx = u8_unit_right_shift_lut(builder);
142        let u8_bitref_lut_idx = u8_bitref_lut(builder);
143        let s = ghash_target(
144            builder,
145            xor_lut,
146            u8_unit_right_shift_lut_idx,
147            u8_bitref_lut_idx,
148            h,
149            ghash_input_vec,
150        );
151
152        // 6. T=MSB(GCTR()))
153        let msb_input = gctr_target(
154            builder,
155            (xor_lut, gf_2_8_mul_lut, sbox_lut),
156            mix_matrix,
157            expanded_key,
158            &j0,
159            &s,
160        );
161        let t_vec = msb_t_target(TAG_LEN / 8, &msb_input);
162
163        std::iter::zip(tag, t_vec).for_each(|(a, b)| builder.connect(a.0, b.0));
164
165        Self {
166            key,
167            nonce,
168            pt,
169            ct,
170            tag,
171        }
172    }
173
174    pub fn set_targets(
175        &self,
176        pw: &mut PartialWitness<F>,
177        key: &[u8; NK * NB],
178        nonce: &[u8; 12],
179        pt: &[u8],
180        ct: &[u8],
181        tag: &[u8],
182    ) -> Result<()> {
183        assert!(pt.len() <= L);
184        assert!(ct.len() <= L);
185        assert!(tag.len() <= TAG_LEN / 8);
186
187        // extend to expected sizes
188        let mut pt_arr: [u8; L] = [0u8; L];
189        pt_arr.copy_from_slice(pt);
190        let mut ct_arr: [u8; L] = [0u8; L];
191        ct_arr.copy_from_slice(ct);
192        let mut tag_arr: [u8; TAG_LEN / 8] = [0u8; TAG_LEN / 8];
193        tag_arr.copy_from_slice(tag);
194
195        std::iter::zip(self.key, key).try_for_each(|(t, v)| pw.set_byte_target(t, *v))?;
196        std::iter::zip(self.nonce, nonce).try_for_each(|(t, v)| pw.set_byte_target(t, *v))?;
197        std::iter::zip(self.pt, pt_arr).try_for_each(|(t, v)| pw.set_byte_target(t, v))?;
198
199        std::iter::zip(self.ct, ct_arr).try_for_each(|(t, v)| pw.set_byte_target(t, v))?;
200        if TAG {
201            std::iter::zip(self.tag, tag_arr).try_for_each(|(t, v)| pw.set_byte_target(t, v))?;
202        } else {
203            self.tag
204                .iter()
205                .try_for_each(|t| pw.set_byte_target(*t, 0u8))?;
206        }
207        Ok(())
208    }
209}
210
211/// L: max size of plaintext to cipher.
212fn gctr_target<const NR: usize, const L: usize>(
213    builder: &mut CircuitBuilder<F, D>,
214    lut_indices: (usize, usize, usize),
215    mix_matrix: [[ByteTarget; 4]; 4],
216    key: [[ByteTarget; 4]; 4 * (NR + 1)],
217    icb: &[ByteTarget; 16],
218    x: &[ByteTarget; L],
219) -> [ByteTarget; L] {
220    let (xor_lut_idx, gf_2_8_mul_lut_idx, sbox_lut_idx) = lut_indices;
221    let n = ((L * 8) as f64 / 128_f64).ceil() as usize;
222
223    let mut y: [ByteTarget; L] = *x;
224    let mut cb_i = *icb;
225
226    let zero_byte = builder.zero_byte();
227    for (i, x_i_raw) in x.chunks(16).enumerate() {
228        if i > 0 {
229            cb_i = inc32_target(builder, cb_i);
230        }
231
232        let l = x_i_raw.len().min(16);
233        let mut x_i = [zero_byte; 16];
234        x_i[..l].copy_from_slice(x_i_raw);
235
236        let ciph_cb_i = builder
237            .encrypt_block::<NR>(
238                xor_lut_idx,
239                gf_2_8_mul_lut_idx,
240                sbox_lut_idx,
241                mix_matrix,
242                StateTarget::from_flat(cb_i),
243                key,
244            )
245            .flatten();
246
247        let (y_i, n_bytes) = if i < n && x_i_raw.len() == 16 {
248            (xor_blocks(builder, xor_lut_idx, x_i, ciph_cb_i), 16)
249        } else {
250            // last chunk, might be smaller than 16 bytes (L%16 bytes)
251            let msb_res = msb_t_target(L % 16, ciph_cb_i.as_slice());
252            let l = msb_res.len().min(16);
253            let mut m: [ByteTarget; 16] = [zero_byte; 16];
254            m[..l].copy_from_slice(&msb_res);
255            (xor_blocks(builder, xor_lut_idx, x_i, m), l)
256        };
257        y[i * 16..i * 16 + n_bytes].clone_from_slice(&y_i[..n_bytes]);
258    }
259    y
260}
261
262fn ghash_target(
263    builder: &mut CircuitBuilder<F, D>,
264    xor_lut_idx: usize,
265    u8_unit_right_shift_lut_idx: usize,
266    u8_bitref_lut_idx: usize,
267    h: [ByteTarget; 16],
268    x: Vec<ByteTarget>,
269) -> [ByteTarget; 16] {
270    assert!(x.len().is_multiple_of(16)); // multiple of 128 bits
271    let m = x.len() / 16;
272
273    let zero_byte = builder.zero_byte();
274    let mut y = [zero_byte; 16]; // (128 bits)
275    for i in 0..m {
276        let mut xi = [zero_byte; 16];
277        xi.clone_from_slice(&x[i * 16..i * 16 + 16]);
278        let y_xi = xor_blocks(builder, xor_lut_idx, y, xi);
279        y = gf_2_128_mul_target(
280            builder,
281            xor_lut_idx,
282            u8_unit_right_shift_lut_idx,
283            u8_bitref_lut_idx,
284            y_xi,
285            h,
286        );
287    }
288    y
289}
290fn gf_2_128_mul_target(
291    builder: &mut CircuitBuilder<F, D>,
292    xor_lut_idx: usize,
293    u8_unit_right_shift_lut_idx: usize,
294    u8_bitref_lut_idx: usize,
295    x: [ByteTarget; 16],
296    y: [ByteTarget; 16],
297) -> [ByteTarget; 16] {
298    let zero = builder.zero_byte();
299
300    // R: 10000111 (= 225) || 0^120 (in little-endian)
301    let r_first = builder.byte_constant(225);
302
303    let mut z = [zero; 16];
304    let mut v = y;
305    for i in 0..128 {
306        let byte_index = i / 8;
307        let bit_index = 7 - (i % 8);
308        let bit_idx_target = builder.byte_constant(bit_index as u8);
309        let xi = u8_bitref(builder, u8_bitref_lut_idx, x[byte_index], bit_idx_target);
310
311        // set z = if xi==1: z^v, else: z
312        for b in 0..16 {
313            let z_xor_v = byte_xor(builder, xor_lut_idx, z[b], v[b]);
314            z[b] = ByteTarget(builder.select(xi, z_xor_v.0, z[b].0));
315        }
316
317        let lsb = u8_bitref(builder, u8_bitref_lut_idx, v[15], zero); // (little-endian)
318
319        v = right_shift_one_target(builder, u8_unit_right_shift_lut_idx, &v);
320
321        // if lsb==1: v=v^R, else: v
322        let v_xor_r = byte_xor(builder, xor_lut_idx, v[0], r_first);
323        v[0] = ByteTarget(builder.select(lsb, v_xor_r.0, v[0].0));
324    }
325    z
326}
327pub fn right_shift_one_target(
328    builder: &mut CircuitBuilder<F, D>,
329    u8_unit_right_shift_lut_idx: usize,
330    v: &[ByteTarget; 16],
331) -> [ByteTarget; 16] {
332    let mut r: [ByteTarget; 16] = *v;
333
334    let mut carry = builder.zero();
335    for i in 0..16 {
336        let current = v[i];
337        let mut shifted = u8_unit_right_shift(builder, u8_unit_right_shift_lut_idx, current);
338        let next_carry =
339            builder.mul_const_add(F::from_canonical_u64(2) * F::NEG_ONE, shifted.0, current.0);
340
341        shifted =
342            ByteTarget(builder.mul_const_add(F::from_canonical_u64(1 << 7), carry, shifted.0));
343
344        r[i] = shifted;
345        carry = next_carry;
346    }
347    r
348}
349
350pub fn inc32_target(
351    builder: &mut CircuitBuilder<F, D>,
352    block: [ByteTarget; 16],
353) -> [ByteTarget; 16] {
354    let mut r = block;
355    let zero = builder.zero();
356    let u8_max = builder.constant(F::from_canonical_u8(u8::MAX));
357
358    let mut carry = builder.one();
359    for byte_index in (12..16).rev() {
360        let a = block[byte_index];
361        let sum = builder.add(a.0, carry);
362        let a_is_u8_max = builder.is_equal(a.0, u8_max);
363        let carry_out = builder.mul(carry, a_is_u8_max.target);
364        r[byte_index] = ByteTarget(builder.select(a_is_u8_max, zero, sum));
365        carry = carry_out;
366    }
367    r
368}
369
370fn xor_blocks(
371    builder: &mut CircuitBuilder<F, D>,
372    xor_lut_idx: usize,
373    b1: [ByteTarget; 16],
374    b2: [ByteTarget; 16],
375) -> [ByteTarget; 16] {
376    array::from_fn(|i| byte_xor(builder, xor_lut_idx, b1[i], b2[i]))
377}
378
379pub fn msb_t_target(t_bytes: usize, block: &[ByteTarget]) -> Vec<ByteTarget> {
380    // Always dealing with full bytes in circuit.
381    let full_bytes = t_bytes;
382
383    let mut out: Vec<ByteTarget> = Vec::new();
384    for block_i in block.iter().take(full_bytes) {
385        out.push(*block_i);
386    }
387
388    out
389}
390
391/// Lookup table for unit right-shift of elements of u8.
392pub fn u8_unit_right_shift_lut(builder: &mut CircuitBuilder<F, D>) -> usize {
393    let u8_unit_right_shift_table: Vec<(u16, u16)> = (0..=u8::MAX as usize)
394        .map(|x| (x as u16, (x as u16) >> 1))
395        .collect();
396    builder.add_lookup_table_from_pairs(Arc::new(u8_unit_right_shift_table))
397}
398pub fn u8_unit_right_shift(
399    builder: &mut CircuitBuilder<F, D>,
400    u8_unit_right_shift_lut_idx: usize,
401    x: ByteTarget,
402) -> ByteTarget {
403    ByteTarget(builder.add_lookup_from_index(x.0, u8_unit_right_shift_lut_idx))
404}
405
406/// Lookup table for u8 bit referencing.
407pub fn u8_bitref_lut(builder: &mut CircuitBuilder<F, D>) -> usize {
408    let u8_bitref_table: Vec<(u16, u16)> = (0..=u8::MAX as usize)
409        .flat_map(|x| {
410            (0..8)
411                .map(|i| (((x as u16) << 3) + i as u16, ((x as u16) >> i) & 1))
412                .collect::<Vec<_>>()
413        })
414        .collect();
415    builder.add_lookup_table_from_pairs(Arc::new(u8_bitref_table))
416}
417pub fn u8_bitref(
418    builder: &mut CircuitBuilder<F, D>,
419    u8_bitref_lut_idx: usize,
420    x: ByteTarget,
421    i: ByteTarget,
422) -> BoolTarget {
423    let lookup_idx = builder.mul_const_add(F::from_canonical_u64(8), x.0, i.0);
424    BoolTarget::new_unsafe(builder.add_lookup_from_index(lookup_idx, u8_bitref_lut_idx))
425}
426
427fn le_bits_from_byte(v: u8) -> [bool; 8] {
428    let v_bits = (0..8)
429        .scan(v, |st, _| {
430            let cur_bit = (*st % 2) != 0;
431            *st >>= 1;
432            Some(cur_bit)
433        })
434        .collect::<Vec<_>>();
435    array::from_fn(|i| v_bits[i])
436}
437
438#[cfg(test)]
439mod tests {
440    use std::array;
441
442    use anyhow::Result;
443    use plonky2::{
444        field::goldilocks_field::GoldilocksField as F,
445        iop::witness::PartialWitness,
446        plonk::{
447            circuit_builder::CircuitBuilder, circuit_data::CircuitConfig,
448            config::PoseidonGoldilocksConfig,
449        },
450    };
451
452    use super::{D, *};
453    use crate::{
454        circuit_aes::{PartialWitnessByteArray, byte_xor_lut, gf_2_8_mul_lut},
455        native_aes::key_expansion,
456        native_gcm::{gctr, gf_2_128_mul, ghash, right_shift_one},
457    };
458
459    #[test]
460    fn test_gctr() -> Result<()> {
461        // AES-GCM-128
462        test_gctr_opt::<4, 4, 10, 13>()?;
463        test_gctr_opt::<4, 4, 10, 17>()?;
464
465        // AES-GCM-256
466        test_gctr_opt::<8, 4, 14, 13>()?;
467
468        Ok(())
469    }
470    fn test_gctr_opt<const NK: usize, const NB: usize, const NR: usize, const L: usize>()
471    -> Result<()>
472    where
473        [(); 4 * (NR + 1)]:,
474        [(); NK * NB]:,
475    {
476        let key: &[u8; NK * NB] = &[42; NK * NB];
477        let expanded_key = key_expansion::<NK, NB, NR>(key);
478        let nonce: &[u8; 12] = &[111; 12];
479        let icb: &[u8; 16] = &[222; 16];
480        let pt: &[u8; L] = &[42u8; L];
481
482        let expected = gctr(expanded_key, icb, pt);
483
484        // Circuit declaration
485        let config = CircuitConfig::standard_recursion_config();
486        let mut builder = CircuitBuilder::<F, D>::new(config);
487
488        let key_target: [ByteTarget; NK * NB] =
489            array::from_fn(|_| builder.add_virtual_byte_target_unsafe());
490        let nonce_target: [ByteTarget; 12] =
491            array::from_fn(|_| builder.add_virtual_byte_target_unsafe());
492        let icb_target: [ByteTarget; 16] =
493            array::from_fn(|_| builder.add_virtual_byte_target_unsafe());
494        let pt_target: [ByteTarget; L] =
495            array::from_fn(|_| builder.add_virtual_byte_target_unsafe());
496        let sbox_lut = sbox_lut(&mut builder);
497        let xor_lut = byte_xor_lut(&mut builder);
498        let gf_2_8_mul_lut = gf_2_8_mul_lut(&mut builder);
499        let mix_matrix = state_mix_matrix(&mut builder);
500
501        let key_expanded_target =
502            builder.key_expansion::<NK, NB, NR>(xor_lut, sbox_lut, key_target);
503        let gctr_out_target = gctr_target::<NR, L>(
504            &mut builder,
505            (xor_lut, gf_2_8_mul_lut, sbox_lut),
506            mix_matrix,
507            key_expanded_target,
508            &icb_target,
509            &pt_target,
510        );
511
512        println!(
513            "gctr (NK:{}, NB:{}, NR:{}, L:{}) num_gates: {}",
514            NK,
515            NB,
516            NR,
517            L,
518            builder.num_gates()
519        );
520        let data = builder.build::<PoseidonGoldilocksConfig>();
521
522        // set values to circuit
523        let mut pw = PartialWitness::<F>::new();
524        std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_target(t, *v))?;
525        std::iter::zip(nonce_target, nonce).try_for_each(|(t, v)| pw.set_byte_target(t, *v))?;
526        std::iter::zip(icb_target, icb).try_for_each(|(t, v)| pw.set_byte_target(t, *v))?;
527        std::iter::zip(pt_target, pt).try_for_each(|(t, v)| pw.set_byte_target(t, *v))?;
528
529        std::iter::zip(gctr_out_target, expected)
530            .try_for_each(|(t, v)| pw.set_byte_target(t, v))?;
531
532        let proof = data.prove(pw)?;
533        data.verify(proof)
534    }
535
536    #[test]
537    fn test_right_shift_one() -> Result<()> {
538        let x: [u8; 16] = [111; 16];
539
540        let mut expected = x;
541        right_shift_one(&mut expected);
542
543        // Circuit declaration
544        let config = CircuitConfig::standard_recursion_config();
545        let mut builder = CircuitBuilder::<F, D>::new(config);
546
547        let u8_unit_right_shift_lut_idx = u8_unit_right_shift_lut(&mut builder);
548
549        let x_target: [ByteTarget; 16] =
550            array::from_fn(|_| builder.add_virtual_byte_target_unsafe());
551
552        let out_target: [ByteTarget; 16] =
553            right_shift_one_target(&mut builder, u8_unit_right_shift_lut_idx, &x_target);
554
555        println!("right_shift_one num_gates: {}", builder.num_gates());
556        let data = builder.build::<PoseidonGoldilocksConfig>();
557
558        // set values to circuit
559        let mut pw = PartialWitness::<F>::new();
560        std::iter::zip(x_target, x).try_for_each(|(t, v)| pw.set_byte_target(t, v))?;
561
562        std::iter::zip(out_target, expected).try_for_each(|(t, v)| pw.set_byte_target(t, v))?;
563
564        let proof = data.prove(pw)?;
565        data.verify(proof)?;
566        Ok(())
567    }
568
569    #[test]
570    fn test_gf_mul() -> Result<()> {
571        // AES-GCM-128
572        test_gf_mul_opt::<4, 4, 10, 16>()?;
573        test_gf_mul_opt::<4, 4, 10, 32>()?;
574
575        // AES-GCM-256
576        test_gf_mul_opt::<8, 4, 14, 16>()?;
577
578        Ok(())
579    }
580    fn test_gf_mul_opt<const NK: usize, const NB: usize, const NR: usize, const L: usize>()
581    -> Result<()>
582    where
583        [(); 4 * (NR + 1)]:,
584        [(); NK * NB]:,
585    {
586        let x: [u8; 16] = [111; 16];
587        let y: [u8; 16] = [222; 16];
588
589        let expected = gf_2_128_mul(x, y);
590        let x_bits: [[bool; 8]; 16] = array::from_fn(|b| super::le_bits_from_byte(x[b]));
591        let y_bits: [[bool; 8]; 16] = array::from_fn(|b| super::le_bits_from_byte(y[b]));
592
593        // sanity check
594        let expected2 = crate::native_gcm::gf_2_128_mul_circuit_version(x_bits, y_bits);
595        let expected_bits: [[bool; 8]; 16] =
596            array::from_fn(|b| super::le_bits_from_byte(expected[b]));
597        assert_eq!(expected2, expected_bits);
598
599        // Circuit declaration
600        let config = CircuitConfig::standard_recursion_config();
601        let mut builder = CircuitBuilder::<F, D>::new(config);
602        let xor_lut_idx = byte_xor_lut(&mut builder);
603        let u8_unit_right_shift_lut_idx = u8_unit_right_shift_lut(&mut builder);
604        let u8_bitref_lut_idx = u8_bitref_lut(&mut builder);
605
606        let x_target: [ByteTarget; 16] =
607            array::from_fn(|_| builder.add_virtual_byte_target_unsafe());
608        let y_target: [ByteTarget; 16] =
609            array::from_fn(|_| builder.add_virtual_byte_target_unsafe());
610
611        let out_target = gf_2_128_mul_target(
612            &mut builder,
613            xor_lut_idx,
614            u8_unit_right_shift_lut_idx,
615            u8_bitref_lut_idx,
616            x_target,
617            y_target,
618        );
619
620        println!("gf_2_128_mul (L:{}) num_gates: {}", L, builder.num_gates());
621        let data = builder.build::<PoseidonGoldilocksConfig>();
622
623        // set values to circuit
624        let mut pw = PartialWitness::<F>::new();
625        std::iter::zip(x_target, x).try_for_each(|(t, v)| pw.set_byte_target(t, v))?;
626        std::iter::zip(y_target, y).try_for_each(|(t, v)| pw.set_byte_target(t, v))?;
627
628        std::iter::zip(out_target, expected).try_for_each(|(t, v)| pw.set_byte_target(t, v))?;
629
630        let proof = data.prove(pw)?;
631        data.verify(proof)
632    }
633
634    #[test]
635    fn test_ghash() -> Result<()> {
636        // AES-GCM-128
637        test_ghash_opt::<4, 4, 10, 16>()?;
638        test_ghash_opt::<4, 4, 10, 32>()?;
639
640        // AES-GCM-256
641        test_ghash_opt::<8, 4, 14, 16>()?;
642
643        Ok(())
644    }
645    fn test_ghash_opt<const NK: usize, const NB: usize, const NR: usize, const L: usize>()
646    -> Result<()>
647    where
648        [(); 4 * (NR + 1)]:,
649        [(); NK * NB]:,
650    {
651        let h: [u8; 16] = [222; 16];
652        let x: &[u8; L] = &[42u8; L];
653
654        let expected = ghash(h, x);
655
656        // Circuit declaration
657        let config = CircuitConfig::standard_recursion_config();
658        let mut builder = CircuitBuilder::<F, D>::new(config);
659
660        let xor_lut_idx = byte_xor_lut(&mut builder);
661        let u8_unit_right_shift_lut_idx = u8_unit_right_shift_lut(&mut builder);
662        let u8_bitref_lut_idx = u8_bitref_lut(&mut builder);
663
664        let h_target: [ByteTarget; 16] =
665            array::from_fn(|_| builder.add_virtual_byte_target_unsafe());
666        let x_target: [ByteTarget; L] =
667            array::from_fn(|_| builder.add_virtual_byte_target_unsafe());
668
669        // let ghash_out_target = ghash_target::<L>(&mut builder, h_target, &x_target);
670        let ghash_out_target = ghash_target(
671            &mut builder,
672            xor_lut_idx,
673            u8_unit_right_shift_lut_idx,
674            u8_bitref_lut_idx,
675            h_target,
676            x_target.to_vec(),
677        );
678
679        println!("ghash (L:{}) num_gates: {}", L, builder.num_gates());
680        let data = builder.build::<PoseidonGoldilocksConfig>();
681
682        // set values to circuit
683        let mut pw = PartialWitness::<F>::new();
684        std::iter::zip(h_target, h).try_for_each(|(t, v)| pw.set_byte_target(t, v))?;
685        std::iter::zip(x_target, x).try_for_each(|(t, v)| pw.set_byte_target(t, *v))?;
686
687        std::iter::zip(ghash_out_target, expected)
688            .try_for_each(|(t, v)| pw.set_byte_target(t, v))?;
689
690        let proof = data.prove(pw)?;
691        data.verify(proof)
692    }
693
694    #[test]
695    fn test_encrypt() -> Result<()> {
696        // AES-GCM-128
697        test_encrypt_opt::<4, 4, 10, 13, false>(false)?;
698        test_encrypt_opt::<4, 4, 10, 13, true>(false)?; // with tag
699        test_encrypt_opt::<4, 4, 10, 17, false>(false)?;
700
701        // AES-GCM-256
702        test_encrypt_opt::<8, 4, 14, 13, false>(false)?;
703        test_encrypt_opt::<8, 4, 14, 13, true>(false)?; // with tag
704
705        Ok(())
706    }
707    // run it with: cargo test -- --ignored
708    #[ignore]
709    #[test]
710    fn test_encrypt_report_sizes() -> Result<()> {
711        // AES-GCM-128
712        test_encrypt_opt::<4, 4, 10, 16, false>(true)?;
713        test_encrypt_opt::<4, 4, 10, 17, false>(true)?;
714        test_encrypt_opt::<4, 4, 10, 32, false>(true)?;
715        test_encrypt_opt::<4, 4, 10, 33, false>(true)?;
716        test_encrypt_opt::<4, 4, 10, 64, false>(true)?;
717        test_encrypt_opt::<4, 4, 10, 128, false>(true)?;
718        test_encrypt_opt::<4, 4, 10, 256, false>(true)?;
719        test_encrypt_opt::<4, 4, 10, 512, false>(true)?;
720        test_encrypt_opt::<4, 4, 10, 1024, false>(true)?;
721        test_encrypt_opt::<4, 4, 10, 2048, false>(true)?;
722
723        // AES-GCM-256
724        test_encrypt_opt::<8, 4, 14, 16, false>(true)?;
725        test_encrypt_opt::<8, 4, 14, 17, false>(true)?;
726        test_encrypt_opt::<8, 4, 14, 32, false>(true)?;
727        test_encrypt_opt::<8, 4, 14, 33, false>(true)?;
728        test_encrypt_opt::<8, 4, 14, 64, false>(true)?;
729        test_encrypt_opt::<8, 4, 14, 128, false>(true)?;
730        test_encrypt_opt::<8, 4, 14, 256, false>(true)?;
731        test_encrypt_opt::<8, 4, 14, 512, false>(true)?;
732        test_encrypt_opt::<8, 4, 14, 1024, false>(true)?;
733        test_encrypt_opt::<8, 4, 14, 2048, false>(true)?;
734
735        Ok(())
736    }
737    fn test_encrypt_opt<
738        const NK: usize,
739        const NB: usize,
740        const NR: usize,
741        const L: usize,
742        const TAG: bool,
743    >(
744        only_build: bool,
745    ) -> Result<()>
746    where
747        [(); 4 * (NR + 1)]:,
748        [(); NK * NB]:,
749    {
750        let key: &[u8; NK * NB] = &[42; NK * NB];
751        let nonce: &[u8; 12] = &[111; 12];
752        let pt: &[u8; L] = &[42u8; L];
753
754        let (ct, tag) = crate::native_gcm::encrypt::<NK, NB, NR>(key, nonce, pt);
755
756        // Circuit declaration
757        let config = CircuitConfig::standard_recursion_config();
758        let mut builder = CircuitBuilder::<F, D>::new(config);
759
760        let aes_targets = AesGcmTarget::<NK, NB, NR, L, TAG>::build(&mut builder);
761
762        println!(
763            "encrypt (NK:{}, NB:{}, NR:{}, L:{}, TAG:{}) num_gates: {}",
764            NK,
765            NB,
766            NR,
767            L,
768            TAG,
769            builder.num_gates()
770        );
771        let data = builder.build::<PoseidonGoldilocksConfig>();
772
773        if only_build {
774            return Ok(());
775        }
776
777        // set values to circuit
778        let mut pw = PartialWitness::<F>::new();
779        aes_targets.set_targets(&mut pw, key, nonce, pt, &ct, &tag)?;
780
781        let proof = data.prove(pw)?;
782        data.verify(proof)
783    }
784}