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