1use std::array;
  2
  3use plonky2::{
  4    field::extension::Extendable,
  5    hash::{hash_types::RichField, poseidon::PoseidonHash},
  6    iop::target::Target,
  7    plonk::circuit_builder::CircuitBuilder,
  8};
  9use pod2::backends::plonky2::{
 10    basetypes::{D, F},
 11    primitives::ec::{
 12        bits::BigUInt320Target,
 13        curve::{CircuitBuilderElliptic, Point, PointTarget},
 14    },
 15};
 16
 17use crate::circuit::PublicKeyTarget;
 18
 19pub trait CircuitBuilderHashedElGamal<F: RichField + Extendable<D>, const D: usize> {
 20    /// Encrypts an EcGFp5 point according to the 'hashed ElGamal '
 21    /// scheme described in `super::hashed_elgamal_encrypt`.
 22    fn hashed_elgamal_encrypt(
 23        &mut self,
 24        pk: &PublicKeyTarget,
 25        nonce: &BigUInt320Target,
 26        msg: [Target; 5],
 27    ) -> (PointTarget, [Target; 5]);
 28}
 29
 30impl CircuitBuilderHashedElGamal<F, D> for CircuitBuilder<F, D> {
 31    fn hashed_elgamal_encrypt(
 32        &mut self,
 33        pk: &PublicKeyTarget,
 34        nonce: &BigUInt320Target,
 35        msg: [Target; 5],
 36    ) -> (PointTarget, [Target; 5]) {
 37        let generator = self.constant_point(Point::generator());
 38        let nonce_times_gen = self.multiply_point(&nonce.bits, &generator);
 39        let nonce_times_pk = self.multiply_point(&nonce.bits, pk);
 40        let h = self.hash_n_to_m_no_pad::<PoseidonHash>(
 41            [nonce_times_pk.x.components, nonce_times_pk.u.components].concat(),
 42            5,
 43        );
 44        (nonce_times_gen, array::from_fn(|i| self.add(msg[i], h[i])))
 45    }
 46}
 47
 48#[cfg(test)]
 49mod tests {
 50    use num_bigint::RandBigInt;
 51    use plonky2::{
 52        field::types::{Field, Field64},
 53        iop::witness::{PartialWitness, WitnessWrite},
 54        plonk::{
 55            circuit_builder::CircuitBuilder, circuit_data::CircuitConfig,
 56            config::PoseidonGoldilocksConfig,
 57        },
 58    };
 59    use pod2::backends::plonky2::{
 60        basetypes::{D, F},
 61        primitives::ec::{
 62            bits::CircuitBuilderBits,
 63            curve::{CircuitBuilderElliptic, GROUP_ORDER, WitnessWriteCurve},
 64        },
 65    };
 66    use rand::{Rng, rngs::OsRng};
 67
 68    use crate::{
 69        ECGFP5SecretKey,
 70        hashed_elgamal::{circuit::CircuitBuilderHashedElGamal, hashed_elgamal_encrypt},
 71    };
 72
 73    #[test]
 74    fn hashed_elgamal_encryption() -> anyhow::Result<()> {
 75        let config = CircuitConfig::standard_recursion_config();
 76        let mut builder = CircuitBuilder::<F, D>::new(config);
 77
 78        let pk_target = builder.add_virtual_point_target();
 79        let nonce_target = builder.add_virtual_biguint320_target();
 80        let msg_target = builder.add_virtual_target_arr();
 81
 82        let ct_target = builder.hashed_elgamal_encrypt(&pk_target, &nonce_target, msg_target);
 83
 84        let data = builder.build::<PoseidonGoldilocksConfig>();
 85
 86        (0..10).try_for_each(|_| {
 87            let sk = ECGFP5SecretKey::rand();
 88            let pk = sk.public_key();
 89
 90            let msg = std::array::from_fn(|_| F::from_canonical_u64(OsRng.gen_range(0..F::ORDER)));
 91            let nonce = OsRng.gen_biguint_below(&GROUP_ORDER);
 92
 93            let encrypted_msg = hashed_elgamal_encrypt(pk, &nonce, msg);
 94
 95            let mut pw = PartialWitness::new();
 96            pw.set_point_target(&pk_target, &pk)?;
 97            pw.set_target_arr(&msg_target, &msg)?;
 98            pw.set_biguint320_target(&nonce_target, &nonce)?;
 99            pw.set_point_target(&ct_target.0, &encrypted_msg.0)?;
100            pw.set_target_arr(&ct_target.1, &encrypted_msg.1)?;
101
102            let proof = data.prove(pw)?;
103            data.verify(proof)
104        })
105    }
106}