From 2359150b9c39d19d966ec2fa1ba96c320973efc1 Mon Sep 17 00:00:00 2001
From: Sam Kim <skim13@cs.stanford.edu>
Date: Thu, 14 Oct 2021 09:12:30 -0400
Subject: [PATCH] incorporate aes ciphertext for zk-proof instructions

---
 zk-token-sdk/src/instruction/transfer.rs | 9 +++++++++
 zk-token-sdk/src/instruction/withdraw.rs | 8 ++++++++
 2 files changed, 17 insertions(+)

diff --git a/zk-token-sdk/src/instruction/transfer.rs b/zk-token-sdk/src/instruction/transfer.rs
index 934813574a..d36d5195c5 100644
--- a/zk-token-sdk/src/instruction/transfer.rs
+++ b/zk-token-sdk/src/instruction/transfer.rs
@@ -6,6 +6,7 @@ use {
 use {
     crate::{
         encryption::{
+            aes::AESCiphertext,
             elgamal::{ElGamalCiphertext, ElGamalPubkey, ElGamalSecretKey},
             pedersen::{
                 Pedersen, PedersenBase, PedersenCommitment, PedersenDecryptHandle, PedersenOpening,
@@ -46,10 +47,14 @@ pub struct TransferData {
 
     /// Zero-knowledge proofs for Transfer
     pub proof: TransferProof,
+
+    /// The new decryptable balance
+    pub aes_ciphertext: pod::OptionAESCiphertext, // 17 bytes
 }
 
 #[cfg(not(target_arch = "bpf"))]
 impl TransferData {
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         transfer_amount: u64,
         spendable_balance: u64,
@@ -58,6 +63,7 @@ impl TransferData {
         source_sk: &ElGamalSecretKey,
         dest_pk: ElGamalPubkey,
         auditor_pk: ElGamalPubkey,
+        aes_ciphertext: Option<AESCiphertext>,
     ) -> Self {
         // split and encrypt transfer amount
         //
@@ -136,6 +142,7 @@ impl TransferData {
             new_spendable_ct: new_spendable_ct.into(),
             transfer_public_keys,
             proof,
+            aes_ciphertext: aes_ciphertext.into(),
         }
     }
 
@@ -493,6 +500,7 @@ mod test {
             &source_sk,
             dest_pk,
             auditor_pk,
+            None,
         );
 
         assert!(transfer_data.verify().is_ok());
@@ -527,6 +535,7 @@ mod test {
             &source_sk,
             dest_pk,
             auditor_pk,
+            None,
         );
 
         let decryption_data = decode_u32_precomputation_for_G();
diff --git a/zk-token-sdk/src/instruction/withdraw.rs b/zk-token-sdk/src/instruction/withdraw.rs
index 17f708357c..1f6985202d 100644
--- a/zk-token-sdk/src/instruction/withdraw.rs
+++ b/zk-token-sdk/src/instruction/withdraw.rs
@@ -8,6 +8,7 @@ use {
         encryption::{
             elgamal::{ElGamalCiphertext, ElGamalPubkey, ElGamalSecretKey},
             pedersen::{PedersenBase, PedersenOpening},
+            aes::AESCiphertext,
         },
         errors::ProofError,
         instruction::Verifiable,
@@ -36,6 +37,9 @@ pub struct WithdrawData {
 
     /// Proof that the account is solvent
     pub proof: WithdrawProof, // 736 bytes
+
+    /// The new decryptable balance component
+    pub aes_ciphertext: pod::OptionAESCiphertext, // 17 bytes
 }
 
 impl WithdrawData {
@@ -46,6 +50,7 @@ impl WithdrawData {
         source_sk: &ElGamalSecretKey,
         current_balance: u64,
         current_balance_ct: ElGamalCiphertext,
+        aes_ciphertext: Option<AESCiphertext>,
     ) -> Self {
         // subtract withdraw amount from current balance
         //
@@ -62,6 +67,7 @@ impl WithdrawData {
         Self {
             final_balance_ct: final_balance_ct.into(),
             proof,
+            aes_ciphertext: aes_ciphertext.into(),
         }
     }
 }
@@ -192,6 +198,7 @@ mod test {
             &secret,
             current_balance,
             current_balance_ct,
+            None,
         );
         assert!(data.verify().is_ok());
 
@@ -203,6 +210,7 @@ mod test {
             &secret,
             wrong_balance,
             current_balance_ct,
+            None,
         );
         assert!(data.verify().is_err());