[zk-token-sdk, clap-v3-utils] Implement `EncodableKey` for encryption keys (#31496)

* implement EncodableKey for ElGamalKeypair

* implement EncodableKey for AeKey

* add keypair_from_path and keypair_from_seed support for encryption keys

* remove duplicate methods from traits

Cargo.lock

@@ -5292,6 +5292,7 @@ dependencies = [
  "solana-perf",
  "solana-remote-wallet",
  "solana-sdk 1.16.0",
+ "solana-zk-token-sdk 1.16.0",
  "tempfile",
  "thiserror",
  "tiny-bip39",

clap-v3-utils/Cargo.toml

@@ -16,6 +16,7 @@ rpassword = { workspace = true }
 solana-perf = { workspace = true }
 solana-remote-wallet = { workspace = true }
 solana-sdk = { workspace = true }
+solana-zk-token-sdk = { workspace = true }
 thiserror = { workspace = true }
 tiny-bip39 = { workspace = true }
 uriparse = { workspace = true }

clap-v3-utils/src/keypair.rs

@@ -33,6 +33,7 @@ use {
             EncodableKey, Keypair, NullSigner, Presigner, Signature, Signer,
         },
     },
+    solana_zk_token_sdk::encryption::{auth_encryption::AeKey, elgamal::ElGamalKeypair},
     std::{
         cell::RefCell,
         convert::TryFrom,
@@ -1020,6 +1021,103 @@ fn confirm_keypair_pubkey(keypair: &Keypair) {
     }
 }
 
+/// Loads an [ElGamalKeypair] from one of several possible sources.
+///
+/// If `confirm_pubkey` is `true` then after deriving the keypair, the user will
+/// be prompted to confirm that the ElGamal pubkey is as expected.
+///
+/// The way this function interprets its arguments is analogous to that of
+/// [`signer_from_path`].
+///
+/// The bip32 hierarchical derivation of an ElGamal keypair is not currently
+/// supported.
+///
+/// # Examples
+///
+/// ```no_run`
+/// use clap::{Arg, Command};
+/// use solana_clap_v3_utils::keypair::elgamal_keypair_from_path;
+///
+/// let clap_app = Command::new("my-program")
+///     // The argument we'll parse as a signer "path"
+///     .arg(Arg::new("elgamal-keypair")
+///         .required(true)
+///         .help("The default signer"));
+///
+/// let clap_matches = clap_app.get_matches();
+/// let elgamal_keypair_str: String = clap_matches.value_of_t_or_exit("elgamal-keypair");
+///
+/// let elgamal_keypair = elgamal_keypair_from_path(
+///     &clap_matches,
+///     &elgamal_keypair_str,
+///     "elgamal-keypair",
+///     false,
+/// )?;
+/// # Ok::<(), Box<dyn std::error::Error>>(())
+/// ```
+pub fn elgamal_keypair_from_path(
+    matches: &ArgMatches,
+    path: &str,
+    elgamal_keypair_name: &str,
+    confirm_pubkey: bool,
+) -> Result<ElGamalKeypair, Box<dyn error::Error>> {
+    let elgamal_keypair = encodable_key_from_path(matches, path, elgamal_keypair_name)?;
+    if confirm_pubkey {
+        confirm_elgamal_keypair_pubkey(&elgamal_keypair);
+    }
+    Ok(elgamal_keypair)
+}
+
+fn confirm_elgamal_keypair_pubkey(keypair: &ElGamalKeypair) {
+    let elgamal_pubkey = keypair.public;
+    print!("Recovered ElGamal pubkey `{elgamal_pubkey:?}`. Continue? (y/n): ");
+    let _ignored = stdout().flush();
+    let mut input = String::new();
+    stdin().read_line(&mut input).expect("Unexpected input");
+    if input.to_lowercase().trim() != "y" {
+        println!("Exiting");
+        exit(1);
+    }
+}
+
+/// Loads an [AeKey] from one of several possible sources.
+///
+/// The way this function interprets its arguments is analogous to that of
+/// [`signer_from_path`].
+///
+/// The bip32 hierarchical derivation of an authenticated encryption key is not
+/// currently supported.
+///
+/// # Examples
+///
+/// ```no_run`
+/// use clap::{Arg, Command};
+/// use solana_clap_v3_utils::keypair::ae_key_from_path;
+///
+/// let clap_app = Command::new("my-program")
+///     // The argument we'll parse as a signer "path"
+///     .arg(Arg::new("ae-key")
+///         .required(true)
+///         .help("The default signer"));
+///
+/// let clap_matches = clap_app.get_matches();
+/// let ae_key_str: String = clap_matches.value_of_t_or_exit("ae-key");
+///
+/// let ae_key = ae_key_from_path(
+///     &clap_matches,
+///     &ae_key_str,
+///     "ae-key",
+/// )?;
+/// # Ok::<(), Box<dyn std::error::Error>>(())
+/// ```
+pub fn ae_key_from_path(
+    matches: &ArgMatches,
+    path: &str,
+    key_name: &str,
+) -> Result<AeKey, Box<dyn error::Error>> {
+    encodable_key_from_path(matches, path, key_name)
+}
+
 fn encodable_key_from_path<K: EncodableKey>(
     matches: &ArgMatches,
     path: &str,
@@ -1082,6 +1180,40 @@ pub fn keypair_from_seed_phrase(
     Ok(keypair)
 }
 
+/// Reads user input from stdin to retrieve a seed phrase and passphrase for ElGamal keypair
+/// derivation.
+///
+/// Optionally skips validation of seed phrase. Optionally confirms recovered public key.
+pub fn elgamal_keypair_from_seed_phrase(
+    elgamal_keypair_name: &str,
+    skip_validation: bool,
+    confirm_pubkey: bool,
+    derivation_path: Option<DerivationPath>,
+    legacy: bool,
+) -> Result<ElGamalKeypair, Box<dyn error::Error>> {
+    let elgamal_keypair: ElGamalKeypair = encodable_key_from_seed_phrase(
+        elgamal_keypair_name,
+        skip_validation,
+        derivation_path,
+        legacy,
+    )?;
+    if confirm_pubkey {
+        confirm_elgamal_keypair_pubkey(&elgamal_keypair);
+    }
+    Ok(elgamal_keypair)
+}
+
+/// Reads user input from stdin to retrieve a seed phrase and passphrase for an authenticated
+/// encryption keypair derivation.
+pub fn ae_key_from_seed_phrase(
+    keypair_name: &str,
+    skip_validation: bool,
+    derivation_path: Option<DerivationPath>,
+    legacy: bool,
+) -> Result<ElGamalKeypair, Box<dyn error::Error>> {
+    encodable_key_from_seed_phrase(keypair_name, skip_validation, derivation_path, legacy)
+}
+
 fn encodable_key_from_seed_phrase<K: EncodableKey>(
     key_name: &str,
     skip_validation: bool,

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

@@ -3,6 +3,7 @@
 //! This module is a simple wrapper of the `Aes128GcmSiv` implementation.
 #[cfg(not(target_os = "solana"))]
 use {
+    crate::encryption::errors::AuthenticatedEncryptionError,
     aes_gcm_siv::{
         aead::{Aead, NewAead},
         Aes128GcmSiv,
@@ -11,14 +12,23 @@ use {
 };
 use {
     arrayref::{array_ref, array_refs},
+    sha3::{Digest, Sha3_512},
     solana_sdk::{
+        derivation_path::DerivationPath,
         instruction::Instruction,
         message::Message,
         pubkey::Pubkey,
         signature::Signature,
-        signer::{Signer, SignerError},
+        signer::{
+            keypair::generate_seed_from_seed_phrase_and_passphrase, EncodableKey, Signer,
+            SignerError,
+        },
+    },
+    std::{
+        convert::TryInto,
+        error, fmt,
+        io::{Read, Write},
     },
-    std::{convert::TryInto, fmt},
     subtle::ConstantTimeEq,
     zeroize::Zeroize,
 };
@@ -94,6 +104,51 @@ impl AeKey {
     }
 }
 
+impl EncodableKey for AeKey {
+    fn read<R: Read>(reader: &mut R) -> Result<Self, Box<dyn error::Error>> {
+        let bytes: [u8; 16] = serde_json::from_reader(reader)?;
+        Ok(Self(bytes))
+    }
+
+    fn write<W: Write>(&self, writer: &mut W) -> Result<String, Box<dyn error::Error>> {
+        let bytes = self.0;
+        let json = serde_json::to_string(&bytes.to_vec())?;
+        writer.write_all(&json.clone().into_bytes())?;
+        Ok(json)
+    }
+
+    fn from_seed(seed: &[u8]) -> Result<Self, Box<dyn error::Error>> {
+        const MINIMUM_SEED_LEN: usize = 16;
+
+        if seed.len() < MINIMUM_SEED_LEN {
+            return Err("Seed is too short".into());
+        }
+
+        let mut hasher = Sha3_512::new();
+        hasher.update(seed);
+        let result = hasher.finalize();
+
+        Ok(Self(result[..16].try_into()?))
+    }
+
+    fn from_seed_and_derivation_path(
+        _seed: &[u8],
+        _derivation_path: Option<DerivationPath>,
+    ) -> Result<Self, Box<dyn error::Error>> {
+        Err(AuthenticatedEncryptionError::DerivationMethodNotSupported.into())
+    }
+
+    fn from_seed_phrase_and_passphrase(
+        seed_phrase: &str,
+        passphrase: &str,
+    ) -> Result<Self, Box<dyn error::Error>> {
+        Self::from_seed(&generate_seed_from_seed_phrase_and_passphrase(
+            seed_phrase,
+            passphrase,
+        ))
+    }
+}
+
 /// For the purpose of encrypting balances for the spl token accounts, the nonce and ciphertext
 /// sizes should always be fixed.
 pub type Nonce = [u8; 12];

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

@@ -16,6 +16,7 @@
 use {
     crate::encryption::{
         discrete_log::DiscreteLog,
+        errors::ElGamalError,
         pedersen::{Pedersen, PedersenCommitment, PedersenOpening, G, H},
     },
     core::ops::{Add, Mul, Sub},
@@ -26,11 +27,15 @@ use {
     },
     serde::{Deserialize, Serialize},
     solana_sdk::{
+        derivation_path::DerivationPath,
         instruction::Instruction,
         message::Message,
         pubkey::Pubkey,
         signature::Signature,
-        signer::{Signer, SignerError},
+        signer::{
+            keypair::generate_seed_from_seed_phrase_and_passphrase, EncodableKey, Signer,
+            SignerError,
+        },
     },
     std::convert::TryInto,
     subtle::{Choice, ConstantTimeEq},
@@ -41,8 +46,7 @@ use {
     rand::rngs::OsRng,
     sha3::Sha3_512,
     std::{
-        fmt,
-        fs::{self, File, OpenOptions},
+        error, fmt,
         io::{Read, Write},
         path::Path,
     },
@@ -200,7 +204,7 @@ impl ElGamalKeypair {
     }
 
     /// Reads a JSON-encoded keypair from a `Reader` implementor
-    pub fn read_json<R: Read>(reader: &mut R) -> Result<Self, Box<dyn std::error::Error>> {
+    pub fn read_json<R: Read>(reader: &mut R) -> Result<Self, Box<dyn error::Error>> {
         let bytes: Vec<u8> = serde_json::from_reader(reader)?;
         Self::from_bytes(&bytes).ok_or_else(|| {
             std::io::Error::new(std::io::ErrorKind::Other, "Invalid ElGamalKeypair").into()
@@ -208,16 +212,12 @@ impl ElGamalKeypair {
     }
 
     /// Reads keypair from a file
-    pub fn read_json_file<F: AsRef<Path>>(path: F) -> Result<Self, Box<dyn std::error::Error>> {
-        let mut file = File::open(path.as_ref())?;
-        Self::read_json(&mut file)
+    pub fn read_json_file<F: AsRef<Path>>(path: F) -> Result<Self, Box<dyn error::Error>> {
+        Self::read_from_file(path)
     }
 
     /// Writes to a `Write` implementer with JSON-encoding
-    pub fn write_json<W: Write>(
-        &self,
-        writer: &mut W,
-    ) -> Result<String, Box<dyn std::error::Error>> {
+    pub fn write_json<W: Write>(&self, writer: &mut W) -> Result<String, Box<dyn error::Error>> {
         let bytes = self.to_bytes();
         let json = serde_json::to_string(&bytes.to_vec())?;
         writer.write_all(&json.clone().into_bytes())?;
@@ -229,29 +229,40 @@ impl ElGamalKeypair {
         &self,
         outfile: F,
     ) -> Result<String, Box<dyn std::error::Error>> {
-        let outfile = outfile.as_ref();
+        self.write_to_file(outfile)
+    }
+}
 
-        if let Some(outdir) = outfile.parent() {
-            fs::create_dir_all(outdir)?;
-        }
+impl EncodableKey for ElGamalKeypair {
+    fn read<R: Read>(reader: &mut R) -> Result<Self, Box<dyn error::Error>> {
+        Self::read_json(reader)
+    }
 
-        let mut f = {
-            #[cfg(not(unix))]
-            {
-                OpenOptions::new()
-            }
-            #[cfg(unix)]
-            {
-                use std::os::unix::fs::OpenOptionsExt;
-                OpenOptions::new().mode(0o600)
-            }
-        }
-        .write(true)
-        .truncate(true)
-        .create(true)
-        .open(outfile)?;
+    fn write<W: Write>(&self, writer: &mut W) -> Result<String, Box<dyn error::Error>> {
+        self.write_json(writer)
+    }
 
-        self.write_json(&mut f)
+    fn from_seed(seed: &[u8]) -> Result<Self, Box<dyn error::Error>> {
+        let secret = ElGamalSecretKey::from_seed(seed)?;
+        let public = ElGamalPubkey::new(&secret);
+        Ok(ElGamalKeypair { public, secret })
+    }
+
+    fn from_seed_and_derivation_path(
+        _seed: &[u8],
+        _derivation_path: Option<DerivationPath>,
+    ) -> Result<Self, Box<dyn error::Error>> {
+        Err(ElGamalError::DerivationMethodNotSupported.into())
+    }
+
+    fn from_seed_phrase_and_passphrase(
+        seed_phrase: &str,
+        passphrase: &str,
+    ) -> Result<Self, Box<dyn error::Error>> {
+        Self::from_seed(&generate_seed_from_seed_phrase_and_passphrase(
+            seed_phrase,
+            passphrase,
+        ))
     }
 }
 
@@ -356,6 +367,16 @@ impl ElGamalSecretKey {
         ElGamalSecretKey(Scalar::random(&mut OsRng))
     }
 
+    /// Derive an ElGamal secret key from an entropy seed.
+    pub fn from_seed(seed: &[u8]) -> Result<Self, Box<dyn error::Error>> {
+        const MINIMUM_SEED_LEN: usize = 32;
+
+        if seed.len() < MINIMUM_SEED_LEN {
+            return Err("Seed is too short".into());
+        }
+        Ok(ElGamalSecretKey(Scalar::hash_from_bytes::<Sha3_512>(seed)))
+    }
+
     pub fn get_scalar(&self) -> &Scalar {
         &self.0
     }
@@ -622,6 +643,7 @@ mod tests {
         super::*,
         crate::encryption::pedersen::Pedersen,
         solana_sdk::{signature::Keypair, signer::null_signer::NullSigner},
+        std::fs::{self, File},
     };
 
     #[test]

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

@@ -8,3 +8,18 @@ pub enum DiscreteLogError {
     #[error("discrete log batch size too large")]
     DiscreteLogBatchSize,
 }
+
+#[derive(Error, Clone, Debug, Eq, PartialEq)]
+pub enum ElGamalError {
+    #[error("key derivation method not supported")]
+    DerivationMethodNotSupported,
+}
+
+#[derive(Error, Clone, Debug, Eq, PartialEq)]
+pub enum AuthenticatedEncryptionError {
+    #[error("key derivation method not supported")]
+    DerivationMethodNotSupported,
+
+    #[error("pubkey does not exist")]
+    PubkeyDoesNotExist,
+}