1//! For LICENSE check out https://github.com/0xPARC/plonky2-crypto-gadgets/blob/main/LICENSE
  2//!
  3//! Implements Poseidon encryption as described at
  4//! https://drive.google.com/file/d/1EVrP3DzoGbmzkRmYnyEDcIQcXVU7GlOd/view .
  5
  6#![allow(incomplete_features)]
  7#![feature(generic_const_exprs)]
  8#![allow(dead_code)] // TMP
  9#![allow(non_snake_case)]
 10
 11use std::array;
 12
 13use num::bigint::BigUint;
 14use num_bigint::RandBigInt;
 15use plonky2::{
 16    field::{
 17        extension::{quintic::QuinticExtension, FieldExtension},
 18        goldilocks_field::GoldilocksField as F,
 19        types::{Field, Field64},
 20    },
 21    hash::{hashing::hash_n_to_hash_no_pad, poseidon::PoseidonPermutation},
 22};
 23use pod2::backends::plonky2::primitives::ec::curve::{Point, GROUP_ORDER};
 24use rand::rngs::OsRng;
 25
 26pub mod circuit;
 27
 28// ECField
 29type Fq = QuinticExtension<F>;
 30
 31pub fn new_key() -> (BigUint, Point) {
 32    let k: BigUint = OsRng.gen_biguint_below(&GROUP_ORDER);
 33    let K = &k * Point::generator();
 34    (k, K)
 35}
 36pub fn expanded_key(K: Point) -> Point {
 37    let r: BigUint = OsRng.gen_biguint_below(&GROUP_ORDER);
 38    &r * K
 39}
 40
 41pub(crate) const TWO128: Fq =
 42    QuinticExtension::<F>([F(F::ORDER - 1), F(F::ORDER - 1), F::ZERO, F::ZERO, F::ZERO]);
 43
 44// /// 2^128
 45// pub(crate) fn two_128() -> Fq {
 46//     // WIP TODO
 47//     let two32 = Fq::from(F::from_canonical_u64(2_u64.pow(32)));
 48//     let two64: Fq = two32 * two32;
 49//     let two128: Fq = two64 * two64;
 50//     // note: two64 gets two32-1
 51//     // dbg!(2_u64.pow(32));
 52//     // dbg!(F::from_canonical_u64(2_u64.pow(32)));
 53//     // dbg!(&two32);
 54//     // dbg!(&two64);
 55//     // dbg!(&two128);
 56//     two128
 57// }
 58
 59pub fn f_vec_to_fq_vec(f: &[F]) -> Vec<Fq> {
 60    f.chunks(5)
 61        .map(|c| {
 62            let mut e: [F; 5] = [F::ZERO; 5];
 63            e[0..c.len()].copy_from_slice(c);
 64            Fq::from_basefield_array(e)
 65        })
 66        .collect()
 67}
 68
 69pub fn encrypt_f(ks: Point, msg_f: &[F], nonce: [F; 2]) -> Vec<F> {
 70    let l_f = msg_f.len();
 71    dbg!(msg_f.len());
 72    dbg!(&msg_f);
 73    let msg: Vec<Fq> = f_vec_to_fq_vec(msg_f);
 74    dbg!(msg.len());
 75    dbg!(&msg);
 76    let ct = encrypt(ks, &msg, nonce);
 77    ct.into_iter().flat_map(|ct_i| ct_i.0).collect()
 78}
 79pub fn encrypt(ks: Point, msg: &[Fq], nonce: [F; 2]) -> Vec<Fq> {
 80    let mut m = msg.to_vec(); // pad m
 81    m.resize(m.len().next_multiple_of(3), Fq::ZERO);
 82    assert!(m.len() < F::ORDER as usize);
 83
 84    let l = Fq::from(F::from_canonical_u64(msg.len() as u64));
 85    let nonce_5: [F; 5] = [nonce[0], nonce[1], F::ZERO, F::ZERO, F::ZERO];
 86    let n = Fq::from_basefield_array(nonce_5);
 87    let nl128: Fq = n + l * TWO128;
 88    let mut s: [Fq; 4] = [Fq::ZERO, ks.x, ks.u, nl128];
 89
 90    let mut ct: Vec<Fq> = vec![Fq::ZERO; m.len() + 1];
 91    for i in 0..m.len() / 3 {
 92        s = hash_state(s);
 93
 94        // absorb
 95        s[1] += m[i * 3];
 96        s[2] += m[i * 3 + 1];
 97        s[3] += m[i * 3 + 2];
 98
 99        // release
100        ct[i * 3] = s[1];
101        ct[i * 3 + 1] = s[2];
102        ct[i * 3 + 2] = s[3];
103    }
104    s = hash_state(s);
105    ct[m.len()] = s[1];
106    ct
107}
108
109pub fn decrypt_f(ks: Point, ct_f: &[F], nonce: [F; 2], l: usize) -> Vec<F> {
110    let ct: Vec<Fq> = f_vec_to_fq_vec(ct_f);
111    let m = decrypt(ks, &ct, nonce, l);
112    m.to_vec().into_iter().flat_map(|e| e.0).collect()
113}
114pub fn decrypt(ks: Point, ct: &[Fq], nonce: [F; 2], l: usize) -> Vec<Fq> {
115    let l_fq = Fq::from(F::from_canonical_u64(l as u64));
116    let nonce_5: [F; 5] = [nonce[0], nonce[1], F::ZERO, F::ZERO, F::ZERO];
117    let n = Fq::from_basefield_array(nonce_5);
118    let nl128: Fq = n + l_fq * TWO128;
119    let mut s: [Fq; 4] = [Fq::ZERO, ks.x, ks.u, nl128];
120
121    let mut m: Vec<Fq> = vec![Fq::ZERO; ct.len() - 1];
122    for i in 0..ct.len() / 3 {
123        s = hash_state(s);
124
125        // release
126        m[3 * i] = s[1] + ct[3 * i];
127        m[3 * i + 1] = s[2] + ct[3 * i + 1];
128        m[3 * i + 2] = s[3] + ct[3 * i + 2];
129
130        // modify state
131        s[1] = ct[3 * i];
132        s[2] = ct[3 * i + 1];
133        s[3] = ct[3 * i + 2];
134    }
135
136    if l > 3 {
137        if l % 3 == 2 {
138            assert_eq!(m[m.len() - 1], Fq::ZERO);
139        } else if l % 3 == 1 {
140            assert_eq!(m[m.len() - 1], Fq::ZERO);
141            assert_eq!(m[m.len() - 2], Fq::ZERO);
142        }
143    }
144    s = hash_state(s);
145    assert_eq!(ct[ct.len() - 1], s[1]);
146    m[0..l].to_vec()
147}
148
149fn hash_state(s: [Fq; 4]) -> [Fq; 4] {
150    let elems: [F; 4 * 5] = array::from_fn(|i| s[i / 5].0[i % 5]);
151    // Note: here we're using plonky2's normal poseidon, not duplex-sponge
152    let h = hash_n_to_hash_no_pad::<F, PoseidonPermutation<_>>(&elems).elements;
153    let mut h5: [F; 5] = [F::ZERO; 5];
154    h5[..4].copy_from_slice(&h);
155    [Fq::from_basefield_array(h5), Fq::ZERO, Fq::ZERO, Fq::ZERO]
156}
157
158#[cfg(test)]
159mod tests {
160    use plonky2::field::types::Sample;
161
162    use super::*;
163
164    #[test]
165    fn test_encrypt_decrypt() {
166        test_encrypt_decrypt_op(9);
167        test_encrypt_decrypt_op(10);
168        test_encrypt_decrypt_op(11);
169        test_encrypt_decrypt_op(12);
170        test_encrypt_decrypt_op(128);
171        test_encrypt_decrypt_op(129);
172        test_encrypt_decrypt_op(1024);
173        test_encrypt_decrypt_op(1023);
174        test_encrypt_decrypt_op(1025);
175    }
176    fn test_encrypt_decrypt_op(msg_len: usize) {
177        let mut rng = rand::thread_rng();
178
179        let (_k, K) = new_key();
180        let ks = expanded_key(K);
181        let nonce: [F; 2] = [F::rand(), F::rand()];
182        // let msg: Vec<Fq> = (0..msg_len).map(|_| Fq::sample(&mut rng)).collect();
183        let msg: Vec<F> = (0..msg_len).map(|_| F::sample(&mut rng)).collect();
184        let l = msg.len().div_ceil(5);
185        dbg!(msg.len(), l);
186
187        let ct = encrypt_f(ks, &msg, nonce);
188
189        let m = decrypt_f(ks, &ct, nonce, l);
190        assert_eq!(m, msg);
191    }
192}