clean up authenticated encryption implementation and also rename aes to auth_encryption

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"

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;
+struct AuthenticatedEncryption;
-impl Aes {
+impl AuthenticatedEncryption {
     #[cfg(not(target_arch = "bpf"))]
     #[allow(clippy::new_ret_no_self)]
-    fn keygen<T: RngCore + CryptoRng>(rng: &mut T) -> AesKey {
+    fn keygen<T: RngCore + CryptoRng>(rng: &mut T) -> AeKey {
-        let random_bytes = rng.gen::<[u8; 16]>();
+        AeKey(rng.gen::<[u8; 16]>())
-        AesKey(random_bytes)
     }
 
     #[cfg(not(target_arch = "bpf"))]
-    fn encrypt(sk: &AesKey, amount: u64) -> AesCiphertext {
+    fn encrypt(key: &AeKey, balance: u64) -> AeCiphertext {
-        let plaintext = amount.to_le_bytes();
+        let mut plaintext = balance.to_le_bytes();
-        let nonce = OsRng.gen::<[u8; 12]>();
+        let nonce: Nonce = OsRng.gen::<[u8; 12]>();
 
-        // TODO: it seems like encryption cannot fail, but will need to double check
+        // The balance and the nonce have fixed length and therefore, encryption should not fail.
-        let ciphertext = Aes128Gcm::new(&sk.0.into())
+        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<u64> {
+    fn decrypt(key: &AeKey, ct: &AeCiphertext) -> Option<u64> {
         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]);
+pub struct AeKey([u8; 16]);
-impl AesKey {
+impl AeKey {
     pub fn new(signer: &dyn Signer, address: &Pubkey) -> Result<Self, SignerError> {
         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<T: RngCore + CryptoRng>(rng: &mut T) -> Self {
-        Aes::keygen(rng)
+        AuthenticatedEncryption::keygen(rng)
     }
 
-    pub fn encrypt(&self, amount: u64) -> AesCiphertext {
+    pub fn encrypt(&self, amount: u64) -> AeCiphertext {
-        Aes::encrypt(self, amount)
+        AuthenticatedEncryption::encrypt(self, amount)
     }
 
-    pub fn decrypt(&self, ct: &AesCiphertext) -> Option<u64> {
+    pub fn decrypt(&self, ct: &AeCiphertext) -> Option<u64> {
-        Aes::decrypt(self, ct)
+        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 struct AeCiphertext {
-    pub nonce: [u8; 12],
+    pub nonce: Nonce,
-    pub ciphertext: [u8; 24],
+    pub ciphertext: Ciphertext,
 }
-impl AesCiphertext {
+impl AeCiphertext {
-    pub fn decrypt(&self, key: &AesKey) -> Option<u64> {
+    pub fn decrypt(&self, key: &AeKey) -> Option<u64> {
-        Aes::decrypt(key, self)
+        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<AesCiphertext> {
+    pub fn from_bytes(bytes: &[u8]) -> Option<AeCiphertext> {
         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,
+            AeKey::new(&keypair1, &Pubkey::default()).unwrap().0,
-            AesKey::new(&keypair2, &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());
     }
 }

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;

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<AesCiphertext> for pod::AesCiphertext {
+    impl From<AeCiphertext> for pod::AeCiphertext {
-        fn from(ct: AesCiphertext) -> Self {
+        fn from(ct: AeCiphertext) -> Self {
             Self(ct.to_bytes())
         }
     }
 
-    impl TryFrom<pod::AesCiphertext> for AesCiphertext {
+    impl TryFrom<pod::AeCiphertext> for AeCiphertext {
         type Error = ProofError;
 
-        fn try_from(ct: pod::AesCiphertext) -> Result<Self, Self::Error> {
+        fn try_from(ct: pod::AeCiphertext) -> Result<Self, Self::Error> {
             Self::from_bytes(&ct.0).ok_or(ProofError::InconsistentCTData)
         }
     }

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 Zeroable for AeCiphertext {}
-unsafe impl Pod for AesCiphertext {}
+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)
     }