[zk-token-sdk] Fix ElGamal key derivation (#28792)

* fix ElGamal key derivation

* cargo fmt

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 secret = ElGamalSecretKey(*s);
+        let public = ElGamalPubkey::new(&secret);
 
-        let P = s.invert() * &(*H);
+        ElGamalKeypair { public, secret }
-
-        ElGamalKeypair {
-            public: ElGamalPubkey(P),
-            secret: ElGamalSecretKey(*s),
-        }
     }
 
     /// 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 {

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)