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