[zk-token-sdk] Generalize encryption key derivation from signers (#31784)

* generalize ElGamal keypair derivation from signer

* generalize AeKey derivation from signer

* add `tiny-bip39` as a dev dependency for tests

Cargo.lock

@@ -7129,6 +7129,7 @@ dependencies = [
  "solana-sdk",
  "subtle",
  "thiserror",
+ "tiny-bip39",
  "zeroize",
 ]
 

zk-token-sdk/Cargo.toml

@@ -14,6 +14,9 @@ num-derive = { workspace = true }
 num-traits = { workspace = true }
 solana-program = { workspace = true }
 
+[dev-dependencies]
+tiny-bip39 = { workspace = true }
+
 [target.'cfg(not(target_os = "solana"))'.dependencies]
 aes-gcm-siv = { workspace = true }
 arrayref = { workspace = true }

zk-token-sdk/src/encryption/auth_encryption.rs

@@ -16,9 +16,6 @@ use {
     sha3::{Digest, Sha3_512},
     solana_sdk::{
         derivation_path::DerivationPath,
-        instruction::Instruction,
-        message::Message,
-        pubkey::Pubkey,
         signature::Signature,
         signer::{
             keypair::generate_seed_from_seed_phrase_and_passphrase, EncodableKey, SeedDerivable,
@@ -85,20 +82,26 @@ impl AuthenticatedEncryption {
 #[derive(Debug, Zeroize)]
 pub struct AeKey([u8; 16]);
 impl AeKey {
-    pub fn new(signer: &dyn Signer, address: &Pubkey) -> Result<Self, SignerError> {
+    pub fn new_from_signer(signer: &dyn Signer, tag: &[u8]) -> Result<Self, Box<dyn error::Error>> {
-        let message = Message::new(
+        let seed = Self::seed_from_signer(signer, tag)?;
-            &[Instruction::new_with_bytes(*address, b"AeKey", vec![])],
+        Self::from_seed(&seed)
-            Some(&signer.try_pubkey()?),
+    }
-        );
+
-        let signature = signer.try_sign_message(&message.serialize())?;
+    pub fn seed_from_signer(signer: &dyn Signer, tag: &[u8]) -> Result<Vec<u8>, SignerError> {
+        let message = [b"AeKey", tag].concat();
+        let signature = signer.try_sign_message(&message)?;
 
         // Some `Signer` implementations return the default signature, which is not suitable for
         // use as key material
         if bool::from(signature.as_ref().ct_eq(Signature::default().as_ref())) {
-            Err(SignerError::Custom("Rejecting default signature".into()))
+            return Err(SignerError::Custom("Rejecting default signature".into()));
-        } else {
-            Ok(AeKey(signature.as_ref()[..16].try_into().unwrap()))
         }
+
+        let mut hasher = Sha3_512::new();
+        hasher.update(signature.as_ref());
+        let result = hasher.finalize();
+
+        Ok(result.to_vec())
     }
 
     pub fn random<T: RngCore + CryptoRng>(rng: &mut T) -> Self {
@@ -209,7 +212,7 @@ impl fmt::Display for AeCiphertext {
 mod tests {
     use {
         super::*,
-        solana_sdk::{signature::Keypair, signer::null_signer::NullSigner},
+        solana_sdk::{pubkey::Pubkey, signature::Keypair, signer::null_signer::NullSigner},
     };
 
     #[test]
@@ -229,11 +232,15 @@ mod tests {
         let keypair2 = Keypair::new();
 
         assert_ne!(
-            AeKey::new(&keypair1, &Pubkey::default()).unwrap().0,
+            AeKey::new_from_signer(&keypair1, Pubkey::default().as_ref())
-            AeKey::new(&keypair2, &Pubkey::default()).unwrap().0,
+                .unwrap()
+                .0,
+            AeKey::new_from_signer(&keypair2, Pubkey::default().as_ref())
+                .unwrap()
+                .0,
         );
 
         let null_signer = NullSigner::new(&Pubkey::default());
-        assert!(AeKey::new(&null_signer, &Pubkey::default()).is_err());
+        assert!(AeKey::new_from_signer(&null_signer, Pubkey::default().as_ref()).is_err());
     }
 }

zk-token-sdk/src/encryption/elgamal.rs

@@ -28,9 +28,6 @@ use {
     serde::{Deserialize, Serialize},
     solana_sdk::{
         derivation_path::DerivationPath,
-        instruction::Instruction,
-        message::Message,
-        pubkey::Pubkey,
         signature::Signature,
         signer::{
             keypair::generate_seed_from_seed_phrase_and_passphrase, EncodableKey, EncodableKeypair,
@@ -45,7 +42,7 @@ use {
 #[cfg(not(target_os = "solana"))]
 use {
     rand::rngs::OsRng,
-    sha3::Sha3_512,
+    sha3::{Digest, Sha3_512},
     std::{
         error, fmt,
         io::{Read, Write},
@@ -162,24 +159,22 @@ pub struct ElGamalKeypair {
 }
 
 impl ElGamalKeypair {
-    /// Deterministically derives an ElGamal keypair from an Ed25519 signing key and a Solana
+    /// Deterministically derives an ElGamal keypair from a Solana signer and a tag.
-    /// address.
     ///
-    /// This function exists for applications where a user may not wish to maintin a Solana
+    /// This function exists for applications where a user may not wish to maintain a Solana signer
-    /// (Ed25519) keypair and an ElGamal keypair separately. A user may wish to solely maintain the
+    /// and an ElGamal keypair separately. Instead, a user can derive the ElGamal keypair
-    /// Solana keypair and then derive the ElGamal keypair on-the-fly whenever
+    /// on-the-fly whenever encryption/decryption is needed.
-    /// encryption/decryption is needed.
     ///
-    /// For the spl token-2022 confidential extension application, the ElGamal encryption public
+    /// For the spl-token-2022 confidential extension, the ElGamal public key is
-    /// key is specified in a token account address. A natural way to derive an ElGamal keypair is
+    /// specified in a token account. A natural way to derive an ElGamal keypair is to define it
-    /// then to define it from the hash of a Solana keypair and a Solana address. However, for
+    /// from the hash of a Solana keypair and a Solana address as the tag. However, for general
-    /// general hardware wallets, the signing key is not exposed in the API. Therefore, this
+    /// hardware wallets, the signing key is not exposed in the API. Therefore, this function uses
-    /// function uses a signer to sign a pre-specified message with respect to a Solana address.
+    /// a signer to sign a pre-specified message with respect to a Solana address. The resulting
-    /// The resulting signature is then hashed to derive an ElGamal keypair.
+    /// signature is then hashed to derive an ElGamal keypair.
     #[cfg(not(target_os = "solana"))]
     #[allow(non_snake_case)]
-    pub fn new(signer: &dyn Signer, address: &Pubkey) -> Result<Self, SignerError> {
+    pub fn new_from_signer(signer: &dyn Signer, tag: &[u8]) -> Result<Self, Box<dyn error::Error>> {
-        let secret = ElGamalSecretKey::new(signer, address)?;
+        let secret = ElGamalSecretKey::new_from_signer(signer, tag)?;
         let public = ElGamalPubkey::new(&secret);
         Ok(ElGamalKeypair { public, secret })
     }
@@ -367,20 +362,18 @@ impl fmt::Display for ElGamalPubkey {
 #[zeroize(drop)]
 pub struct ElGamalSecretKey(Scalar);
 impl ElGamalSecretKey {
-    /// Deterministically derives an ElGamal keypair from an Ed25519 signing key and a Solana
+    /// Deterministically derives an ElGamal secret key from a Solana signer and a tag.
-    /// address.
     ///
-    /// See `ElGamalKeypair::new` for more context on the key derivation.
+    /// See `ElGamalKeypair::new_from_signer` for more context on the key derivation.
-    pub fn new(signer: &dyn Signer, address: &Pubkey) -> Result<Self, SignerError> {
+    pub fn new_from_signer(signer: &dyn Signer, tag: &[u8]) -> Result<Self, Box<dyn error::Error>> {
-        let message = Message::new(
+        let seed = Self::seed_from_signer(signer, tag)?;
-            &[Instruction::new_with_bytes(
+        Self::from_seed(&seed)
-                *address,
+    }
-                b"ElGamalSecretKey",
+
-                vec![],
+    /// Derive a seed from a Solana signer used to generate an ElGamal secret key.
-            )],
+    pub fn seed_from_signer(signer: &dyn Signer, tag: &[u8]) -> Result<Vec<u8>, SignerError> {
-            Some(&signer.try_pubkey()?),
+        let message = [b"ElGamalSecretKey", tag].concat();
-        );
+        let signature = signer.try_sign_message(&message)?;
-        let signature = signer.try_sign_message(&message.serialize())?;
 
         // Some `Signer` implementations return the default signature, which is not suitable for
         // use as key material
@@ -388,9 +381,11 @@ impl ElGamalSecretKey {
             return Err(SignerError::Custom("Rejecting default signatures".into()));
         }
 
-        Ok(ElGamalSecretKey(Scalar::hash_from_bytes::<Sha3_512>(
+        let mut hasher = Sha3_512::new();
-            signature.as_ref(),
+        hasher.update(signature.as_ref());
-        )))
+        let result = hasher.finalize();
+
+        Ok(result.to_vec())
     }
 
     /// Randomly samples an ElGamal secret key.
@@ -714,7 +709,8 @@ mod tests {
     use {
         super::*,
         crate::encryption::pedersen::Pedersen,
-        solana_sdk::{signature::Keypair, signer::null_signer::NullSigner},
+        bip39::{Language, Mnemonic, MnemonicType, Seed},
+        solana_sdk::{pubkey::Pubkey, signature::Keypair, signer::null_signer::NullSigner},
         std::fs::{self, File},
     };
 
@@ -949,21 +945,43 @@ mod tests {
     }
 
     #[test]
-    fn test_secret_key_new() {
+    fn test_secret_key_new_from_signer() {
         let keypair1 = Keypair::new();
         let keypair2 = Keypair::new();
 
         assert_ne!(
-            ElGamalSecretKey::new(&keypair1, &Pubkey::default())
+            ElGamalSecretKey::new_from_signer(&keypair1, Pubkey::default().as_ref())
                 .unwrap()
                 .0,
-            ElGamalSecretKey::new(&keypair2, &Pubkey::default())
+            ElGamalSecretKey::new_from_signer(&keypair2, Pubkey::default().as_ref())
                 .unwrap()
                 .0,
         );
 
         let null_signer = NullSigner::new(&Pubkey::default());
-        assert!(ElGamalSecretKey::new(&null_signer, &Pubkey::default()).is_err());
+        assert!(
+            ElGamalSecretKey::new_from_signer(&null_signer, Pubkey::default().as_ref()).is_err()
+        );
+    }
+
+    #[test]
+    fn test_keypair_from_seed() {
+        let good_seed = vec![0; 32];
+        assert!(ElGamalKeypair::from_seed(&good_seed).is_ok());
+
+        let too_short_seed = vec![0; 31];
+        assert!(ElGamalKeypair::from_seed(&too_short_seed).is_err());
+    }
+
+    #[test]
+    fn test_keypair_from_seed_phrase_and_passphrase() {
+        let mnemonic = Mnemonic::new(MnemonicType::Words12, Language::English);
+        let passphrase = "42";
+        let seed = Seed::new(&mnemonic, passphrase);
+        let expected_keypair = ElGamalKeypair::from_seed(seed.as_bytes()).unwrap();
+        let keypair =
+            ElGamalKeypair::from_seed_phrase_and_passphrase(mnemonic.phrase(), passphrase).unwrap();
+        assert_eq!(keypair.public, expected_keypair.public);
     }
 
     #[test]

zk-token-sdk/src/sigma_proofs/pubkey_proof.rs

@@ -157,7 +157,8 @@ mod test {
             .is_ok());
 
         // derived ElGamal keypair
-        let keypair = ElGamalKeypair::new(&Keypair::new(), &Pubkey::default()).unwrap();
+        let keypair =
+            ElGamalKeypair::new_from_signer(&Keypair::new(), Pubkey::default().as_ref()).unwrap();
 
         let mut prover_transcript = Transcript::new(b"test");
         let mut verifier_transcript = Transcript::new(b"test");