1use std::array;
  2
  3use num::bigint::BigUint;
  4use num_bigint::RandBigInt;
  5use plonky2::field::{
  6    extension::quintic::QuinticExtension,
  7    goldilocks_field::GoldilocksField as F,
  8    types::{Field, Field64},
  9};
 10use pod2::backends::plonky2::primitives::ec::curve::{GROUP_ORDER, Point};
 11use rand::{Rng, rngs::OsRng};
 12
 13pub mod circuit;
 14
 15// Message encoding constants
 16pub const BITS_PER_LIMB: usize = 32;
 17pub const TOTAL_BITS: usize = 5 * BITS_PER_LIMB;
 18
 19pub struct ECGFP5SecretKey(pub BigUint);
 20
 21impl ECGFP5SecretKey {
 22    pub fn new(s: BigUint) -> Self {
 23        assert!(&s < &GROUP_ORDER);
 24        Self(s)
 25    }
 26
 27    pub fn rand() -> Self {
 28        Self(OsRng.gen_biguint_below(&GROUP_ORDER))
 29    }
 30
 31    pub fn public_key(&self) -> Point {
 32        &self.0 * Point::generator()
 33    }
 34}
 35
 36pub fn elgamal_encrypt(pk: Point, nonce: &BigUint, msg: Point) -> (Point, Point) {
 37    assert!(nonce < &GROUP_ORDER);
 38    let nonce_times_gen = nonce * Point::generator();
 39    let nonce_times_pk = nonce * pk;
 40    (nonce_times_gen, msg + nonce_times_pk)
 41}
 42
 43pub fn elgamal_decrypt(sk: ECGFP5SecretKey, ct: (Point, Point)) -> Point {
 44    let nonce_times_pk = &sk.0 * ct.0;
 45    ct.1 + nonce_times_pk.inverse()
 46}
 47
 48/// Convert `TOTAL_BITS`-bit biguint to a point by padding with random bits.
 49pub fn biguint_to_point(x: &BigUint) -> Point {
 50    let num_bits = x.bits() as usize;
 51    assert!(num_bits <= TOTAL_BITS);
 52
 53    let bs_bits = [x.to_radix_le(2), vec![0; TOTAL_BITS - num_bits]].concat();
 54    let bs_fields: Vec<F> = bs_bits
 55        .chunks(BITS_PER_LIMB)
 56        .map(|bits| {
 57            F::from_noncanonical_biguint(
 58                BigUint::from_radix_le(bits, 2).expect("Should be a valid bit decomposition"),
 59            )
 60        })
 61        .collect();
 62
 63    std::iter::repeat_with(|| {
 64        let ec_fields: [_; 5] = array::from_fn(|i| {
 65            let random_element: u64 = OsRng.gen_range(0..F::ORDER - bs_fields[i].0);
 66            F::from_canonical_u64(
 67                bs_fields[i].0 + ((random_element >> BITS_PER_LIMB) << BITS_PER_LIMB),
 68            )
 69        });
 70        QuinticExtension(ec_fields)
 71    })
 72    .find_map(|w| Point::decompress_into_subgroup(&w).ok())
 73    .expect("Should be a match!")
 74}
 75
 76pub fn point_to_biguint(p: Point) -> BigUint {
 77    let msg_bits: Vec<_> = p
 78        .compress_from_subgroup()
 79        .unwrap()
 80        .0
 81        .into_iter()
 82        .flat_map(|f| {
 83            BigUint::from(f.0)
 84                .to_radix_le(2)
 85                .into_iter()
 86                .take(BITS_PER_LIMB)
 87                .collect::<Vec<_>>()
 88        })
 89        .collect();
 90    BigUint::from_bytes_le(
 91        &BigUint::from_radix_le(&msg_bits, 2)
 92            .expect("Should be a valid bit slice.")
 93            .to_radix_le(256),
 94    )
 95}
 96
 97#[cfg(test)]
 98mod tests {
 99    use num::BigUint;
100    use num_bigint::RandBigInt;
101    use pod2::backends::plonky2::primitives::ec::curve::{GROUP_ORDER, Point};
102    use rand::rngs::OsRng;
103
104    use crate::elgamal::{
105        ECGFP5SecretKey, TOTAL_BITS, biguint_to_point, elgamal_decrypt, elgamal_encrypt,
106        point_to_biguint,
107    };
108
109    #[test]
110    fn msg_encoding_roundtrip() {
111        (0..100).for_each(|_| {
112            let test_val = OsRng.gen_biguint_below(&BigUint::from(2u32).pow(TOTAL_BITS as u32));
113            let pt = biguint_to_point(&test_val);
114            let maybe_test_val = point_to_biguint(pt);
115            assert_eq!(test_val, maybe_test_val);
116        })
117    }
118
119    #[test]
120    fn round_trip_encryption() {
121        (0..100).for_each(|_| {
122            let sk = ECGFP5SecretKey::rand();
123            let pk = sk.public_key();
124
125            let msg = Point::new_rand_from_subgroup();
126            let nonce = OsRng.gen_biguint_below(&GROUP_ORDER);
127
128            let encrypted_msg = elgamal_encrypt(pk, &nonce, msg);
129            let decrypted_msg = elgamal_decrypt(sk, encrypted_msg);
130
131            assert_eq!(msg, decrypted_msg);
132        })
133    }
134}