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::{FieldExtension, quintic::QuinticExtension},
 18        goldilocks_field::GoldilocksField as F,
 19        types::{Field, Field64},
 20    },
 21    hash::{hashing::hash_n_to_m_no_pad, poseidon::PoseidonPermutation},
 22};
 23use pod2::backends::plonky2::primitives::ec::curve::{GROUP_ORDER, Point};
 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 fn encrypt(ks: Point, msg: &[Fq], nonce: [F; 2]) -> Vec<Fq> {
 42    let mut m = msg.to_vec(); // pad m
 43    m.resize(m.len().next_multiple_of(3), Fq::ZERO);
 44    assert!(m.len() < F::ORDER as usize);
 45
 46    // bind the msg length and nonce into the initial state
 47    let n_l = Fq::from_basefield_array([
 48        nonce[0],
 49        nonce[1],
 50        F::from_canonical_u64(msg.len() as u64),
 51        F::ZERO,
 52        F::ZERO,
 53    ]);
 54    let mut s: [Fq; 4] = [Fq::ZERO, ks.x, ks.u, n_l];
 55
 56    let mut ct: Vec<Fq> = vec![Fq::ZERO; m.len() + 1];
 57    for i in 0..m.len() / 3 {
 58        s = hash_state(s);
 59
 60        // absorb
 61        s[1] += m[i * 3];
 62        s[2] += m[i * 3 + 1];
 63        s[3] += m[i * 3 + 2];
 64
 65        // release
 66        ct[i * 3] = s[1];
 67        ct[i * 3 + 1] = s[2];
 68        ct[i * 3 + 2] = s[3];
 69    }
 70    s = hash_state(s);
 71    ct[m.len()] = s[1];
 72    ct
 73}
 74
 75pub fn decrypt(ks: Point, ct: &[Fq], nonce: [F; 2], l: usize) -> Vec<Fq> {
 76    let n_l = Fq::from_basefield_array([
 77        nonce[0],
 78        nonce[1],
 79        F::from_canonical_u64(l as u64),
 80        F::ZERO,
 81        F::ZERO,
 82    ]);
 83    let mut s: [Fq; 4] = [Fq::ZERO, ks.x, ks.u, n_l];
 84
 85    let mut m: Vec<Fq> = vec![Fq::ZERO; ct.len() - 1];
 86    for i in 0..ct.len() / 3 {
 87        s = hash_state(s);
 88
 89        // release
 90        m[3 * i] = ct[3 * i] - s[1];
 91        m[3 * i + 1] = ct[3 * i + 1] - s[2];
 92        m[3 * i + 2] = ct[3 * i + 2] - s[3];
 93
 94        // modify state
 95        s[1] = ct[3 * i];
 96        s[2] = ct[3 * i + 1];
 97        s[3] = ct[3 * i + 2];
 98    }
 99
100    if l > 3 {
101        if l % 3 == 2 {
102            assert_eq!(m[m.len() - 1], Fq::ZERO);
103        } else if l % 3 == 1 {
104            assert_eq!(m[m.len() - 1], Fq::ZERO);
105            assert_eq!(m[m.len() - 2], Fq::ZERO);
106        }
107    }
108    s = hash_state(s);
109    assert_eq!(ct[ct.len() - 1], s[1]);
110    m[0..l].to_vec()
111}
112
113fn hash_state(s: [Fq; 4]) -> [Fq; 4] {
114    let elems: [F; 4 * 5] = array::from_fn(|i| s[i / 5].0[i % 5]);
115    // Note: here we're using plonky2's normal poseidon, not duplex-sponge
116    let h = hash_n_to_m_no_pad::<F, PoseidonPermutation<_>>(&elems, 4 * 5);
117    array::from_fn(|i| Fq::from_basefield_array(array::from_fn::<_, 5, _>(|j| h[j + 5 * i])))
118}
119
120#[cfg(test)]
121mod tests {
122    use plonky2::field::types::Sample;
123
124    use super::*;
125
126    #[test]
127    fn test_encrypt_decrypt() {
128        test_encrypt_decrypt_op(9);
129        test_encrypt_decrypt_op(10);
130        test_encrypt_decrypt_op(11);
131        test_encrypt_decrypt_op(12);
132        test_encrypt_decrypt_op(128);
133        test_encrypt_decrypt_op(129);
134        test_encrypt_decrypt_op(1024);
135        test_encrypt_decrypt_op(1023);
136        test_encrypt_decrypt_op(1025);
137    }
138    fn test_encrypt_decrypt_op(msg_len: usize) {
139        let mut rng = rand::thread_rng();
140
141        let (_k, K) = new_key();
142        let ks = expanded_key(K);
143        let nonce: [F; 2] = [F::rand(), F::rand()];
144        let msg: Vec<Fq> = (0..msg_len).map(|_| Fq::sample(&mut rng)).collect();
145        let l = msg.len();
146
147        let ct = encrypt(ks, &msg, nonce);
148
149        let m = decrypt(ks, &ct, nonce, l);
150        assert_eq!(m, msg);
151        assert_ne!(&msg, &ct[..msg_len]);
152    }
153}