Adjust ElGamal::new() signature

zk-token-sdk/src/instruction/close_account.rs

@@ -152,10 +152,10 @@ mod test {
 
     #[test]
     fn test_close_account_correctness() {
-        let (source_pk, source_sk) = ElGamal::new();
+        let source = ElGamal::new();
 
         // invalid ciphertexts
-        let balance = source_pk.encrypt(0_u64);
+        let balance = source.pk.encrypt(0_u64);
 
         let zeroed_comm = Pedersen::with(0_u64, &PedersenOpening::default());
         let handle = balance.decrypt_handle;
@@ -165,7 +165,7 @@ mod test {
             decrypt_handle: handle,
         };
 
-        let proof = CloseAccountProof::new(&source_sk, &zeroed_comm_ciphertext);
+        let proof = CloseAccountProof::new(&source.sk, &zeroed_comm_ciphertext);
         assert!(proof.verify(&zeroed_comm_ciphertext).is_err());
 
         let zeroed_handle_ciphertext = ElGamalCiphertext {
@@ -173,24 +173,24 @@ mod test {
             decrypt_handle: PedersenDecryptHandle::default(),
         };
 
-        let proof = CloseAccountProof::new(&source_sk, &zeroed_handle_ciphertext);
+        let proof = CloseAccountProof::new(&source.sk, &zeroed_handle_ciphertext);
         assert!(proof.verify(&zeroed_handle_ciphertext).is_err());
 
         // valid ciphertext, but encryption of non-zero amount
-        let balance = source_pk.encrypt(55_u64);
+        let balance = source.pk.encrypt(55_u64);
 
-        let proof = CloseAccountProof::new(&source_sk, &balance);
+        let proof = CloseAccountProof::new(&source.sk, &balance);
         assert!(proof.verify(&balance).is_err());
 
         // all-zeroed ciphertext interpretted as a valid encryption of zero
         let zeroed_ct: ElGamalCiphertext = pod::ElGamalCiphertext::zeroed().try_into().unwrap();
-        let proof = CloseAccountProof::new(&source_sk, &zeroed_ct);
+        let proof = CloseAccountProof::new(&source.sk, &zeroed_ct);
         assert!(proof.verify(&zeroed_ct).is_ok());
 
         // general case: valid encryption of zero
-        let balance = source_pk.encrypt(0_u64);
+        let balance = source.pk.encrypt(0_u64);
 
-        let proof = CloseAccountProof::new(&source_sk, &balance);
+        let proof = CloseAccountProof::new(&source.sk, &balance);
         assert!(proof.verify(&balance).is_ok());
     }
 }

zk-token-sdk/src/instruction/transfer.rs

@@ -468,9 +468,12 @@ mod test {
     #[test]
     fn test_transfer_correctness() {
         // ElGamal keys for source, destination, and auditor accounts
-        let (source_pk, source_sk) = ElGamal::new();
-        let (dest_pk, _) = ElGamal::new();
-        let (auditor_pk, _) = ElGamal::new();
+        let ElGamal {
+            pk: source_pk,
+            sk: source_sk,
+        } = ElGamal::new();
+        let dest_pk = ElGamal::new().pk;
+        let auditor_pk = ElGamal::new().pk;
 
         // create source account spendable ciphertext
         let spendable_balance: u64 = 77;
@@ -496,9 +499,15 @@ mod test {
     #[test]
     fn test_source_dest_ciphertext() {
         // ElGamal keys for source, destination, and auditor accounts
-        let (source_pk, source_sk) = ElGamal::new();
-        let (dest_pk, dest_sk) = ElGamal::new();
-        let (auditor_pk, _) = ElGamal::new();
+        let ElGamal {
+            pk: source_pk,
+            sk: source_sk,
+        } = ElGamal::new();
+        let ElGamal {
+            pk: dest_pk,
+            sk: dest_sk,
+        } = ElGamal::new();
+        let auditor_pk = ElGamal::new().pk;
 
         // create source account spendable ciphertext
         let spendable_balance: u64 = 77;

zk-token-sdk/src/instruction/update_account_pk.rs

@@ -219,50 +219,50 @@ mod test {
 
     #[test]
     fn test_update_account_public_key_general_cases() {
-        let (current_pk, current_sk) = ElGamal::new();
-        let (new_pk, new_sk) = ElGamal::new();
+        let current = ElGamal::new();
+        let new = ElGamal::new();
 
         // If current_ct and new_ct encrypt same values, then the proof verification should succeed
         let balance: u64 = 77;
-        let current_ct = current_pk.encrypt(balance);
-        let new_ct = new_pk.encrypt(balance);
+        let current_ct = current.pk.encrypt(balance);
+        let new_ct = new.pk.encrypt(balance);
 
-        let proof = UpdateAccountPkProof::new(balance, &current_sk, &new_sk, &current_ct, &new_ct);
+        let proof = UpdateAccountPkProof::new(balance, &current.sk, &new.sk, &current_ct, &new_ct);
         assert!(proof.verify(&current_ct, &new_ct).is_ok());
 
         // If current_ct and new_ct encrypt different values, then the proof verification should fail
-        let new_ct = new_pk.encrypt(55_u64);
+        let new_ct = new.pk.encrypt(55_u64);
 
-        let proof = UpdateAccountPkProof::new(balance, &current_sk, &new_sk, &current_ct, &new_ct);
+        let proof = UpdateAccountPkProof::new(balance, &current.sk, &new.sk, &current_ct, &new_ct);
         assert!(proof.verify(&current_ct, &new_ct).is_err());
     }
 
     #[test]
     fn test_update_account_public_key_zeroed_ciphertexts() {
-        let (current_pk, current_sk) = ElGamal::new();
-        let (new_pk, new_sk) = ElGamal::new();
+        let current = ElGamal::new();
+        let new = ElGamal::new();
 
         // A zeroed cipehrtext should be considered as an account balance of 0
         let balance: u64 = 0;
         let zeroed_ct_as_current_ct: ElGamalCiphertext =
             pod::ElGamalCiphertext::zeroed().try_into().unwrap();
-        let new_ct: ElGamalCiphertext = new_pk.encrypt(balance);
+        let new_ct: ElGamalCiphertext = new.pk.encrypt(balance);
         let proof = UpdateAccountPkProof::new(
             balance,
-            &current_sk,
-            &new_sk,
+            &current.sk,
+            &new.sk,
             &zeroed_ct_as_current_ct,
             &new_ct,
         );
         assert!(proof.verify(&zeroed_ct_as_current_ct, &new_ct).is_ok());
 
-        let current_ct = current_pk.encrypt(balance);
+        let current_ct = current.pk.encrypt(balance);
         let zeroed_ct_as_new_ct: ElGamalCiphertext =
             pod::ElGamalCiphertext::zeroed().try_into().unwrap();
         let proof = UpdateAccountPkProof::new(
             balance,
-            &current_sk,
-            &new_sk,
+            &current.sk,
+            &new.sk,
             &current_ct,
             &zeroed_ct_as_new_ct,
         );
@@ -274,8 +274,8 @@ mod test {
             pod::ElGamalCiphertext::zeroed().try_into().unwrap();
         let proof = UpdateAccountPkProof::new(
             balance,
-            &current_sk,
-            &new_sk,
+            &current.sk,
+            &new.sk,
             &zeroed_ct_as_current_ct,
             &zeroed_ct_as_new_ct,
         );
@@ -286,11 +286,11 @@ mod test {
 
     #[test]
     fn test_update_account_public_key_partially_zeroed_ciphertexts() {
-        let (current_pk, current_sk) = ElGamal::new();
-        let (new_pk, new_sk) = ElGamal::new();
+        let current = ElGamal::new();
+        let new = ElGamal::new();
 
         let balance = 0_u64;
-        let balance_ciphertext = new_pk.encrypt(balance);
+        let balance_ciphertext = new.pk.encrypt(balance);
 
         let zeroed_comm = Pedersen::with(0_u64, &PedersenOpening::default());
         let handle = balance_ciphertext.decrypt_handle;
@@ -300,12 +300,12 @@ mod test {
             message_comm: zeroed_comm,
             decrypt_handle: handle,
         };
-        let new_ct: ElGamalCiphertext = new_pk.encrypt(balance);
+        let new_ct: ElGamalCiphertext = new.pk.encrypt(balance);
 
         let proof = UpdateAccountPkProof::new(
             balance,
-            &current_sk,
-            &new_sk,
+            &current.sk,
+            &new.sk,
             &zeroed_comm_ciphertext,
             &new_ct,
         );
@@ -318,8 +318,8 @@ mod test {
 
         let proof = UpdateAccountPkProof::new(
             balance,
-            &current_sk,
-            &new_sk,
+            &current.sk,
+            &new.sk,
             &zeroed_handle_ciphertext,
             &new_ct,
         );
@@ -330,12 +330,12 @@ mod test {
             message_comm: zeroed_comm,
             decrypt_handle: handle,
         };
-        let current_ct: ElGamalCiphertext = current_pk.encrypt(balance);
+        let current_ct: ElGamalCiphertext = current.pk.encrypt(balance);
 
         let proof = UpdateAccountPkProof::new(
             balance,
-            &current_sk,
-            &new_sk,
+            &current.sk,
+            &new.sk,
             &current_ct,
             &zeroed_comm_ciphertext,
         );
@@ -348,8 +348,8 @@ mod test {
 
         let proof = UpdateAccountPkProof::new(
             balance,
-            &current_sk,
-            &new_sk,
+            &current.sk,
+            &new.sk,
             &current_ct,
             &zeroed_handle_ciphertext,
         );

zk-token-sdk/src/instruction/withdraw.rs

@@ -179,17 +179,17 @@ mod test {
     #[ignore]
     fn test_withdraw_correctness() {
         // generate and verify proof for the proper setting
-        let (source_pk, source_sk) = ElGamal::new();
+        let ElGamal { pk, sk } = ElGamal::new();
 
         let current_balance: u64 = 77;
-        let current_balance_ct = source_pk.encrypt(current_balance);
+        let current_balance_ct = pk.encrypt(current_balance);
 
         let withdraw_amount: u64 = 55;
 
         let data = WithdrawData::new(
             withdraw_amount,
-            source_pk,
-            &source_sk,
+            pk,
+            &sk,
             current_balance,
             current_balance_ct,
         );
@@ -197,13 +197,7 @@ mod test {
 
         // generate and verify proof with wrong balance
         let wrong_balance: u64 = 99;
-        let data = WithdrawData::new(
-            withdraw_amount,
-            source_pk,
-            &source_sk,
-            wrong_balance,
-            current_balance_ct,
-        );
+        let data = WithdrawData::new(withdraw_amount, pk, &sk, wrong_balance, current_balance_ct);
         assert!(data.verify().is_err());
 
         // TODO: test for ciphertexts that encrypt numbers outside the 0, 2^64 range