Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[zk-token-sdk] Implement FromStr for ElGamalPubkey, ElGamalCiphertext, and AeCiphertext #130

Merged
merged 14 commits into from
Mar 15, 2024
Merged
2 changes: 1 addition & 1 deletion zk-token-sdk/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand All @@ -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]
Expand Down
27 changes: 26 additions & 1 deletion zk-token-sdk/src/zk_token_elgamal/pod/auth_encryption.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand All @@ -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()
Expand All @@ -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);
}
}
49 changes: 48 additions & 1 deletion zk-token-sdk/src/zk_token_elgamal/pod/elgamal.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand All @@ -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)]
Expand All @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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);
}
}
33 changes: 33 additions & 0 deletions zk-token-sdk/src/zk_token_elgamal/pod/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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]);
Expand Down Expand Up @@ -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;
Loading