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