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::BigUint;
 83    use num_bigint::RandBigInt;
 84    use plonky2::{
 85        iop::witness::PartialWitness,
 86        plonk::{
 87            circuit_builder::CircuitBuilder, circuit_data::CircuitConfig,
 88            config::PoseidonGoldilocksConfig,
 89        },
 90    };
 91    use pod2::backends::plonky2::{
 92        basetypes::{D, F},
 93        primitives::ec::{
 94            bits::CircuitBuilderBits,
 95            curve::{CircuitBuilderElliptic, GROUP_ORDER, Point, WitnessWriteCurve},
 96        },
 97    };
 98    use rand::rngs::OsRng;
 99
100    use crate::elgamal::{
101        ECGFP5SecretKey, TOTAL_BITS,
102        circuit::{
103            CircuitBuilderECGFP5PublicKey, CircuitBuilderElGamal, PartialWitnessECGFP5PublicKey,
104        },
105        elgamal_decrypt, elgamal_encrypt,
106    };
107
108    #[test]
109    fn public_key_calculation() -> anyhow::Result<()> {
110        let config = CircuitConfig::standard_recursion_config();
111        let mut builder = CircuitBuilder::<F, D>::new(config);
112
113        let sk_target = builder.add_secret_key();
114        let pk_target = builder.public_key(&sk_target);
115
116        let data = builder.build::<PoseidonGoldilocksConfig>();
117        (0..10).try_for_each(|_| {
118            let sk = ECGFP5SecretKey::rand();
119            let pk = sk.public_key();
120
121            let mut pw = PartialWitness::new();
122            pw.set_secret_key_target(&sk_target, &sk)?;
123            pw.set_point_target(&pk_target, &pk)?;
124
125            let proof = data.prove(pw)?;
126            data.verify(proof)
127        })
128    }
129
130    #[test]
131    fn elgamal_encryption() -> anyhow::Result<()> {
132        let config = CircuitConfig::standard_recursion_config();
133        let mut builder = CircuitBuilder::<F, D>::new(config);
134
135        let pk_target = builder.add_virtual_point_target();
136        let nonce_target = builder.add_virtual_biguint320_target();
137        let msg_target = builder.add_virtual_point_target();
138
139        let ct_target = builder.elgamal_encrypt(&pk_target, &nonce_target, &msg_target);
140
141        let data = builder.build::<PoseidonGoldilocksConfig>();
142
143        (0..10).try_for_each(|_| {
144            let sk = ECGFP5SecretKey::rand();
145            let pk = sk.public_key();
146
147            let msg = Point::new_rand_from_subgroup();
148            let nonce = OsRng.gen_biguint_below(&GROUP_ORDER);
149
150            let encrypted_msg = elgamal_encrypt(pk, &nonce, msg);
151
152            let mut pw = PartialWitness::new();
153            pw.set_point_target(&pk_target, &pk)?;
154            pw.set_point_target(&msg_target, &msg)?;
155            pw.set_biguint320_target(&nonce_target, &nonce)?;
156            pw.set_point_target(&ct_target.0, &encrypted_msg.0)?;
157            pw.set_point_target(&ct_target.1, &encrypted_msg.1)?;
158
159            let proof = data.prove(pw)?;
160            data.verify(proof)
161        })
162    }
163}