1//! For LICENSE check out https://github.com/0xPARC/plonky2-crypto-gadgets/blob/main/LICENSE
 2
 3use std::array;
 4
 5use num::bigint::BigUint;
 6use plonky2::{
 7    field::goldilocks_field::GoldilocksField as F,
 8    hash::{hashing::hash_n_to_m_no_pad, poseidon::PoseidonPermutation},
 9};
10use pod2::backends::plonky2::primitives::ec::curve::{GROUP_ORDER, Point};
11
12use crate::ECGFP5SecretKey;
13
14pub mod circuit;
15
16/// A variant of 'hashed ElGamal' with addition in place of
17/// XORing. Assumes a message of 5 field elements and uses the
18/// Poseidon hash function outputting 5 field elements.
19pub fn hashed_elgamal_encrypt(pk: Point, nonce: &BigUint, msg: [F; 5]) -> (Point, [F; 5]) {
20    assert!(nonce < &GROUP_ORDER);
21    let nonce_times_gen = nonce * Point::generator();
22    let nonce_times_pk = nonce * pk;
23    let h = hash_n_to_m_no_pad::<F, PoseidonPermutation<_>>(&nonce_times_pk.as_fields(), 5);
24    (nonce_times_gen, array::from_fn(|i| msg[i] + h[i]))
25}
26
27/// Inverse to `hashed_elgamal_encrypt`.
28pub fn hashed_elgamal_decrypt(sk: ECGFP5SecretKey, ct: (Point, [F; 5])) -> [F; 5] {
29    let nonce_times_pk = &sk.0 * ct.0;
30    let h = hash_n_to_m_no_pad::<F, PoseidonPermutation<_>>(&nonce_times_pk.as_fields(), 5);
31    array::from_fn(|i| ct.1[i] - h[i])
32}
33
34#[cfg(test)]
35mod tests {
36    use num_bigint::RandBigInt;
37    use plonky2::field::{
38        goldilocks_field::GoldilocksField as F,
39        types::{Field, Field64},
40    };
41    use pod2::backends::plonky2::primitives::ec::curve::GROUP_ORDER;
42    use rand::{Rng, rngs::OsRng};
43
44    use crate::{
45        ECGFP5SecretKey,
46        hashed_elgamal::{hashed_elgamal_decrypt, hashed_elgamal_encrypt},
47    };
48
49    #[test]
50    fn round_trip_encryption() {
51        (0..100).for_each(|_| {
52            let sk = ECGFP5SecretKey::rand();
53            let pk = sk.public_key();
54
55            let msg = std::array::from_fn(|_| F::from_canonical_u64(OsRng.gen_range(0..F::ORDER)));
56            let nonce = OsRng.gen_biguint_below(&GROUP_ORDER);
57
58            let encrypted_msg = hashed_elgamal_encrypt(pk, &nonce, msg);
59            let decrypted_msg = hashed_elgamal_decrypt(sk, encrypted_msg);
60
61            assert_eq!(msg, decrypted_msg);
62        })
63    }
64}