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