From 57103c515bf4279db5a680f8473fc95863d008dc Mon Sep 17 00:00:00 2001
From: Sam Kim <skim13@cs.stanford.edu>
Date: Mon, 11 Oct 2021 18:43:10 -0400
Subject: [PATCH] update applying pending balance for aes ciphertext

---
 zk-token-sdk/src/encryption/aes.rs           |  2 +-
 zk-token-sdk/src/zk_token_elgamal/convert.rs | 26 +++++++++++++++++++-
 zk-token-sdk/src/zk_token_elgamal/pod.rs     | 15 +++++++++++
 3 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/zk-token-sdk/src/encryption/aes.rs b/zk-token-sdk/src/encryption/aes.rs
index df620c5cb6..0541b67470 100644
--- a/zk-token-sdk/src/encryption/aes.rs
+++ b/zk-token-sdk/src/encryption/aes.rs
@@ -47,7 +47,7 @@ impl AESKey {
 }
 
 #[derive(Debug)]
-pub struct AESCiphertext([u8; 16]);
+pub struct AESCiphertext(pub [u8; 16]);
 impl AESCiphertext {
     pub fn decrypt(&self, sk: &AESKey) -> u64 {
         AES::decrypt(sk, self)
diff --git a/zk-token-sdk/src/zk_token_elgamal/convert.rs b/zk-token-sdk/src/zk_token_elgamal/convert.rs
index e773c72af5..840b9c02fc 100644
--- a/zk-token-sdk/src/zk_token_elgamal/convert.rs
+++ b/zk-token-sdk/src/zk_token_elgamal/convert.rs
@@ -22,7 +22,7 @@ mod target_arch {
             range_proof::RangeProof,
         },
         curve25519_dalek::{ristretto::CompressedRistretto, scalar::Scalar},
-        std::convert::TryFrom,
+        std::convert::{TryFrom, TryInto},
     };
 
     impl From<Scalar> for pod::Scalar {
@@ -136,6 +136,30 @@ mod target_arch {
         }
     }
 
+    impl From<Option<AESCiphertext>> for pod::OptionAESCiphertext {
+        fn from(ct: Option<AESCiphertext>) -> Self {
+            let mut buf = [0_u8; 17];
+            match ct {
+                Some(ct) => {
+                    buf[0] = 1_u8;
+                    buf[1..].copy_from_slice(&ct.0);
+                    Self(buf)
+                },
+                None => Self(buf),
+            }
+        }
+    }
+
+    impl From<pod::OptionAESCiphertext> for Option<AESCiphertext> {
+        fn from(ct: pod::OptionAESCiphertext) -> Self {
+            if ct.0[0] == 0 {
+                None
+            } else {
+                Some(AESCiphertext(ct.0[1..17].try_into().unwrap()))
+            }
+        }
+    }
+
     impl TryFrom<RangeProof> for pod::RangeProof64 {
         type Error = ProofError;
 
diff --git a/zk-token-sdk/src/zk_token_elgamal/pod.rs b/zk-token-sdk/src/zk_token_elgamal/pod.rs
index 2ae931c611..780a17dd69 100644
--- a/zk-token-sdk/src/zk_token_elgamal/pod.rs
+++ b/zk-token-sdk/src/zk_token_elgamal/pod.rs
@@ -80,3 +80,18 @@ impl fmt::Debug for AESCiphertext {
         write!(f, "{:?}", self.0)
     }
 }
+
+/// Temporary serialization of Option<AESCiphertext>
+#[derive(Clone, Copy, Pod, Zeroable, PartialEq)]
+#[repr(transparent)]
+pub struct OptionAESCiphertext(pub [u8; 17]);
+
+impl fmt::Debug for OptionAESCiphertext {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        if self.0[0] == 1_u8 {
+            write!(f, "Some({:?})", &self.0[1..17])
+        } else {
+            write!(f, "None")
+        }
+    }
+}