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