From fba70c8504984ca42d5f31b08c9873688999d884 Mon Sep 17 00:00:00 2001 From: samkim-crypto Date: Fri, 15 Mar 2024 09:09:41 +0900 Subject: [PATCH] [zk-token-sdk] Implement `FromStr` for `ElGamalPubkey`, `ElGamalCiphertext`, 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 * remove unnecessary imports * remove the need for `ParseError` dependency --------- Co-authored-by: Jon C --- zk-token-sdk/Cargo.toml | 2 +- .../zk_token_elgamal/pod/auth_encryption.rs | 27 +++++++++- .../src/zk_token_elgamal/pod/elgamal.rs | 49 ++++++++++++++++++- zk-token-sdk/src/zk_token_elgamal/pod/mod.rs | 33 +++++++++++++ 4 files changed, 108 insertions(+), 3 deletions(-) diff --git a/zk-token-sdk/Cargo.toml b/zk-token-sdk/Cargo.toml index e78e8c18f..4d4ff1a21 100644 --- 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 index 45534218f..9868e6e9f 100644 --- 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,7 +3,7 @@ #[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, }; @@ -11,6 +11,9 @@ use { /// 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 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 index 4473ab3ee..251729c08 100644 --- 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 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 for ElGamalPubkey { fn from(decoded_pubkey: decoded::ElGamalPubkey) -> Self { @@ -129,3 +147,32 @@ impl TryFrom 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 index 864fc7dde..d782672c8 100644 --- 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 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 { + 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;