From 09ad0d9075ce04a0a23eeb6420b3701426cc5c55 Mon Sep 17 00:00:00 2001
From: samkim-crypto <skim13@cs.stanford.edu>
Date: Sun, 24 Sep 2023 12:46:50 -0700
Subject: [PATCH] hard-code elgamal keypair support

---
 Cargo.lock                          |  1 +
 token/cli/Cargo.toml                | 21 +++++++----
 token/cli/src/encryption_keypair.rs | 55 +++++++++++++++++++++++++++++
 3 files changed, 71 insertions(+), 6 deletions(-)
 create mode 100644 token/cli/src/encryption_keypair.rs

diff --git a/Cargo.lock b/Cargo.lock
index 7747a849498..6c808632b3c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -7054,6 +7054,7 @@ name = "spl-token-cli"
 version = "3.1.1"
 dependencies = [
  "assert_cmd",
+ "base64 0.21.4",
  "clap 2.34.0",
  "console",
  "serde",
diff --git a/token/cli/Cargo.toml b/token/cli/Cargo.toml
index 9e553511e5a..a325819858c 100644
--- a/token/cli/Cargo.toml
+++ b/token/cli/Cargo.toml
@@ -12,6 +12,7 @@ version = "3.1.1"
 walkdir = "2"
 
 [dependencies]
+base64 = "0.21.4"
 clap = "2.33.3"
 console = "0.15.7"
 serde = "1.0.188"
@@ -26,12 +27,20 @@ solana-logger = "=1.16.13"
 solana-remote-wallet = "=1.16.13"
 solana-sdk = "=1.16.13"
 solana-transaction-status = "=1.16.13"
-spl-token = { version = "4.0", path="../program", features = [ "no-entrypoint" ] }
-spl-token-2022 = { version = "0.8", path="../program-2022", features = [ "no-entrypoint" ] }
-spl-token-client = { version = "0.6", path="../client" }
-spl-token-metadata-interface = { version = "0.2", path="../../token-metadata/interface" }
-spl-associated-token-account = { version = "2.0", path="../../associated-token-account/program", features = [ "no-entrypoint" ] }
-spl-memo = { version = "4.0.0", path="../../memo/program", features = ["no-entrypoint"] }
+spl-token = { version = "4.0", path = "../program", features = [
+  "no-entrypoint",
+] }
+spl-token-2022 = { version = "0.8", path = "../program-2022", features = [
+  "no-entrypoint",
+] }
+spl-token-client = { version = "0.6", path = "../client" }
+spl-token-metadata-interface = { version = "0.2", path = "../../token-metadata/interface" }
+spl-associated-token-account = { version = "2.0", path = "../../associated-token-account/program", features = [
+  "no-entrypoint",
+] }
+spl-memo = { version = "4.0.0", path = "../../memo/program", features = [
+  "no-entrypoint",
+] }
 strum = "0.25"
 strum_macros = "0.25"
 tokio = "1.14"
diff --git a/token/cli/src/encryption_keypair.rs b/token/cli/src/encryption_keypair.rs
new file mode 100644
index 00000000000..c0b4c534dc4
--- /dev/null
+++ b/token/cli/src/encryption_keypair.rs
@@ -0,0 +1,55 @@
+use {
+    base64::{prelude::BASE64_STANDARD, Engine},
+    clap::ArgMatches,
+    spl_token_2022::solana_zk_token_sdk::{
+        encryption::elgamal::{ElGamalKeypair, ElGamalPubkey},
+        zk_token_elgamal::pod::ElGamalPubkey as PodElGamalPubkey,
+    },
+};
+
+const ELGAMAL_PUBKEY_MAX_BASE64_LEN: usize = 44;
+
+pub(crate) fn elgamal_pubkey_or_none(
+    matches: &ArgMatches,
+    name: &str,
+) -> Result<Option<PodElGamalPubkey>, String> {
+    let arg_str = matches.value_of(name).unwrap();
+    if arg_str == "none" {
+        return Ok(None);
+    }
+    elgamal_pubkey_of(matches, name).map(|pubkey| Some(pubkey))
+}
+
+pub(crate) fn elgamal_pubkey_of(
+    matches: &ArgMatches,
+    name: &str,
+) -> Result<PodElGamalPubkey, String> {
+    if let Ok(keypair) = elgamal_keypair_of(matches, name) {
+        let elgamal_pubkey = (*keypair.pubkey()).into();
+        Ok(elgamal_pubkey)
+    } else {
+        let arg_str = matches.value_of(name).unwrap();
+        if let Some(pubkey) = elgamal_pubkey_from_str(arg_str) {
+            Ok(pubkey)
+        } else {
+            Err("failed to read ElGamal pubkey".to_string())
+        }
+    }
+}
+
+pub(crate) fn elgamal_keypair_of(
+    matches: &ArgMatches,
+    name: &str,
+) -> Result<ElGamalKeypair, String> {
+    let path = matches.value_of(name).unwrap();
+    ElGamalKeypair::read_json_file(path).map_err(|e| e.to_string())
+}
+
+fn elgamal_pubkey_from_str(s: &str) -> Option<PodElGamalPubkey> {
+    if s.len() > ELGAMAL_PUBKEY_MAX_BASE64_LEN {
+        return None;
+    }
+    let pubkey_vec = BASE64_STANDARD.decode(s).ok()?;
+    let elgamal_pubkey = ElGamalPubkey::from_bytes(&pubkey_vec)?;
+    Some(elgamal_pubkey.into())
+}