1//! For LICENSE check out https://github.com/0xPARC/plonky2-crypto-gadgets/blob/main/LICENSE
  2
  3#![allow(incomplete_features)]
  4#![feature(generic_const_exprs)]
  5#![allow(dead_code)] // TMP
  6#![allow(non_snake_case)]
  7
  8use std::array;
  9
 10use num::bigint::BigUint;
 11use num_bigint::RandBigInt;
 12use plonky2::{
 13    field::{
 14        extension::{FieldExtension, quintic::QuinticExtension},
 15        goldilocks_field::GoldilocksField as F,
 16        types::Field,
 17    },
 18    hash::{hashing::hash_n_to_hash_no_pad, poseidon::PoseidonPermutation},
 19};
 20use pod2::backends::plonky2::primitives::ec::curve::{GROUP_ORDER, Point};
 21use rand::rngs::OsRng;
 22
 23pub mod circuit;
 24
 25// ECField
 26type Fq = QuinticExtension<F>;
 27
 28pub fn new_key() -> (BigUint, Point) {
 29    let k: BigUint = OsRng.gen_biguint_below(&GROUP_ORDER);
 30    let K = &k * Point::generator();
 31    (k, K)
 32}
 33pub fn expanded_key(K: Point) -> Point {
 34    let r: BigUint = OsRng.gen_biguint_below(&GROUP_ORDER);
 35    &r * K
 36}
 37
 38pub(crate) fn two_128() -> Fq {
 39    // WIP TODO
 40    let two32 = Fq::from(F::from_canonical_u64(2_u64.pow(32)));
 41    // dbg!(2_u64.pow(32));
 42    // dbg!(F::from_canonical_u64(2_u64.pow(32)));
 43    // dbg!(&two32);
 44    let two64: Fq = two32 * two32;
 45    // dbg!(&two64);
 46    let two128: Fq = two64 * two64;
 47    // dbg!(&two128);
 48    two128
 49}
 50
 51pub fn encrypt(ks: Point, msg: &[Fq], nonce: F) -> Vec<Fq> {
 52    // pad m
 53    let mut m = msg.to_vec();
 54    m.resize(m.len().next_multiple_of(3), Fq::ZERO);
 55
 56    let l = Fq::from(F::from_canonical_u64(msg.len() as u64));
 57    let n = Fq::from(nonce);
 58    let nl128: Fq = n + l * two_128();
 59    let mut s: [Fq; 4] = [Fq::ZERO, ks.x, ks.u, nl128];
 60
 61    let mut ct: Vec<Fq> = vec![Fq::ZERO; m.len() + 1];
 62    for i in 0..m.len() / 3 {
 63        s = hash_state(s);
 64
 65        // absorb
 66        s[1] += m[i * 3];
 67        s[2] += m[i * 3 + 1];
 68        s[3] += m[i * 3 + 2];
 69
 70        // release
 71        ct[i * 3] = s[1];
 72        ct[i * 3 + 1] = s[2];
 73        ct[i * 3 + 2] = s[3];
 74    }
 75    s = hash_state(s);
 76    ct[m.len()] = s[1];
 77    ct
 78}
 79
 80pub fn decrypt(ks: Point, ct: &[Fq], nonce: F, l: usize) -> Vec<Fq> {
 81    let l_fq = Fq::from(F::from_canonical_u64(l as u64));
 82    let n = Fq::from(nonce);
 83    let nl128: Fq = n + l_fq * two_128();
 84    let mut s: [Fq; 4] = [Fq::ZERO, ks.x, ks.u, nl128];
 85
 86    let mut m: Vec<Fq> = vec![Fq::ZERO; ct.len() - 1];
 87    for i in 0..ct.len() / 3 {
 88        s = hash_state(s);
 89
 90        // release
 91        m[3 * i] = s[1] + ct[3 * i];
 92        m[3 * i + 1] = s[2] + ct[3 * i + 1];
 93        m[3 * i + 2] = s[3] + ct[3 * i + 2];
 94
 95        // modify state
 96        s[1] = ct[3 * i];
 97        s[2] = ct[3 * i + 1];
 98        s[3] = ct[3 * i + 2];
 99    }
100
101    if l > 3 {
102        if l % 3 == 2 {
103            assert_eq!(m[m.len() - 1], Fq::ZERO);
104        } else if l % 3 == 1 {
105            assert_eq!(m[m.len() - 1], Fq::ZERO);
106            assert_eq!(m[m.len() - 2], Fq::ZERO);
107        }
108    }
109    s = hash_state(s);
110    assert_eq!(ct[ct.len() - 1], s[1]);
111    m[0..l].to_vec()
112}
113
114fn hash_state(s: [Fq; 4]) -> [Fq; 4] {
115    let elems: [F; 4 * 5] = array::from_fn(|i| s[i / 5].0[i % 5]);
116    let h = hash_n_to_hash_no_pad::<F, PoseidonPermutation<_>>(&elems).elements;
117    let mut h5: [F; 5] = [F::ZERO; 5];
118    h5[..4].copy_from_slice(&h);
119    [Fq::from_basefield_array(h5), Fq::ZERO, Fq::ZERO, Fq::ZERO]
120}
121
122#[cfg(test)]
123mod tests {
124    use plonky2::field::types::Sample;
125
126    use super::*;
127
128    #[test]
129    fn test_encrypt_decrypt() {
130        test_encrypt_decrypt_op(9);
131        test_encrypt_decrypt_op(10);
132        test_encrypt_decrypt_op(11);
133        test_encrypt_decrypt_op(12);
134        test_encrypt_decrypt_op(128);
135        test_encrypt_decrypt_op(129);
136        test_encrypt_decrypt_op(1024);
137        test_encrypt_decrypt_op(1023);
138        test_encrypt_decrypt_op(1025);
139    }
140    fn test_encrypt_decrypt_op(msg_len: usize) {
141        let mut rng = rand::thread_rng();
142
143        let (_k, K) = new_key();
144        let ks = expanded_key(K);
145        let nonce = F::rand();
146        let msg: Vec<Fq> = (0..msg_len).map(|_| Fq::sample(&mut rng)).collect();
147        let l = msg.len();
148
149        let ct = encrypt(ks, &msg, nonce);
150
151        let m = decrypt(ks, &ct, nonce, l);
152        assert_eq!(m, msg);
153    }
154}