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