From 69fab16e83006b7a9740ced1350efded3be38c77 Mon Sep 17 00:00:00 2001
From: Michael Vines <mvines@gmail.com>
Date: Fri, 29 Oct 2021 10:47:49 -0700
Subject: [PATCH] ElGamalKeypair::new() now generates valid keypairs

---
 zk-token-sdk/src/encryption/elgamal.rs | 33 ++++++++++++++++++++++----
 1 file changed, 28 insertions(+), 5 deletions(-)

diff --git a/zk-token-sdk/src/encryption/elgamal.rs b/zk-token-sdk/src/encryption/elgamal.rs
index e7ed5d4ab..bcdb090e2 100644
--- a/zk-token-sdk/src/encryption/elgamal.rs
+++ b/zk-token-sdk/src/encryption/elgamal.rs
@@ -37,8 +37,8 @@ use {
 
 struct ElGamal;
 impl ElGamal {
-    /// On input a randomness generator, the function generates the public and
-    /// secret keys for ElGamal encryption.
+    /// The function generates the public and secret keys for ElGamal encryption from the provided
+    /// randomness generator
     #[cfg(not(target_arch = "bpf"))]
     #[allow(non_snake_case)]
     fn keygen<T: RngCore + CryptoRng>(rng: &mut T) -> ElGamalKeypair {
@@ -52,6 +52,15 @@ impl ElGamal {
             }
         }
 
+        Self::keygen_with_scalar(s)
+    }
+
+    /// Generates the public and secret keys for ElGamal encryption from a non-zero Scalar
+    #[cfg(not(target_arch = "bpf"))]
+    #[allow(non_snake_case)]
+    fn keygen_with_scalar(s: Scalar) -> ElGamalKeypair {
+        assert!(s != Scalar::zero());
+
         let H = PedersenBase::default().H;
         let P = s.invert() * H;
 
@@ -142,10 +151,24 @@ impl ElGamalKeypair {
     #[cfg(not(target_arch = "bpf"))]
     #[allow(non_snake_case)]
     pub fn new(signer: &dyn Signer, address: &Pubkey) -> Result<Self, SignerError> {
-        let secret = ElGamalSecretKey::new(signer, address)?;
-        let public = ElGamalPubkey::new(&secret);
+        let message = Message::new(
+            &[Instruction::new_with_bytes(
+                *address,
+                b"ElGamalSecretKey",
+                vec![],
+            )],
+            Some(&signer.try_pubkey()?),
+        );
+        let signature = signer.try_sign_message(&message.serialize())?;
 
-        Ok(Self { public, secret })
+        // Some `Signer` implementations return the default signature, which is not suitable for
+        // use as key material
+        if signature == Signature::default() {
+            return Err(SignerError::Custom("Rejecting default signature".into()));
+        }
+
+        let scalar = Scalar::hash_from_bytes::<Sha3_512>(signature.as_ref());
+        Ok(ElGamal::keygen_with_scalar(scalar))
     }
 
     /// Generates the public and secret keys for ElGamal encryption.