From e0c168ef3fd4f9feadf7bb51981f286f23f4e73c Mon Sep 17 00:00:00 2001
From: Sam Kim <skim13@cs.stanford.edu>
Date: Fri, 8 Oct 2021 09:12:54 -0400
Subject: [PATCH] add aes encryption

---
 zk-token-sdk/Cargo.toml            |  1 +
 zk-token-sdk/src/encryption/aes.rs | 45 ++++++++++++++++++++++++++----
 zk-token-sdk/src/encryption/mod.rs |  1 +
 3 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/zk-token-sdk/Cargo.toml b/zk-token-sdk/Cargo.toml
index f4808ab998..57390397d3 100644
--- a/zk-token-sdk/Cargo.toml
+++ b/zk-token-sdk/Cargo.toml
@@ -16,6 +16,7 @@ num-traits = "0.2"
 solana-program = "=1.7.15"
 
 [target.'cfg(not(target_arch = "bpf"))'.dependencies]
+aes = "0.7.5"
 arrayref = "0.3.6"
 bincode = "1"
 byteorder = "1"
diff --git a/zk-token-sdk/src/encryption/aes.rs b/zk-token-sdk/src/encryption/aes.rs
index bb25931c15..00b8a3d9b2 100644
--- a/zk-token-sdk/src/encryption/aes.rs
+++ b/zk-token-sdk/src/encryption/aes.rs
@@ -1,30 +1,65 @@
+#[cfg(not(target_arch = "bpf"))]
+use rand::{rngs::OsRng, Rng};
 
+use aes::cipher::{BlockDecrypt, BlockEncrypt, NewBlockCipher};
+use aes::{Aes128, Block};
+
+use arrayref::array_ref;
 
 pub struct AES;
 impl AES {
     pub fn new() -> AESKey {
-        AESKey
+        let random_bytes = OsRng.gen::<[u8; 16]>();
+        AESKey(random_bytes)
     }
 
     pub fn encrypt(sk: &AESKey, amount: u64) -> AESCiphertext {
-        AESCiphertext
+        let amount_bytes = amount.to_le_bytes();
+
+        let mut aes_block: Block = [0_u8; 16].into();
+        aes_block[..8].copy_from_slice(&amount_bytes);
+
+        Aes128::new(&sk.0.into()).encrypt_block(&mut aes_block);
+        AESCiphertext(aes_block.into())
     }
 
     pub fn decrypt(sk: &AESKey, ct: &AESCiphertext) -> u64 {
-        0_u64
+        let mut aes_block: Block = ct.0.into();
+        Aes128::new(&sk.0.into()).decrypt_block(&mut aes_block);
+
+        let amount_bytes = array_ref![aes_block[..8], 0, 8];
+        u64::from_le_bytes(*amount_bytes)
     }
 }
 
-pub struct AESKey;
+#[derive(Debug)]
+pub struct AESKey([u8; 16]);
 impl AESKey {
     pub fn encrypt(&self, amount: u64) -> AESCiphertext {
         AES::encrypt(self, amount)
     }
 }
 
-pub struct AESCiphertext;
+#[derive(Debug)]
+pub struct AESCiphertext([u8; 16]);
 impl AESCiphertext {
     pub fn decrypt(&self, sk: &AESKey) -> u64 {
         AES::decrypt(sk, self)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_aes_encrypt_decrypt_correctness() {
+        let sk = AES::new();
+        let amount = 55;
+
+        let ct = sk.encrypt(amount);
+        let decrypted_amount = ct.decrypt(&sk);
+
+        assert_eq!(amount, decrypted_amount);
+    }
+}
diff --git a/zk-token-sdk/src/encryption/mod.rs b/zk-token-sdk/src/encryption/mod.rs
index ea34448d31..5605bdb773 100644
--- a/zk-token-sdk/src/encryption/mod.rs
+++ b/zk-token-sdk/src/encryption/mod.rs
@@ -1,3 +1,4 @@
+pub mod aes;
 pub mod discrete_log;
 pub mod elgamal;
 pub mod pedersen;