A poseidon-cipher/src/circuit.rs

@@ -0,0 +1,201 @@
+//! For LICENSE check out https://github.com/0xPARC/plonky2-crypto-gadgets/blob/main/LICENSE
+
+use std::array;
+
+use anyhow::Result;
+use plonky2::{
+    field::{
+        extension::{FieldExtension, quintic::QuinticExtension},
+        goldilocks_field::GoldilocksField as F,
+        types::Field,
+    },
+    hash::poseidon::PoseidonHash,
+    iop::{
+        target::Target,
+        witness::{PartialWitness, WitnessWrite},
+    },
+    plonk::circuit_builder::CircuitBuilder,
+};
+use pod2::backends::plonky2::{
+    basetypes::D,
+    primitives::ec::{
+        curve::{CircuitBuilderElliptic, Point, PointTarget, WitnessWriteCurve},
+        field::{CircuitBuilderNNF, OEFTarget},
+    },
+};
+
+use crate::{Fq, two_128};
+
+// Fq Target
+type FqT = OEFTarget<5, QuinticExtension<F>>;
+
+pub struct PoseidonEncryptTarget<const L: usize>
+where
+    [(); L + 1]:,
+{
+    ks: PointTarget,
+    m: [FqT; L],
+    nonce: Target,
+    ct: [FqT; L + 1],
+}
+
+impl<const L: usize> PoseidonEncryptTarget<L>
+where
+    [(); L + 1]:,
+{
+    pub fn build(builder: &mut CircuitBuilder<F, D>) -> Self {
+        assert!(L.is_multiple_of(3));
+        let fqt_zero = const_fqt_zero(builder);
+
+        // add targets
+        let ks: PointTarget = builder.add_virtual_point_target();
+        let m: [FqT; L] = array::from_fn(|_| builder.add_virtual_nnf_target());
+        let nonce: Target = builder.add_virtual_target();
+        let mut ct: [FqT; L + 1] = array::from_fn(|_| fqt_zero.clone());
+
+        // build the circuit logic
+        let f_zero = builder.constant(F::ZERO);
+        let n = FqT::new([nonce, f_zero, f_zero, f_zero, f_zero]);
+        let two128: Fq = two_128();
+        let l_two128: Fq = Fq::from_basefield_array([
+            F::from_canonical_u64(L as u64),
+            F::ZERO,
+            F::ZERO,
+            F::ZERO,
+            F::ZERO,
+        ]) * two128;
+        let l_two128_target: FqT = const_fqt_from_fq(builder, l_two128);
+        let nl128: FqT = builder.nnf_add(&n, &l_two128_target);
+        let fq_zero: FqT = FqT::new([f_zero, f_zero, f_zero, f_zero, f_zero]);
+        let mut s: [FqT; 4] = [fq_zero, ks.x.clone(), ks.u.clone(), nl128];
+
+        for i in 0..L / 3 {
+            s = hash_state_target(builder, s);
+
+            // absorb
+            s[1] = builder.nnf_add(&s[1], &m[i * 3]);
+            s[2] = builder.nnf_add(&s[2], &m[i * 3 + 1]);
+            s[3] = builder.nnf_add(&s[3], &m[i * 3 + 2]);
+
+            // release
+            ct[i * 3] = s[1].clone();
+            ct[i * 3 + 1] = s[2].clone();
+            ct[i * 3 + 2] = s[3].clone();
+        }
+        s = hash_state_target(builder, s);
+        ct[L] = s[1].clone();
+
+        Self { ks, m, nonce, ct }
+    }
+
+    pub fn set_targets(
+        &self,
+        pw: &mut PartialWitness<F>,
+        ks: Point,
+        m: &[Fq],
+        nonce: F,
+        ct: &[Fq],
+    ) -> Result<()> {
+        assert_eq!(m.len(), L);
+        assert_eq!(m.len() + 1, ct.len());
+
+        pw.set_point_target(&self.ks, &ks)?;
+        #[allow(clippy::needless_range_loop)]
+        for i in 0..L {
+            for j in 0..5 {
+                pw.set_target(self.m[i].components[j], m[i].0[j])?;
+            }
+        }
+        pw.set_target(self.nonce, nonce)?;
+        #[allow(clippy::needless_range_loop)]
+        for i in 0..L + 1 {
+            for j in 0..5 {
+                pw.set_target(self.ct[i].components[j], ct[i].0[j])?;
+            }
+        }
+
+        Ok(())
+    }
+}
+
+fn hash_state_target(builder: &mut CircuitBuilder<F, D>, s: [FqT; 4]) -> [FqT; 4] {
+    let elems: [Target; 4 * 5] = array::from_fn(|i| s[i / 5].components[i % 5]);
+    let h = builder
+        .hash_n_to_hash_no_pad::<PoseidonHash>(elems.to_vec())
+        .elements;
+    let mut h5: [Target; 5] = [builder.zero(); 5];
+    h5[..4].copy_from_slice(&h);
+    let fqt_zero = const_fqt_zero(builder);
+    [
+        FqT::new(h5),
+        fqt_zero.clone(),
+        fqt_zero.clone(),
+        fqt_zero.clone(),
+    ]
+}
+
+fn const_fqt_from_fq(builder: &mut CircuitBuilder<F, D>, fq: Fq) -> FqT {
+    let t: [Target; 5] = array::from_fn(|i| builder.constant(fq.0[i]));
+    FqT::new(t)
+}
+fn const_fqt_from_f(builder: &mut CircuitBuilder<F, D>, f: F) -> FqT {
+    let v = builder.constant(f);
+    let zero = builder.constant(F::ZERO);
+    FqT::new([v, zero, zero, zero, zero])
+}
+fn const_fqt_zero(builder: &mut CircuitBuilder<F, D>) -> FqT {
+    let zero = builder.constant(F::ZERO);
+    FqT::new([zero, zero, zero, zero, zero])
+}
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+    use plonky2::{
+        field::{goldilocks_field::GoldilocksField as F, types::Sample},
+        iop::witness::PartialWitness,
+        plonk::{
+            circuit_builder::CircuitBuilder, circuit_data::CircuitConfig,
+            config::PoseidonGoldilocksConfig,
+        },
+    };
+
+    use super::*;
+
+    #[test]
+    fn test_encrypt() -> Result<()> {
+        let mut rng = rand::thread_rng();
+        const L: usize = 129; // (multiple of 3)
+
+        let (_k, K) = crate::new_key();
+        let ks = crate::expanded_key(K);
+        let nonce = F::rand();
+        let msg: Vec<Fq> = (0..L).map(|_| Fq::sample(&mut rng)).collect();
+        let l = msg.len();
+
+        let ct = crate::encrypt(ks, &msg, nonce);
+
+        let m = crate::decrypt(ks, &ct, nonce, l);
+        assert_eq!(m, msg);
+
+        // circuit declaration
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::new(config);
+
+        let cipher_targets = PoseidonEncryptTarget::<L>::build(&mut builder);
+
+        println!(
+            "PoseidonEncrypt circuit (L:{}) num_gates: {}",
+            L,
+            builder.num_gates()
+        );
+        let data = builder.build::<PoseidonGoldilocksConfig>();
+
+        // set values to circuit
+        let mut pw = PartialWitness::<F>::new();
+        cipher_targets.set_targets(&mut pw, ks, &msg, nonce, &ct)?;
+
+        let proof = data.prove(pw)?;
+        data.verify(proof)
+    }
+}

M poseidon-cipher/src/lib.rs → poseidon-cipher/src/lib.rs

@@ -1,104 +1,96 @@
-//! For LICENSE check out https://github.com/0xPARC/plonky2-aes/blob/main/LICENSE
+//! For LICENSE check out https://github.com/0xPARC/plonky2-crypto-gadgets/blob/main/LICENSE
+
 #![allow(incomplete_features)]
 #![feature(generic_const_exprs)]
 #![allow(dead_code)] // TMP
 #![allow(non_snake_case)]
 
+use std::array;
+
 use num::bigint::BigUint;
 use num_bigint::RandBigInt;
 use plonky2::{
     field::{
-        extension::{quintic::QuinticExtension, FieldExtension},
+        extension::{FieldExtension, quintic::QuinticExtension},
         goldilocks_field::GoldilocksField as F,
-        types::{Field, Field64},
+        types::Field,
     },
     hash::{hashing::hash_n_to_hash_no_pad, poseidon::PoseidonPermutation},
 };
-use pod2::backends::plonky2::primitives::ec::curve::{Point, GROUP_ORDER};
-use rand::{rngs::OsRng, Rng};
+use pod2::backends::plonky2::primitives::ec::curve::{GROUP_ORDER, Point};
+use rand::rngs::OsRng;
+
+pub mod circuit;
 
 // ECField
 type Fq = QuinticExtension<F>;
 
 pub fn new_key() -> (BigUint, Point) {
     let k: BigUint = OsRng.gen_biguint_below(&GROUP_ORDER);
-    // let k: Fq = QuinticExtension(std::array::from_fn(|_| {
-    //     F::from_canonical_u64(OsRng.gen_range(0..F::ORDER))
-    // }));
-    // let k = F::from_canonical_u64(OsRng.gen_range(0..F::ORDER));
     let K = &k * Point::generator();
     (k, K)
 }
 pub fn expanded_key(K: Point) -> Point {
     let r: BigUint = OsRng.gen_biguint_below(&GROUP_ORDER);
-    // let r: Fq = QuinticExtension(std::array::from_fn(|_| {
-    //     F::from_canonical_u64(OsRng.gen_range(0..F::ORDER))
-    // }));
-    &r * K // TODO maybe separate into x, u (ECField)
+    &r * K
+}
+
+pub(crate) fn two_128() -> Fq {
+    // WIP TODO
+    let two32 = Fq::from(F::from_canonical_u64(2_u64.pow(32)));
+    // dbg!(2_u64.pow(32));
+    // dbg!(F::from_canonical_u64(2_u64.pow(32)));
+    // dbg!(&two32);
+    let two64: Fq = two32 * two32;
+    // dbg!(&two64);
+    let two128: Fq = two64 * two64;
+    // dbg!(&two128);
+    two128
 }
 
-pub fn encrypt(ks: Point, msg: Vec<Fq>, nonce: F) -> Vec<Fq> {
+pub fn encrypt(ks: Point, msg: &[Fq], nonce: F) -> Vec<Fq> {
     // pad m
-    dbg!(&msg.len());
-    let mut m = msg.clone();
+    let mut m = msg.to_vec();
     m.resize(m.len().next_multiple_of(3), Fq::ZERO);
-    dbg!(&m.len());
 
     let l = Fq::from(F::from_canonical_u64(msg.len() as u64));
     let n = Fq::from(nonce);
-    // WIP TODO
-    let two32 = Fq::from(F::from_canonical_u64(2_u64.pow(32)));
-    let two64: Fq = two32 * two32;
-    let two128: Fq = two64 * two64;
-    let nl128: Fq = n + l * two128; // TODO
+    let nl128: Fq = n + l * two_128();
     let mut s: [Fq; 4] = [Fq::ZERO, ks.x, ks.u, nl128];
 
-    let mut ct: Vec<Fq> = vec![];
+    let mut ct: Vec<Fq> = vec![Fq::ZERO; m.len() + 1];
     for i in 0..m.len() / 3 {
         s = hash_state(s);
 
         // absorb
-        s[1] = s[1] + m[i * 3];
-        s[2] = s[2] + m[i * 3 + 1];
-        s[3] = s[3] + m[i * 3 + 2];
+        s[1] += m[i * 3];
+        s[2] += m[i * 3 + 1];
+        s[3] += m[i * 3 + 2];
 
         // release
-        ct.push(s[1]);
-        ct.push(s[2]);
-        ct.push(s[3]);
+        ct[i * 3] = s[1];
+        ct[i * 3 + 1] = s[2];
+        ct[i * 3 + 2] = s[3];
     }
     s = hash_state(s);
-    ct.push(s[1]);
+    ct[m.len()] = s[1];
     ct
 }
 
-pub fn decrypt(ks: Point, ct: Vec<Fq>, nonce: F, l: usize) -> Vec<Fq> {
+pub fn decrypt(ks: Point, ct: &[Fq], nonce: F, l: usize) -> Vec<Fq> {
     let l_fq = Fq::from(F::from_canonical_u64(l as u64));
     let n = Fq::from(nonce);
-    // WIP TODO
-    let two32 = Fq::from(F::from_canonical_u64(2_u64.pow(32)));
-    dbg!(2_u64.pow(32));
-    dbg!(F::from_canonical_u64(2_u64.pow(32)));
-    dbg!(&two32);
-    let two64: Fq = two32 * two32;
-    dbg!(&two64);
-    let two128: Fq = two64 * two64;
-    dbg!(&two128);
-    let nl128: Fq = n + l_fq * two128; // TODO
+    let nl128: Fq = n + l_fq * two_128();
     let mut s: [Fq; 4] = [Fq::ZERO, ks.x, ks.u, nl128];
 
-    // let mut m: Vec<Fq> = vec![Fq::ZERO; ct.len() / 3];
-    let mut m: Vec<Fq> = vec![];
+    let mut m: Vec<Fq> = vec![Fq::ZERO; ct.len() - 1];
     for i in 0..ct.len() / 3 {
         s = hash_state(s);
 
         // release
-        // m[3 * i] = s[1] + ct[3 * i];
-        // m[3 * i + 1] = s[2] + ct[3 * i + 1];
-        // m[3 * i + 2] = s[3] + ct[3 * i + 2];
-        m.push(s[1] + ct[3 * i]);
-        m.push(s[2] + ct[3 * i + 1]);
-        m.push(s[3] + ct[3 * i + 2]);
+        m[3 * i] = s[1] + ct[3 * i];
+        m[3 * i + 1] = s[2] + ct[3 * i + 1];
+        m[3 * i + 2] = s[3] + ct[3 * i + 2];
 
         // modify state
         s[1] = ct[3 * i];
@@ -120,17 +112,7 @@ m[0..l].to_vec()
 }
 
 fn hash_state(s: [Fq; 4]) -> [Fq; 4] {
-    let elems: Vec<F> = s
-        .iter()
-        .map(|e| {
-            {
-                let aux: [F; 5] = e.to_basefield_array();
-                aux
-            }
-            .to_vec()
-        })
-        .collect::<Vec<Vec<F>>>()
-        .concat();
+    let elems: [F; 4 * 5] = array::from_fn(|i| s[i / 5].0[i % 5]);
     let h = hash_n_to_hash_no_pad::<F, PoseidonPermutation<_>>(&elems).elements;
     let mut h5: [F; 5] = [F::ZERO; 5];
     h5[..4].copy_from_slice(&h);
@@ -141,7 +123,6 @@ #[cfg(test)]
 mod tests {
     use plonky2::field::types::Sample;
 
-    // use rand::prelude::*;
     use super::*;
 
     #[test]
@@ -152,6 +133,9 @@ test_encrypt_decrypt_op(11);
         test_encrypt_decrypt_op(12);
         test_encrypt_decrypt_op(128);
         test_encrypt_decrypt_op(129);
+        test_encrypt_decrypt_op(1024);
+        test_encrypt_decrypt_op(1023);
+        test_encrypt_decrypt_op(1025);
     }
     fn test_encrypt_decrypt_op(msg_len: usize) {
         let mut rng = rand::thread_rng();
@@ -162,9 +146,9 @@ let nonce = F::rand();
         let msg: Vec<Fq> = (0..msg_len).map(|_| Fq::sample(&mut rng)).collect();
         let l = msg.len();
 
-        let ct = encrypt(ks, msg.clone(), nonce);
+        let ct = encrypt(ks, &msg, nonce);
 
-        let m = decrypt(ks, ct, nonce, l);
+        let m = decrypt(ks, &ct, nonce, l);
         assert_eq!(m, msg);
     }
 }