[zk-token-sdk] Implement FromStr for ElGamalPubkey, `ElGamalCiphe…

…rtext`, and `AeCiphertext` (#130) * add `ParseError` in `zk-token-elgamal` * implement `FromStr` for `ElGamalPubkey` and `ElGamalCiphertext` * implement `FromStr` for `AeCiphertext` * fix target * cargo fmt * use constants for byte length check * make `FromStr` functions available on chain * use macros for the `FromStr` implementations * restrict `from_str` macro to `pub(crate)` * decode directly into array * cargo fmt * Apply suggestions from code review Co-authored-by: Jon C <[email protected]> * remove unnecessary imports * remove the need for `ParseError` dependency --------- Co-authored-by: Jon C <[email protected]>
anza-xyz · Mar 15, 2024 · fba70c8 · fba70c8
1 parent b3fd87f
commit fba70c8
Show file tree

Hide file tree

Showing 4 changed files with 108 additions and 3 deletions.
diff --git a/zk-token-sdk/Cargo.toml b/zk-token-sdk/Cargo.toml
@@ -15,6 +15,7 @@ bytemuck = { workspace = true, features = ["derive"] }
 num-derive = { workspace = true }
 num-traits = { workspace = true }
 solana-program = { workspace = true }
+thiserror = { workspace = true }
 
 [dev-dependencies]
 tiny-bip39 = { workspace = true }
@@ -34,7 +35,6 @@ serde_json = { workspace = true }
 sha3 = "0.9"
 solana-sdk = { workspace = true }
 subtle = { workspace = true }
-thiserror = { workspace = true }
 zeroize = { workspace = true, features = ["zeroize_derive"] }
 
 [lib]

diff --git a/zk-token-sdk/src/zk_token_elgamal/pod/auth_encryption.rs b/zk-token-sdk/src/zk_token_elgamal/pod/auth_encryption.rs
@@ -3,14 +3,17 @@
 #[cfg(not(target_os = "solana"))]
 use crate::encryption::auth_encryption::{self as decoded, AuthenticatedEncryptionError};
 use {
-    crate::zk_token_elgamal::pod::{Pod, Zeroable},
+    crate::zk_token_elgamal::pod::{impl_from_str, Pod, Zeroable},
     base64::{prelude::BASE64_STANDARD, Engine},
     std::fmt,
 };
 
 /// Byte length of an authenticated encryption ciphertext
 const AE_CIPHERTEXT_LEN: usize = 36;
 
+/// Maximum length of a base64 encoded authenticated encryption ciphertext
+const AE_CIPHERTEXT_MAX_BASE64_LEN: usize = 48;
+
 /// The `AeCiphertext` type as a `Pod`.
 #[derive(Clone, Copy, PartialEq, Eq)]
 #[repr(transparent)]
@@ -34,6 +37,12 @@ impl fmt::Display for AeCiphertext {
     }
 }
 
+impl_from_str!(
+    TYPE = AeCiphertext,
+    BYTES_LEN = AE_CIPHERTEXT_LEN,
+    BASE64_LEN = AE_CIPHERTEXT_MAX_BASE64_LEN
+);
+
 impl Default for AeCiphertext {
     fn default() -> Self {
         Self::zeroed()
@@ -55,3 +64,19 @@ impl TryFrom<AeCiphertext> for decoded::AeCiphertext {
         Self::from_bytes(&pod_ciphertext.0).ok_or(AuthenticatedEncryptionError::Deserialization)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use {super::*, crate::encryption::auth_encryption::AeKey, std::str::FromStr};
+
+    #[test]
+    fn ae_ciphertext_fromstr() {
+        let ae_key = AeKey::new_rand();
+        let expected_ae_ciphertext: AeCiphertext = ae_key.encrypt(0_u64).into();
+
+        let ae_ciphertext_base64_str = format!("{}", expected_ae_ciphertext);
+        let computed_ae_ciphertext = AeCiphertext::from_str(&ae_ciphertext_base64_str).unwrap();
+
+        assert_eq!(expected_ae_ciphertext, computed_ae_ciphertext);
+    }
+}
diff --git a/zk-token-sdk/src/zk_token_elgamal/pod/elgamal.rs b/zk-token-sdk/src/zk_token_elgamal/pod/elgamal.rs
@@ -7,7 +7,7 @@ use {
 };
 use {
     crate::{
-        zk_token_elgamal::pod::{pedersen::PEDERSEN_COMMITMENT_LEN, Pod, Zeroable},
+        zk_token_elgamal::pod::{impl_from_str, pedersen::PEDERSEN_COMMITMENT_LEN, Pod, Zeroable},
         RISTRETTO_POINT_LEN,
     },
     base64::{prelude::BASE64_STANDARD, Engine},
@@ -17,12 +17,18 @@ use {
 /// Byte length of an ElGamal public key
 const ELGAMAL_PUBKEY_LEN: usize = RISTRETTO_POINT_LEN;
 
+/// Maximum length of a base64 encoded ElGamal public key
+const ELGAMAL_PUBKEY_MAX_BASE64_LEN: usize = 44;
+
 /// Byte length of a decrypt handle
 pub(crate) const DECRYPT_HANDLE_LEN: usize = RISTRETTO_POINT_LEN;
 
 /// Byte length of an ElGamal ciphertext
 const ELGAMAL_CIPHERTEXT_LEN: usize = PEDERSEN_COMMITMENT_LEN + DECRYPT_HANDLE_LEN;
 
+/// Maximum length of a base64 encoded ElGamal ciphertext
+const ELGAMAL_CIPHERTEXT_MAX_BASE64_LEN: usize = 88;
+
 /// The `ElGamalCiphertext` type as a `Pod`.
 #[derive(Clone, Copy, Pod, Zeroable, PartialEq, Eq)]
 #[repr(transparent)]
@@ -46,6 +52,12 @@ impl Default for ElGamalCiphertext {
     }
 }
 
+impl_from_str!(
+    TYPE = ElGamalCiphertext,
+    BYTES_LEN = ELGAMAL_CIPHERTEXT_LEN,
+    BASE64_LEN = ELGAMAL_CIPHERTEXT_MAX_BASE64_LEN
+);
+
 #[cfg(not(target_os = "solana"))]
 impl From<decoded::ElGamalCiphertext> for ElGamalCiphertext {
     fn from(decoded_ciphertext: decoded::ElGamalCiphertext) -> Self {
@@ -79,6 +91,12 @@ impl fmt::Display for ElGamalPubkey {
     }
 }
 
+impl_from_str!(
+    TYPE = ElGamalPubkey,
+    BYTES_LEN = ELGAMAL_PUBKEY_LEN,
+    BASE64_LEN = ELGAMAL_PUBKEY_MAX_BASE64_LEN
+);
+
 #[cfg(not(target_os = "solana"))]
 impl From<decoded::ElGamalPubkey> for ElGamalPubkey {
     fn from(decoded_pubkey: decoded::ElGamalPubkey) -> Self {
@@ -129,3 +147,32 @@ impl TryFrom<DecryptHandle> for decoded::DecryptHandle {
         Self::from_bytes(&pod_handle.0).ok_or(ElGamalError::CiphertextDeserialization)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use {super::*, crate::encryption::elgamal::ElGamalKeypair, std::str::FromStr};
+
+    #[test]
+    fn elgamal_pubkey_fromstr() {
+        let elgamal_keypair = ElGamalKeypair::new_rand();
+        let expected_elgamal_pubkey: ElGamalPubkey = (*elgamal_keypair.pubkey()).into();
+
+        let elgamal_pubkey_base64_str = format!("{}", expected_elgamal_pubkey);
+        let computed_elgamal_pubkey = ElGamalPubkey::from_str(&elgamal_pubkey_base64_str).unwrap();
+
+        assert_eq!(expected_elgamal_pubkey, computed_elgamal_pubkey);
+    }
+
+    #[test]
+    fn elgamal_ciphertext_fromstr() {
+        let elgamal_keypair = ElGamalKeypair::new_rand();
+        let expected_elgamal_ciphertext: ElGamalCiphertext =
+            elgamal_keypair.pubkey().encrypt(0_u64).into();
+
+        let elgamal_ciphertext_base64_str = format!("{}", expected_elgamal_ciphertext);
+        let computed_elgamal_ciphertext =
+            ElGamalCiphertext::from_str(&elgamal_ciphertext_base64_str).unwrap();
+
+        assert_eq!(expected_elgamal_ciphertext, computed_elgamal_ciphertext);
+    }
+}
diff --git a/zk-token-sdk/src/zk_token_elgamal/pod/mod.rs b/zk-token-sdk/src/zk_token_elgamal/pod/mod.rs
@@ -10,6 +10,7 @@ use {
     crate::zk_token_proof_instruction::ProofType,
     num_traits::{FromPrimitive, ToPrimitive},
     solana_program::instruction::InstructionError,
+    thiserror::Error,
 };
 pub use {
     auth_encryption::AeCiphertext,
@@ -26,6 +27,14 @@ pub use {
     },
 };
 
+#[derive(Error, Debug, Clone, PartialEq, Eq)]
+pub enum ParseError {
+    #[error("String is the wrong size")]
+    WrongSize,
+    #[error("Invalid Base64 string")]
+    Invalid,
+}
+
 #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Pod, Zeroable)]
 #[repr(transparent)]
 pub struct PodU16([u8; 2]);
@@ -73,3 +82,27 @@ impl TryFrom<PodProofType> for ProofType {
 #[derive(Clone, Copy, Pod, Zeroable, PartialEq, Eq)]
 #[repr(transparent)]
 pub struct CompressedRistretto(pub [u8; 32]);
+
+macro_rules! impl_from_str {
+    (TYPE = $type:ident, BYTES_LEN = $bytes_len:expr, BASE64_LEN = $base64_len:expr) => {
+        impl std::str::FromStr for $type {
+            type Err = crate::zk_token_elgamal::pod::ParseError;
+
+            fn from_str(s: &str) -> Result<Self, Self::Err> {
+                if s.len() > $base64_len {
+                    return Err(Self::Err::WrongSize);
+                }
+                let mut bytes = [0u8; $bytes_len];
+                let decoded_len = BASE64_STANDARD
+                    .decode_slice(s, &mut bytes)
+                    .map_err(|_| Self::Err::Invalid)?;
+                if decoded_len != $bytes_len {
+                    Err(Self::Err::WrongSize)
+                } else {
+                    Ok($type(bytes))
+                }
+            }
+        }
+    };
+}
+pub(crate) use impl_from_str;