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
 63pub fn inc32_target(
 64    builder: &mut CircuitBuilder<F, D>,
 65    block: [ByteTarget; 16],
 66) -> [ByteTarget; 16] {
 67    let mut r = block;
 68
 69    let mut carry = BoolTarget::new_unsafe(builder.one());
 70    for byte_index in (12..16).rev() {
 71        for bit_index in 0..8 {
 72            let a = block[byte_index][bit_index];
 73            // builder.assert_bool(a); // TODO we can assume it's valid bool, skip check
 74
 75            let sum = xor(builder, a, carry);
 76
 77            let carry_out = builder.and(a, carry);
 78
 79            r[byte_index][bit_index] = sum;
 80            carry = carry_out;
 81        }
 82    }
 83    r
 84}
 85
 86fn xor_byte(builder: &mut CircuitBuilder<F, D>, b1: ByteTarget, b2: ByteTarget) -> ByteTarget {
 87    array::from_fn(|i| xor(builder, b1[i], b2[i]))
 88}
 89
 90fn xor_blocks(
 91    builder: &mut CircuitBuilder<F, D>,
 92    b1: [ByteTarget; 16],
 93    b2: [ByteTarget; 16],
 94) -> [ByteTarget; 16] {
 95    array::from_fn(|i| xor_byte(builder, b1[i], b2[i]))
 96}
 97
 98pub fn msb_t(
 99    builder: &mut CircuitBuilder<F, D>,
100    t_bytes: usize,
101    block: &[ByteTarget],
102) -> Vec<ByteTarget> {
103    let t = t_bytes * 8;
104
105    let full_bytes = t / 8;
106    let rem_bits = t % 8;
107
108    let mut out: Vec<ByteTarget> = Vec::new();
109
110    for i in 0..full_bytes {
111        out.push(block[i]);
112    }
113
114    let zero_bool = BoolTarget::new_unsafe(builder.zero());
115    if rem_bits != 0 {
116        let mut partial = [zero_bool; 8];
117
118        for bit in 0..8 {
119            if bit < rem_bits {
120                partial[bit] = block[full_bytes][bit];
121            } else {
122                partial[bit] = zero_bool;
123            }
124        }
125
126        out.push(partial);
127    }
128    out
129}
130
131#[cfg(test)]
132mod tests {
133    use std::array;
134
135    use anyhow::Result;
136    use plonky2::{
137        field::{goldilocks_field::GoldilocksField as F, types::Field},
138        iop::witness::{PartialWitness, WitnessWrite},
139        plonk::{
140            circuit_builder::CircuitBuilder, circuit_data::CircuitConfig,
141            config::PoseidonGoldilocksConfig,
142        },
143    };
144    use rand::RngExt;
145
146    use super::*;
147    use crate::{
148        circuit_aes::{sbox_lut, state_mix_matrix_bits, PartialWitnessByteArray},
149        native_aes::key_expansion,
150        native_gcm::gctr,
151    };
152
153    use super::D;
154
155    #[test]
156    fn test_gctr() -> Result<()> {
157        // AES-GCM-128
158        test_gctr_opt::<4, 4, 10, 13>()?;
159        test_gctr_opt::<4, 4, 10, 17>()?;
160
161        // AES-GCM-256
162        test_gctr_opt::<8, 4, 14, 13>()?;
163
164        Ok(())
165    }
166
167    fn test_gctr_opt<const NK: usize, const NB: usize, const NR: usize, const L: usize>(
168    ) -> Result<()>
169    where
170        [(); 4 * (NR + 1)]:,
171        [(); NK * NB]:,
172    {
173        let key: &[u8; NK * NB] = &[42; NK * NB];
174        let expanded_key = key_expansion::<NK, NB, NR>(key);
175        let nonce: &[u8; 12] = &[111; 12];
176        let icb: &[u8; 16] = &[222; 16];
177        let pt: &[u8; L] = &[42u8; L];
178
179        let expected = gctr(expanded_key, &icb, pt);
180
181        // Circuit declaration
182        let config = CircuitConfig::standard_recursion_config();
183        let mut builder = CircuitBuilder::<F, D>::new(config);
184
185        let key_target: [ByteTarget; NK * NB] =
186            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
187        let nonce_target: [ByteTarget; 12] =
188            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
189        let icb_target: [ByteTarget; 16] =
190            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
191        let pt_target: [ByteTarget; L] =
192            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
193        let sbox_lut = sbox_lut(&mut builder);
194        let mix_matrix = state_mix_matrix_bits(&mut builder);
195
196        let key_expanded_target = builder.key_expansion::<NK, NB, NR>(sbox_lut, key_target);
197        let gctr_out_target = gctr_target::<NR, L>(
198            &mut builder,
199            sbox_lut,
200            mix_matrix,
201            key_expanded_target,
202            &icb_target,
203            &pt_target,
204        );
205
206        println!(
207            "encrypt_block (NK:{}, NB:{}, NR:{}, L:{}) num_gates: {}",
208            NK,
209            NB,
210            NR,
211            L,
212            builder.num_gates()
213        );
214        let data = builder.build::<PoseidonGoldilocksConfig>();
215
216        // set values to circuit
217        let mut pw = PartialWitness::<F>::new();
218        std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?;
219        std::iter::zip(nonce_target, nonce)
220            .try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?;
221        std::iter::zip(icb_target, icb).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?;
222        std::iter::zip(pt_target, pt).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?;
223
224        std::iter::zip(gctr_out_target, expected)
225            .try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
226
227        let proof = data.prove(pw)?;
228        data.verify(proof)
229    }
230}