1//! For LICENSE check out https://github.com/0xPARC/plonky2-crypto-gadgets/blob/main/LICENSE
  2
  3use std::array;
  4
  5use anyhow::Result;
  6use plonky2::{
  7    field::{
  8        extension::{FieldExtension, quintic::QuinticExtension},
  9        goldilocks_field::GoldilocksField as F,
 10        types::Field,
 11    },
 12    hash::poseidon::PoseidonHash,
 13    iop::{
 14        target::Target,
 15        witness::{PartialWitness, WitnessWrite},
 16    },
 17    plonk::circuit_builder::CircuitBuilder,
 18};
 19use pod2::backends::plonky2::{
 20    basetypes::D,
 21    primitives::ec::{
 22        curve::{CircuitBuilderElliptic, Point, PointTarget, WitnessWriteCurve},
 23        field::{CircuitBuilderNNF, OEFTarget},
 24    },
 25};
 26
 27use crate::{Fq, two_128};
 28
 29// Fq Target
 30type FqT = OEFTarget<5, QuinticExtension<F>>;
 31
 32pub struct PoseidonEncryptTarget<const L: usize>
 33where
 34    [(); L + 1]:,
 35{
 36    ks: PointTarget,
 37    m: [FqT; L],
 38    nonce: Target,
 39    ct: [FqT; L + 1],
 40}
 41
 42impl<const L: usize> PoseidonEncryptTarget<L>
 43where
 44    [(); L + 1]:,
 45{
 46    pub fn build(builder: &mut CircuitBuilder<F, D>) -> Self {
 47        assert!(L.is_multiple_of(3));
 48        let fqt_zero = const_fqt_zero(builder);
 49
 50        // add targets
 51        let ks: PointTarget = builder.add_virtual_point_target();
 52        let m: [FqT; L] = array::from_fn(|_| builder.add_virtual_nnf_target());
 53        let nonce: Target = builder.add_virtual_target();
 54        let mut ct: [FqT; L + 1] = array::from_fn(|_| fqt_zero.clone());
 55
 56        // build the circuit logic
 57        let f_zero = builder.constant(F::ZERO);
 58        let n = FqT::new([nonce, f_zero, f_zero, f_zero, f_zero]);
 59        let two128: Fq = two_128();
 60        let l_two128: Fq = Fq::from_basefield_array([
 61            F::from_canonical_u64(L as u64),
 62            F::ZERO,
 63            F::ZERO,
 64            F::ZERO,
 65            F::ZERO,
 66        ]) * two128;
 67        let l_two128_target: FqT = const_fqt_from_fq(builder, l_two128);
 68        let nl128: FqT = builder.nnf_add(&n, &l_two128_target);
 69        let fq_zero: FqT = FqT::new([f_zero, f_zero, f_zero, f_zero, f_zero]);
 70        let mut s: [FqT; 4] = [fq_zero, ks.x.clone(), ks.u.clone(), nl128];
 71
 72        for i in 0..L / 3 {
 73            s = hash_state_target(builder, s);
 74
 75            // absorb
 76            s[1] = builder.nnf_add(&s[1], &m[i * 3]);
 77            s[2] = builder.nnf_add(&s[2], &m[i * 3 + 1]);
 78            s[3] = builder.nnf_add(&s[3], &m[i * 3 + 2]);
 79
 80            // release
 81            ct[i * 3] = s[1].clone();
 82            ct[i * 3 + 1] = s[2].clone();
 83            ct[i * 3 + 2] = s[3].clone();
 84        }
 85        s = hash_state_target(builder, s);
 86        ct[L] = s[1].clone();
 87
 88        Self { ks, m, nonce, ct }
 89    }
 90
 91    pub fn set_targets(
 92        &self,
 93        pw: &mut PartialWitness<F>,
 94        ks: Point,
 95        m: &[Fq],
 96        nonce: F,
 97        ct: &[Fq],
 98    ) -> Result<()> {
 99        assert_eq!(m.len(), L);
100        assert_eq!(m.len() + 1, ct.len());
101
102        pw.set_point_target(&self.ks, &ks)?;
103        #[allow(clippy::needless_range_loop)]
104        for i in 0..L {
105            for j in 0..5 {
106                pw.set_target(self.m[i].components[j], m[i].0[j])?;
107            }
108        }
109        pw.set_target(self.nonce, nonce)?;
110        #[allow(clippy::needless_range_loop)]
111        for i in 0..L + 1 {
112            for j in 0..5 {
113                pw.set_target(self.ct[i].components[j], ct[i].0[j])?;
114            }
115        }
116
117        Ok(())
118    }
119}
120
121fn hash_state_target(builder: &mut CircuitBuilder<F, D>, s: [FqT; 4]) -> [FqT; 4] {
122    let elems: [Target; 4 * 5] = array::from_fn(|i| s[i / 5].components[i % 5]);
123    let h = builder
124        .hash_n_to_hash_no_pad::<PoseidonHash>(elems.to_vec())
125        .elements;
126    let mut h5: [Target; 5] = [builder.zero(); 5];
127    h5[..4].copy_from_slice(&h);
128    let fqt_zero = const_fqt_zero(builder);
129    [
130        FqT::new(h5),
131        fqt_zero.clone(),
132        fqt_zero.clone(),
133        fqt_zero.clone(),
134    ]
135}
136
137fn const_fqt_from_fq(builder: &mut CircuitBuilder<F, D>, fq: Fq) -> FqT {
138    let t: [Target; 5] = array::from_fn(|i| builder.constant(fq.0[i]));
139    FqT::new(t)
140}
141fn const_fqt_from_f(builder: &mut CircuitBuilder<F, D>, f: F) -> FqT {
142    let v = builder.constant(f);
143    let zero = builder.constant(F::ZERO);
144    FqT::new([v, zero, zero, zero, zero])
145}
146fn const_fqt_zero(builder: &mut CircuitBuilder<F, D>) -> FqT {
147    let zero = builder.constant(F::ZERO);
148    FqT::new([zero, zero, zero, zero, zero])
149}
150
151#[cfg(test)]
152mod tests {
153    use anyhow::Result;
154    use plonky2::{
155        field::{goldilocks_field::GoldilocksField as F, types::Sample},
156        iop::witness::PartialWitness,
157        plonk::{
158            circuit_builder::CircuitBuilder, circuit_data::CircuitConfig,
159            config::PoseidonGoldilocksConfig,
160        },
161    };
162
163    use super::*;
164
165    #[test]
166    fn test_encrypt() -> Result<()> {
167        let mut rng = rand::thread_rng();
168        const L: usize = 129; // (multiple of 3)
169
170        let (_k, K) = crate::new_key();
171        let ks = crate::expanded_key(K);
172        let nonce = F::rand();
173        let msg: Vec<Fq> = (0..L).map(|_| Fq::sample(&mut rng)).collect();
174        let l = msg.len();
175
176        let ct = crate::encrypt(ks, &msg, nonce);
177
178        let m = crate::decrypt(ks, &ct, nonce, l);
179        assert_eq!(m, msg);
180
181        // circuit declaration
182        let config = CircuitConfig::standard_recursion_config();
183        let mut builder = CircuitBuilder::<F, D>::new(config);
184
185        let cipher_targets = PoseidonEncryptTarget::<L>::build(&mut builder);
186
187        println!(
188            "PoseidonEncrypt circuit (L:{}) num_gates: {}",
189            L,
190            builder.num_gates()
191        );
192        let data = builder.build::<PoseidonGoldilocksConfig>();
193
194        // set values to circuit
195        let mut pw = PartialWitness::<F>::new();
196        cipher_targets.set_targets(&mut pw, ks, &msg, nonce, &ct)?;
197
198        let proof = data.prove(pw)?;
199        data.verify(proof)
200    }
201}