1use plonky2::{
  2    field::extension::Extendable, hash::hash_types::RichField, iop::witness::PartialWitness,
  3    plonk::circuit_builder::CircuitBuilder,
  4};
  5use pod2::backends::plonky2::{
  6    basetypes::{D, F},
  7    primitives::ec::{
  8        bits::{BigUInt320Target, CircuitBuilderBits},
  9        curve::{CircuitBuilderElliptic, Point, PointTarget, WitnessWriteCurve},
 10    },
 11};
 12
 13use crate::elgamal::ECGFP5SecretKey;
 14
 15pub struct ECGFP5SecretKeyTarget(pub BigUInt320Target);
 16
 17pub type PublicKeyTarget = PointTarget;
 18
 19pub trait CircuitBuilderECGFP5PublicKey<F: RichField + Extendable<D>, const D: usize> {
 20    /// Adds secret key target
 21    fn add_secret_key(&mut self) -> ECGFP5SecretKeyTarget;
 22
 23    /// Derives public key from secret key
 24    fn public_key(&mut self, sk: &ECGFP5SecretKeyTarget) -> PublicKeyTarget;
 25}
 26
 27pub trait CircuitBuilderElGamal<F: RichField + Extendable<D>, const D: usize> {
 28    /// Encrypts an EcGFp5 point using a public key and nonce.
 29    fn elgamal_encrypt(
 30        &mut self,
 31        pk: &PublicKeyTarget,
 32        nonce: &BigUInt320Target,
 33        msg: &PointTarget,
 34    ) -> (PointTarget, PointTarget);
 35}
 36
 37impl CircuitBuilderECGFP5PublicKey<F, D> for CircuitBuilder<F, D> {
 38    fn add_secret_key(&mut self) -> ECGFP5SecretKeyTarget {
 39        ECGFP5SecretKeyTarget(self.add_virtual_biguint320_target())
 40    }
 41
 42    fn public_key(&mut self, sk: &ECGFP5SecretKeyTarget) -> PublicKeyTarget {
 43        let generator = self.constant_point(Point::generator());
 44        self.multiply_point(&sk.0.bits, &generator)
 45    }
 46}
 47
 48impl CircuitBuilderElGamal<F, D> for CircuitBuilder<F, D> {
 49    fn elgamal_encrypt(
 50        &mut self,
 51        pk: &PublicKeyTarget,
 52        nonce: &BigUInt320Target,
 53        msg: &PointTarget,
 54    ) -> (PointTarget, PointTarget) {
 55        let generator = self.constant_point(Point::generator());
 56        let nonce_times_gen = self.multiply_point(&nonce.bits, &generator);
 57        let nonce_times_pk = self.multiply_point(&nonce.bits, pk);
 58        (nonce_times_gen, self.add_point(msg, &nonce_times_pk))
 59    }
 60}
 61
 62pub trait PartialWitnessECGFP5PublicKey {
 63    fn set_secret_key_target(
 64        &mut self,
 65        target: &ECGFP5SecretKeyTarget,
 66        value: &ECGFP5SecretKey,
 67    ) -> anyhow::Result<()>;
 68}
 69
 70impl PartialWitnessECGFP5PublicKey for PartialWitness<F> {
 71    fn set_secret_key_target(
 72        &mut self,
 73        target: &ECGFP5SecretKeyTarget,
 74        value: &ECGFP5SecretKey,
 75    ) -> anyhow::Result<()> {
 76        self.set_biguint320_target(&target.0, &value.0)
 77    }
 78}
 79
 80#[cfg(test)]
 81mod tests {
 82    use num_bigint::RandBigInt;
 83    use plonky2::{
 84        iop::witness::PartialWitness,
 85        plonk::{
 86            circuit_builder::CircuitBuilder, circuit_data::CircuitConfig,
 87            config::PoseidonGoldilocksConfig,
 88        },
 89    };
 90    use pod2::backends::plonky2::{
 91        basetypes::{D, F},
 92        primitives::ec::{
 93            bits::CircuitBuilderBits,
 94            curve::{CircuitBuilderElliptic, GROUP_ORDER, Point, WitnessWriteCurve},
 95        },
 96    };
 97    use rand::rngs::OsRng;
 98
 99    use crate::elgamal::{
100        ECGFP5SecretKey,
101        circuit::{
102            CircuitBuilderECGFP5PublicKey, CircuitBuilderElGamal, PartialWitnessECGFP5PublicKey,
103        },
104        elgamal_encrypt,
105    };
106
107    #[test]
108    fn public_key_calculation() -> anyhow::Result<()> {
109        let config = CircuitConfig::standard_recursion_config();
110        let mut builder = CircuitBuilder::<F, D>::new(config);
111
112        let sk_target = builder.add_secret_key();
113        let pk_target = builder.public_key(&sk_target);
114
115        let data = builder.build::<PoseidonGoldilocksConfig>();
116        (0..10).try_for_each(|_| {
117            let sk = ECGFP5SecretKey::rand();
118            let pk = sk.public_key();
119
120            let mut pw = PartialWitness::new();
121            pw.set_secret_key_target(&sk_target, &sk)?;
122            pw.set_point_target(&pk_target, &pk)?;
123
124            let proof = data.prove(pw)?;
125            data.verify(proof)
126        })
127    }
128
129    #[test]
130    fn elgamal_encryption() -> anyhow::Result<()> {
131        let config = CircuitConfig::standard_recursion_config();
132        let mut builder = CircuitBuilder::<F, D>::new(config);
133
134        let pk_target = builder.add_virtual_point_target();
135        let nonce_target = builder.add_virtual_biguint320_target();
136        let msg_target = builder.add_virtual_point_target();
137
138        let ct_target = builder.elgamal_encrypt(&pk_target, &nonce_target, &msg_target);
139
140        let data = builder.build::<PoseidonGoldilocksConfig>();
141
142        (0..10).try_for_each(|_| {
143            let sk = ECGFP5SecretKey::rand();
144            let pk = sk.public_key();
145
146            let msg = Point::new_rand_from_subgroup();
147            let nonce = OsRng.gen_biguint_below(&GROUP_ORDER);
148
149            let encrypted_msg = elgamal_encrypt(pk, &nonce, msg);
150
151            let mut pw = PartialWitness::new();
152            pw.set_point_target(&pk_target, &pk)?;
153            pw.set_point_target(&msg_target, &msg)?;
154            pw.set_biguint320_target(&nonce_target, &nonce)?;
155            pw.set_point_target(&ct_target.0, &encrypted_msg.0)?;
156            pw.set_point_target(&ct_target.1, &encrypted_msg.1)?;
157
158            let proof = data.prove(pw)?;
159            data.verify(proof)
160        })
161    }
162}