From c828031d9a6d3e58f13a086bbf80164936479fff Mon Sep 17 00:00:00 2001
From: samkim-crypto <skim13@cs.stanford.edu>
Date: Fri, 11 Nov 2022 23:44:02 +0000
Subject: [PATCH] [zk-token-sdk] Fix ElGamal key derivation (#28792)

* fix ElGamal key derivation

* cargo fmt
---
 zk-token-sdk/src/encryption/elgamal.rs        | 15 +++++++--------
 zk-token-sdk/src/sigma_proofs/pubkey_proof.rs | 17 ++++++++++++++++-
 2 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/zk-token-sdk/src/encryption/elgamal.rs b/zk-token-sdk/src/encryption/elgamal.rs
index cb312faa4df94a..8b84e24a46ae94 100644
--- a/zk-token-sdk/src/encryption/elgamal.rs
+++ b/zk-token-sdk/src/encryption/elgamal.rs
@@ -71,14 +71,10 @@ impl ElGamal {
     #[cfg(not(target_os = "solana"))]
     #[allow(non_snake_case)]
     fn keygen_with_scalar(s: &Scalar) -> ElGamalKeypair {
-        assert!(s != &Scalar::zero());
-
-        let P = s.invert() * &(*H);
+        let secret = ElGamalSecretKey(*s);
+        let public = ElGamalPubkey::new(&secret);
 
-        ElGamalKeypair {
-            public: ElGamalPubkey(P),
-            secret: ElGamalSecretKey(*s),
-        }
+        ElGamalKeypair { public, secret }
     }
 
     /// On input an ElGamal public key and an amount to be encrypted, the function returns a
@@ -267,7 +263,10 @@ impl ElGamalPubkey {
     /// Derives the `ElGamalPubkey` that uniquely corresponds to an `ElGamalSecretKey`.
     #[allow(non_snake_case)]
     pub fn new(secret: &ElGamalSecretKey) -> Self {
-        ElGamalPubkey(&secret.0 * &(*H))
+        let s = &secret.0;
+        assert!(s != &Scalar::zero());
+
+        ElGamalPubkey(s.invert() * &(*H))
     }
 
     pub fn get_point(&self) -> &RistrettoPoint {
diff --git a/zk-token-sdk/src/sigma_proofs/pubkey_proof.rs b/zk-token-sdk/src/sigma_proofs/pubkey_proof.rs
index 7dd49ad6d91cfe..79378ff80556bc 100644
--- a/zk-token-sdk/src/sigma_proofs/pubkey_proof.rs
+++ b/zk-token-sdk/src/sigma_proofs/pubkey_proof.rs
@@ -136,15 +136,30 @@ impl PubkeySigmaProof {
 
 #[cfg(test)]
 mod test {
-    use super::*;
+    use {
+        super::*,
+        solana_sdk::{pubkey::Pubkey, signature::Keypair},
+    };
 
     #[test]
     fn test_pubkey_proof_correctness() {
+        // random ElGamal keypair
         let keypair = ElGamalKeypair::new_rand();
 
         let mut prover_transcript = Transcript::new(b"test");
         let mut verifier_transcript = Transcript::new(b"test");
 
+        let proof = PubkeySigmaProof::new(&keypair, &mut prover_transcript);
+        assert!(proof
+            .verify(&keypair.public, &mut verifier_transcript)
+            .is_ok());
+
+        // derived ElGamal keypair
+        let keypair = ElGamalKeypair::new(&Keypair::new(), &Pubkey::default()).unwrap();
+
+        let mut prover_transcript = Transcript::new(b"test");
+        let mut verifier_transcript = Transcript::new(b"test");
+
         let proof = PubkeySigmaProof::new(&keypair, &mut prover_transcript);
         assert!(proof
             .verify(&keypair.public, &mut verifier_transcript)