From 9a43fbe3b23ac77fa7ff24ab4ec07a4416669926 Mon Sep 17 00:00:00 2001 From: Sam Kim Date: Mon, 13 Dec 2021 10:56:49 -0500 Subject: [PATCH] clean up authenticated encryption implementation and also rename aes to auth_encryption --- zk-token-sdk/Cargo.toml | 2 +- .../encryption/{aes.rs => auth_encryption.rs} | 87 +++++++++++-------- zk-token-sdk/src/encryption/mod.rs | 2 +- zk-token-sdk/src/zk_token_elgamal/convert.rs | 10 +-- zk-token-sdk/src/zk_token_elgamal/pod.rs | 12 +-- 5 files changed, 62 insertions(+), 51 deletions(-) rename zk-token-sdk/src/encryption/{aes.rs => auth_encryption.rs} (56%) diff --git a/zk-token-sdk/Cargo.toml b/zk-token-sdk/Cargo.toml index 58b4626d7..ef8ea2655 100644 --- a/zk-token-sdk/Cargo.toml +++ b/zk-token-sdk/Cargo.toml @@ -16,7 +16,7 @@ num-traits = "0.2" solana-program = "=1.7.15" [target.'cfg(not(target_arch = "bpf"))'.dependencies] -aes-gcm = "0.9.4" +aes-gcm-siv = "0.10.3" arrayref = "0.3.6" bincode = "1" byteorder = "1" diff --git a/zk-token-sdk/src/encryption/aes.rs b/zk-token-sdk/src/encryption/auth_encryption.rs similarity index 56% rename from zk-token-sdk/src/encryption/aes.rs rename to zk-token-sdk/src/encryption/auth_encryption.rs index f36762a5a..1027b2acd 100644 --- a/zk-token-sdk/src/encryption/aes.rs +++ b/zk-token-sdk/src/encryption/auth_encryption.rs @@ -1,6 +1,9 @@ #[cfg(not(target_arch = "bpf"))] use { - aes_gcm::{aead::Aead, Aes128Gcm, NewAead}, + aes_gcm_siv::{ + aead::{Aead, NewAead}, + Aes128GcmSiv, + }, rand::{rngs::OsRng, CryptoRng, Rng, RngCore}, }; use { @@ -16,35 +19,36 @@ use { zeroize::Zeroize, }; -struct Aes; -impl Aes { +struct AuthenticatedEncryption; +impl AuthenticatedEncryption { #[cfg(not(target_arch = "bpf"))] #[allow(clippy::new_ret_no_self)] - fn keygen(rng: &mut T) -> AesKey { - let random_bytes = rng.gen::<[u8; 16]>(); - AesKey(random_bytes) + fn keygen(rng: &mut T) -> AeKey { + AeKey(rng.gen::<[u8; 16]>()) } #[cfg(not(target_arch = "bpf"))] - fn encrypt(sk: &AesKey, amount: u64) -> AesCiphertext { - let plaintext = amount.to_le_bytes(); - let nonce = OsRng.gen::<[u8; 12]>(); + fn encrypt(key: &AeKey, balance: u64) -> AeCiphertext { + let mut plaintext = balance.to_le_bytes(); + let nonce: Nonce = OsRng.gen::<[u8; 12]>(); - // TODO: it seems like encryption cannot fail, but will need to double check - let ciphertext = Aes128Gcm::new(&sk.0.into()) + // The balance and the nonce have fixed length and therefore, encryption should not fail. + let ciphertext = Aes128GcmSiv::new(&key.0.into()) .encrypt(&nonce.into(), plaintext.as_ref()) - .unwrap(); + .expect("authenticated encryption"); - AesCiphertext { + plaintext.zeroize(); + + AeCiphertext { nonce, ciphertext: ciphertext.try_into().unwrap(), } } #[cfg(not(target_arch = "bpf"))] - fn decrypt(sk: &AesKey, ct: &AesCiphertext) -> Option { + fn decrypt(key: &AeKey, ct: &AeCiphertext) -> Option { let plaintext = - Aes128Gcm::new(&sk.0.into()).decrypt(&ct.nonce.into(), ct.ciphertext.as_ref()); + Aes128GcmSiv::new(&key.0.into()).decrypt(&ct.nonce.into(), ct.ciphertext.as_ref()); if let Ok(plaintext) = plaintext { let amount_bytes: [u8; 8] = plaintext.try_into().unwrap(); @@ -56,11 +60,11 @@ impl Aes { } #[derive(Debug, Zeroize)] -pub struct AesKey([u8; 16]); -impl AesKey { +pub struct AeKey([u8; 16]); +impl AeKey { pub fn new(signer: &dyn Signer, address: &Pubkey) -> Result { let message = Message::new( - &[Instruction::new_with_bytes(*address, b"AesKey", vec![])], + &[Instruction::new_with_bytes(*address, b"AeKey", vec![])], Some(&signer.try_pubkey()?), ); let signature = signer.try_sign_message(&message.serialize())?; @@ -70,31 +74,37 @@ impl AesKey { if signature == Signature::default() { Err(SignerError::Custom("Rejecting default signature".into())) } else { - Ok(AesKey(signature.as_ref()[..16].try_into().unwrap())) + Ok(AeKey(signature.as_ref()[..16].try_into().unwrap())) } } pub fn random(rng: &mut T) -> Self { - Aes::keygen(rng) + AuthenticatedEncryption::keygen(rng) } - pub fn encrypt(&self, amount: u64) -> AesCiphertext { - Aes::encrypt(self, amount) + pub fn encrypt(&self, amount: u64) -> AeCiphertext { + AuthenticatedEncryption::encrypt(self, amount) } - pub fn decrypt(&self, ct: &AesCiphertext) -> Option { - Aes::decrypt(self, ct) + pub fn decrypt(&self, ct: &AeCiphertext) -> Option { + AuthenticatedEncryption::decrypt(self, ct) } } +/// For the purpose of encrypting balances for ZK-Token accounts, the nonce and ciphertext sizes +/// should always be fixed. +pub type Nonce = [u8; 12]; +pub type Ciphertext = [u8; 24]; + +/// Authenticated encryption nonce and ciphertext #[derive(Debug)] -pub struct AesCiphertext { - pub nonce: [u8; 12], - pub ciphertext: [u8; 24], +pub struct AeCiphertext { + pub nonce: Nonce, + pub ciphertext: Ciphertext, } -impl AesCiphertext { - pub fn decrypt(&self, key: &AesKey) -> Option { - Aes::decrypt(key, self) +impl AeCiphertext { + pub fn decrypt(&self, key: &AeKey) -> Option { + AuthenticatedEncryption::decrypt(key, self) } pub fn to_bytes(&self) -> [u8; 36] { @@ -104,23 +114,24 @@ impl AesCiphertext { buf } - pub fn from_bytes(bytes: &[u8]) -> Option { + pub fn from_bytes(bytes: &[u8]) -> Option { if bytes.len() != 36 { return None; } let bytes = array_ref![bytes, 0, 36]; let (nonce, ciphertext) = array_refs![bytes, 12, 24]; - Some(AesCiphertext { + + Some(AeCiphertext { nonce: *nonce, ciphertext: *ciphertext, }) } } -impl Default for AesCiphertext { +impl Default for AeCiphertext { fn default() -> Self { - AesCiphertext { + AeCiphertext { nonce: [0_u8; 12], ciphertext: [0_u8; 24], } @@ -136,7 +147,7 @@ mod tests { #[test] fn test_aes_encrypt_decrypt_correctness() { - let key = AesKey::random(&mut OsRng); + let key = AeKey::random(&mut OsRng); let amount = 55; let ct = key.encrypt(amount); @@ -151,11 +162,11 @@ mod tests { let keypair2 = Keypair::new(); assert_ne!( - AesKey::new(&keypair1, &Pubkey::default()).unwrap().0, - AesKey::new(&keypair2, &Pubkey::default()).unwrap().0, + AeKey::new(&keypair1, &Pubkey::default()).unwrap().0, + AeKey::new(&keypair2, &Pubkey::default()).unwrap().0, ); let null_signer = NullSigner::new(&Pubkey::default()); - assert!(AesKey::new(&null_signer, &Pubkey::default()).is_err()); + assert!(AeKey::new(&null_signer, &Pubkey::default()).is_err()); } } diff --git a/zk-token-sdk/src/encryption/mod.rs b/zk-token-sdk/src/encryption/mod.rs index 5605bdb77..8792bd9a0 100644 --- a/zk-token-sdk/src/encryption/mod.rs +++ b/zk-token-sdk/src/encryption/mod.rs @@ -1,4 +1,4 @@ -pub mod aes; +pub mod auth_encryption; pub mod discrete_log; pub mod elgamal; pub mod pedersen; diff --git a/zk-token-sdk/src/zk_token_elgamal/convert.rs b/zk-token-sdk/src/zk_token_elgamal/convert.rs index 319365696..ff085e14f 100644 --- a/zk-token-sdk/src/zk_token_elgamal/convert.rs +++ b/zk-token-sdk/src/zk_token_elgamal/convert.rs @@ -16,7 +16,7 @@ mod target_arch { super::pod, crate::{ encryption::{ - aes::AesCiphertext, + auth_encryption::AeCiphertext, elgamal::{ElGamalCiphertext, ElGamalPubkey}, pedersen::{PedersenCommitment, PedersenDecryptHandle}, }, @@ -128,16 +128,16 @@ mod target_arch { } } - impl From for pod::AesCiphertext { - fn from(ct: AesCiphertext) -> Self { + impl From for pod::AeCiphertext { + fn from(ct: AeCiphertext) -> Self { Self(ct.to_bytes()) } } - impl TryFrom for AesCiphertext { + impl TryFrom for AeCiphertext { type Error = ProofError; - fn try_from(ct: pod::AesCiphertext) -> Result { + fn try_from(ct: pod::AeCiphertext) -> Result { Self::from_bytes(&ct.0).ok_or(ProofError::InconsistentCTData) } } diff --git a/zk-token-sdk/src/zk_token_elgamal/pod.rs b/zk-token-sdk/src/zk_token_elgamal/pod.rs index 6c62cfed6..06b708df6 100644 --- a/zk-token-sdk/src/zk_token_elgamal/pod.rs +++ b/zk-token-sdk/src/zk_token_elgamal/pod.rs @@ -89,17 +89,17 @@ pub struct RangeProof128(pub [u8; 736]); unsafe impl Zeroable for RangeProof128 {} unsafe impl Pod for RangeProof128 {} -/// Serialization for AesCiphertext +/// Serialization for AeCiphertext #[derive(Clone, Copy, PartialEq)] #[repr(transparent)] -pub struct AesCiphertext(pub [u8; 36]); +pub struct AeCiphertext(pub [u8; 36]); -// `AesCiphertext` is a Pod and Zeroable. +// `AeCiphertext` is a Pod and Zeroable. // Add the marker traits manually because `bytemuck` only adds them for some `u8` arrays -unsafe impl Zeroable for AesCiphertext {} -unsafe impl Pod for AesCiphertext {} +unsafe impl Zeroable for AeCiphertext {} +unsafe impl Pod for AeCiphertext {} -impl fmt::Debug for AesCiphertext { +impl fmt::Debug for AeCiphertext { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { write!(f, "{:?}", self.0) }