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    circuit_aes::{xor, ByteArrayTarget as ByteTarget, CircuitBuilderAESState, StateTarget},
 17    D,
 18};
 19
 20/// L: max size of plaintext to cipher.
 21fn gctr_target<const NR: usize, const L: usize>(
 22    builder: &mut CircuitBuilder<F, D>,
 23    sbox_lut_idx: usize,
 24    mix_matrix: [[ByteTarget; 4]; 4],
 25    key: [[ByteTarget; 4]; 4 * (NR + 1)],
 26    icb: &[ByteTarget; 16],
 27    x: &[ByteTarget; L],
 28) -> [ByteTarget; L] {
 29    let n = ((L * 8) as f64 / 128_f64).ceil() as usize;
 30
 31    let mut y: [ByteTarget; L] = x.clone();
 32    let mut cb_i = *icb;
 33
 34    let zero_byte = builder.zero_byte();
 35    for (i, x_i_raw) in x.chunks(16).enumerate() {
 36        if i > 0 {
 37            cb_i = inc32_target(builder, cb_i);
 38        }
 39
 40        let l = x_i_raw.len().min(16);
 41        let mut x_i = [zero_byte; 16];
 42        x_i[..l].copy_from_slice(x_i_raw);
 43
 44        let ciph_cb_i = builder
 45            .encrypt_block::<NR>(sbox_lut_idx, mix_matrix, StateTarget::from_flat(cb_i), key)
 46            .flatten();
 47
 48        let (y_i, n_bytes) = if i < n && x_i_raw.len() == 16 {
 49            (xor_blocks(builder, x_i, ciph_cb_i), 16)
 50        } else {
 51            // last chunk, might be smaller than 16 bytes (L%16 bytes)
 52            let msb_res = msb_t(builder, L % 16, ciph_cb_i.as_slice());
 53            let l = msb_res.len().min(16);
 54            let mut m: [ByteTarget; 16] = [zero_byte; 16];
 55            m[..l].copy_from_slice(&msb_res);
 56            (xor_blocks(builder, x_i, m), l)
 57        };
 58        y[i * 16..i * 16 + n_bytes].clone_from_slice(&y_i[..n_bytes]);
 59    }
 60    y
 61}
 62
 63fn ghash_target<const L: usize>(
 64    builder: &mut CircuitBuilder<F, D>,
 65    h: [ByteTarget; 16],
 66    x: &[ByteTarget; L],
 67) -> [ByteTarget; 16] {
 68    assert!(L.is_multiple_of(16)); // multiple of 128 bits
 69    let m = L / 16;
 70
 71    let zero_byte = builder.zero_byte();
 72    let mut y = [zero_byte; 16]; // (128 bits)
 73    for i in 0..m {
 74        let mut xi = [zero_byte; 16];
 75        xi.clone_from_slice(&x[i * 16..i * 16 + 16]);
 76        let y_xi = xor_blocks(builder, y, xi);
 77        y = gf_2_128_mul_target(builder, y_xi, h);
 78    }
 79    y
 80}
 81fn gf_2_128_mul_target(
 82    builder: &mut CircuitBuilder<F, D>,
 83    x: [ByteTarget; 16],
 84    y: [ByteTarget; 16],
 85) -> [ByteTarget; 16] {
 86    // WIP TODO
 87
 88    let zero_byte = builder.zero_byte();
 89    let one_bool = builder._true();
 90
 91    // R: 11100001 || 0^120
 92    let mut r = [zero_byte; 16];
 93    r[0][0] = one_bool;
 94    r[0][1] = one_bool;
 95    r[0][2] = one_bool;
 96    r[0][7] = one_bool;
 97
 98    let mut z = [zero_byte; 16];
 99    let mut v = y;
100    for i in 0..128 {
101        let byte_index = i / 8;
102        let bit_index = 7 - (i % 8);
103        // dbg!(byte_index, bit_index);
104        let xi = x[byte_index][bit_index];
105
106        // set z = if xi==1: z^v, else: z
107        for b in 0..16 {
108            for k in 0..8 {
109                let z_xor_v = xor(builder, z[b][k], v[b][k]);
110                z[b][k] =
111                    BoolTarget::new_unsafe(builder.select(xi, z_xor_v.target, z[b][k].target));
112            }
113        }
114
115        let lsb = v[15][7].clone();
116        v = right_shift_one_target(builder, &v);
117
118        // if lsb==1: v=v^R, else: v
119        for b in 0..16 {
120            for k in 0..8 {
121                let v_xor_r = xor(builder, v[b][k], r[b][k]);
122                v[b][k] =
123                    BoolTarget::new_unsafe(builder.select(lsb, v_xor_r.target, v[b][k].target));
124            }
125        }
126    }
127    z
128}
129pub fn right_shift_one_target(
130    builder: &mut CircuitBuilder<F, D>,
131    v: &[ByteTarget; 16],
132) -> [ByteTarget; 16] {
133    let mut r: [ByteTarget; 16] = v.clone();
134
135    let mut carry = builder._false();
136    for i in 0..16 {
137        let current = v[i];
138
139        let next_carry = current[0];
140
141        let mut shifted = [builder._false(); 8];
142        for j in 0..7 {
143            shifted[j] = current[j + 1];
144        }
145        shifted[7] = carry;
146        r[i] = shifted;
147        carry = next_carry;
148    }
149    r
150}
151
152pub fn inc32_target(
153    builder: &mut CircuitBuilder<F, D>,
154    block: [ByteTarget; 16],
155) -> [ByteTarget; 16] {
156    let mut r = block;
157
158    let mut carry = BoolTarget::new_unsafe(builder.one());
159    for byte_index in (12..16).rev() {
160        for bit_index in 0..8 {
161            let a = block[byte_index][bit_index];
162            // builder.assert_bool(a); // TODO we can assume it's valid bool, skip check
163
164            let sum = xor(builder, a, carry);
165            let carry_out = builder.and(a, carry);
166            r[byte_index][bit_index] = sum;
167            carry = carry_out;
168        }
169    }
170    r
171}
172
173fn xor_byte(builder: &mut CircuitBuilder<F, D>, b1: ByteTarget, b2: ByteTarget) -> ByteTarget {
174    array::from_fn(|i| xor(builder, b1[i], b2[i]))
175}
176
177fn xor_blocks(
178    builder: &mut CircuitBuilder<F, D>,
179    b1: [ByteTarget; 16],
180    b2: [ByteTarget; 16],
181) -> [ByteTarget; 16] {
182    array::from_fn(|i| xor_byte(builder, b1[i], b2[i]))
183}
184
185pub fn msb_t(
186    builder: &mut CircuitBuilder<F, D>,
187    t_bytes: usize,
188    block: &[ByteTarget],
189) -> Vec<ByteTarget> {
190    let t = t_bytes * 8;
191
192    let full_bytes = t / 8;
193    let rem_bits = t % 8;
194
195    let mut out: Vec<ByteTarget> = Vec::new();
196    for i in 0..full_bytes {
197        out.push(block[i]);
198    }
199
200    let zero_bool = BoolTarget::new_unsafe(builder.zero());
201    if rem_bits != 0 {
202        let mut partial = [zero_bool; 8];
203
204        for bit in 0..8 {
205            if bit < rem_bits {
206                partial[bit] = block[full_bytes][bit];
207            } else {
208                partial[bit] = zero_bool;
209            }
210        }
211        out.push(partial);
212    }
213    out
214}
215
216#[cfg(test)]
217mod tests {
218    use std::array;
219
220    use anyhow::Result;
221    use plonky2::{
222        field::{goldilocks_field::GoldilocksField as F, types::Field},
223        iop::witness::{PartialWitness, WitnessWrite},
224        plonk::{
225            circuit_builder::CircuitBuilder, circuit_data::CircuitConfig,
226            config::PoseidonGoldilocksConfig,
227        },
228    };
229    use rand::RngExt;
230
231    use super::*;
232    use crate::{
233        circuit_aes::{sbox_lut, state_mix_matrix_bits, PartialWitnessByteArray},
234        native_aes::key_expansion,
235        native_gcm::{gctr, gf_2_128_mul, ghash, right_shift_one},
236    };
237
238    use super::D;
239
240    #[test]
241    fn test_gctr() -> Result<()> {
242        // AES-GCM-128
243        test_gctr_opt::<4, 4, 10, 13>()?;
244        test_gctr_opt::<4, 4, 10, 17>()?;
245
246        // AES-GCM-256
247        test_gctr_opt::<8, 4, 14, 13>()?;
248
249        Ok(())
250    }
251    fn test_gctr_opt<const NK: usize, const NB: usize, const NR: usize, const L: usize>(
252    ) -> Result<()>
253    where
254        [(); 4 * (NR + 1)]:,
255        [(); NK * NB]:,
256    {
257        let key: &[u8; NK * NB] = &[42; NK * NB];
258        let expanded_key = key_expansion::<NK, NB, NR>(key);
259        let nonce: &[u8; 12] = &[111; 12];
260        let icb: &[u8; 16] = &[222; 16];
261        let pt: &[u8; L] = &[42u8; L];
262
263        let expected = gctr(expanded_key, &icb, pt);
264
265        // Circuit declaration
266        let config = CircuitConfig::standard_recursion_config();
267        let mut builder = CircuitBuilder::<F, D>::new(config);
268
269        let key_target: [ByteTarget; NK * NB] =
270            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
271        let nonce_target: [ByteTarget; 12] =
272            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
273        let icb_target: [ByteTarget; 16] =
274            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
275        let pt_target: [ByteTarget; L] =
276            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
277        let sbox_lut = sbox_lut(&mut builder);
278        let mix_matrix = state_mix_matrix_bits(&mut builder);
279
280        let key_expanded_target = builder.key_expansion::<NK, NB, NR>(sbox_lut, key_target);
281        let gctr_out_target = gctr_target::<NR, L>(
282            &mut builder,
283            sbox_lut,
284            mix_matrix,
285            key_expanded_target,
286            &icb_target,
287            &pt_target,
288        );
289
290        println!(
291            "gctr (NK:{}, NB:{}, NR:{}, L:{}) num_gates: {}",
292            NK,
293            NB,
294            NR,
295            L,
296            builder.num_gates()
297        );
298        let data = builder.build::<PoseidonGoldilocksConfig>();
299
300        // set values to circuit
301        let mut pw = PartialWitness::<F>::new();
302        std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?;
303        std::iter::zip(nonce_target, nonce)
304            .try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?;
305        std::iter::zip(icb_target, icb).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?;
306        std::iter::zip(pt_target, pt).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?;
307
308        std::iter::zip(gctr_out_target, expected)
309            .try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
310
311        let proof = data.prove(pw)?;
312        data.verify(proof)
313    }
314
315    #[test]
316    fn test_right_shift_one() -> Result<()> {
317        let x: [u8; 16] = [111; 16];
318
319        let mut expected = x.clone();
320        right_shift_one(&mut expected);
321
322        // Circuit declaration
323        let config = CircuitConfig::standard_recursion_config();
324        let mut builder = CircuitBuilder::<F, D>::new(config);
325
326        let x_target: [ByteTarget; 16] =
327            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
328
329        let out_target: [ByteTarget; 16] = right_shift_one_target(&mut builder, &x_target);
330
331        println!("right_shift_one num_gates: {}", builder.num_gates());
332        let data = builder.build::<PoseidonGoldilocksConfig>();
333
334        // set values to circuit
335        let mut pw = PartialWitness::<F>::new();
336        std::iter::zip(x_target, x).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
337
338        std::iter::zip(out_target, expected)
339            .try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
340
341        let proof = data.prove(pw)?;
342        data.verify(proof)?;
343        Ok(())
344    }
345
346    #[test]
347    fn test_gf_mul() -> Result<()> {
348        // AES-GCM-128
349        test_gf_mul_opt::<4, 4, 10, 16>()?;
350        // test_gf_mul_opt::<4, 4, 10, 32>()?;
351        //
352        // // AES-GCM-256
353        // test_gf_mul_opt::<8, 4, 14, 16>()?;
354
355        Ok(())
356    }
357    fn test_gf_mul_opt<const NK: usize, const NB: usize, const NR: usize, const L: usize>(
358    ) -> Result<()>
359    where
360        [(); 4 * (NR + 1)]:,
361        [(); NK * NB]:,
362    {
363        let x: [u8; 16] = [111; 16];
364        let y: [u8; 16] = [222; 16];
365
366        let expected = gf_2_128_mul(x, y);
367
368        // Circuit declaration
369        let config = CircuitConfig::standard_recursion_config();
370        let mut builder = CircuitBuilder::<F, D>::new(config);
371
372        let x_target: [ByteTarget; 16] =
373            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
374        let y_target: [ByteTarget; 16] =
375            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
376
377        let out_target = gf_2_128_mul_target(&mut builder, x_target, y_target);
378
379        println!("gf_2_128_mul (L:{}) num_gates: {}", L, builder.num_gates());
380        let data = builder.build::<PoseidonGoldilocksConfig>();
381
382        // set values to circuit
383        let mut pw = PartialWitness::<F>::new();
384        std::iter::zip(x_target, x).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
385        std::iter::zip(y_target, y).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
386
387        std::iter::zip(out_target, expected)
388            .try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
389
390        let proof = data.prove(pw)?;
391        data.verify(proof)
392    }
393
394    #[test]
395    fn test_ghash() -> Result<()> {
396        // AES-GCM-128
397        test_ghash_opt::<4, 4, 10, 16>()?;
398        test_ghash_opt::<4, 4, 10, 32>()?;
399
400        // AES-GCM-256
401        test_ghash_opt::<8, 4, 14, 16>()?;
402
403        Ok(())
404    }
405    fn test_ghash_opt<const NK: usize, const NB: usize, const NR: usize, const L: usize>(
406    ) -> Result<()>
407    where
408        [(); 4 * (NR + 1)]:,
409        [(); NK * NB]:,
410    {
411        let h: [u8; 16] = [222; 16];
412        let x: &[u8; L] = &[42u8; L];
413
414        let expected = ghash(h, x);
415
416        // Circuit declaration
417        let config = CircuitConfig::standard_recursion_config();
418        let mut builder = CircuitBuilder::<F, D>::new(config);
419
420        let h_target: [ByteTarget; 16] =
421            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
422        let x_target: [ByteTarget; L] =
423            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
424
425        let ghash_out_target = ghash_target::<L>(&mut builder, h_target, &x_target);
426
427        println!("ghash (L:{}) num_gates: {}", L, builder.num_gates());
428        let data = builder.build::<PoseidonGoldilocksConfig>();
429
430        // set values to circuit
431        let mut pw = PartialWitness::<F>::new();
432        std::iter::zip(h_target, h).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
433        std::iter::zip(x_target, x).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?;
434
435        std::iter::zip(ghash_out_target, expected)
436            .try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
437
438        let proof = data.prove(pw)?;
439        data.verify(proof)
440    }
441}