Additional tests for proof verification when ElGamal pubkey is zeroed (#24243)

* zk-token-sdk: add edge case tests for withdraw withheld proof

* zk-token-sdk: add test cases for proof verification when pubkeys are invalid

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

@@ -596,6 +596,40 @@ mod test {
         );
 
         assert!(transfer_data.is_err());
+
+        // Case 5: invalid destination or auditor pubkey
+        let spendable_balance: u64 = 0;
+        let spendable_ciphertext = source_keypair.public.encrypt(spendable_balance);
+
+        let transfer_amount: u64 = 0;
+
+        // destination pubkey invalid
+        let dest_pk = pod::ElGamalPubkey::zeroed().try_into().unwrap();
+        let auditor_pk = ElGamalKeypair::new_rand().public;
+
+        let transfer_data = TransferData::new(
+            transfer_amount,
+            (spendable_balance, &spendable_ciphertext),
+            &source_keypair,
+            (&dest_pk, &auditor_pk),
+        )
+        .unwrap();
+
+        assert!(transfer_data.verify().is_err());
+
+        // auditor pubkey invalid
+        let dest_pk = ElGamalKeypair::new_rand().public;
+        let auditor_pk = pod::ElGamalPubkey::zeroed().try_into().unwrap();
+
+        let transfer_data = TransferData::new(
+            transfer_amount,
+            (spendable_balance, &spendable_ciphertext),
+            &source_keypair,
+            (&dest_pk, &auditor_pk),
+        )
+        .unwrap();
+
+        assert!(transfer_data.verify().is_err());
     }
 
     #[test]

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

@@ -848,5 +848,67 @@ mod test {
         );
 
         assert!(fee_data.is_err());
+
+        // Case 5: invalid destination, auditor, or withdraw authority pubkeys
+        let spendable_balance: u64 = 120;
+        let spendable_ciphertext = source_keypair.public.encrypt(spendable_balance);
+
+        let transfer_amount: u64 = 0;
+
+        let fee_parameters = FeeParameters {
+            fee_rate_basis_points: 400,
+            maximum_fee: 3,
+        };
+
+        // destination pubkey invalid
+        let destination_pubkey: ElGamalPubkey = pod::ElGamalPubkey::zeroed().try_into().unwrap();
+        let auditor_pubkey = ElGamalKeypair::new_rand().public;
+        let withdraw_withheld_authority_pubkey = ElGamalKeypair::new_rand().public;
+
+        let fee_data = TransferWithFeeData::new(
+            transfer_amount,
+            (spendable_balance, &spendable_ciphertext),
+            &source_keypair,
+            (&destination_pubkey, &auditor_pubkey),
+            fee_parameters,
+            &withdraw_withheld_authority_pubkey,
+        )
+        .unwrap();
+
+        assert!(fee_data.verify().is_err());
+
+        // auditor pubkey invalid
+        let destination_pubkey: ElGamalPubkey = ElGamalKeypair::new_rand().public;
+        let auditor_pubkey = pod::ElGamalPubkey::zeroed().try_into().unwrap();
+        let withdraw_withheld_authority_pubkey = ElGamalKeypair::new_rand().public;
+
+        let fee_data = TransferWithFeeData::new(
+            transfer_amount,
+            (spendable_balance, &spendable_ciphertext),
+            &source_keypair,
+            (&destination_pubkey, &auditor_pubkey),
+            fee_parameters,
+            &withdraw_withheld_authority_pubkey,
+        )
+        .unwrap();
+
+        assert!(fee_data.verify().is_err());
+
+        // withdraw authority invalid
+        let destination_pubkey: ElGamalPubkey = ElGamalKeypair::new_rand().public;
+        let auditor_pubkey = ElGamalKeypair::new_rand().public;
+        let withdraw_withheld_authority_pubkey = pod::ElGamalPubkey::zeroed().try_into().unwrap();
+
+        let fee_data = TransferWithFeeData::new(
+            transfer_amount,
+            (spendable_balance, &spendable_ciphertext),
+            &source_keypair,
+            (&destination_pubkey, &auditor_pubkey),
+            fee_parameters,
+            &withdraw_withheld_authority_pubkey,
+        )
+        .unwrap();
+
+        assert!(fee_data.verify().is_err());
     }
 }

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

@@ -47,6 +47,7 @@ impl WithdrawWithheldTokensData {
         withdraw_withheld_authority_ciphertext: &ElGamalCiphertext,
         amount: u64,
     ) -> Result<Self, ProofError> {
+        // encrypt withdraw amount under destination public key
         let destination_opening = PedersenOpening::new_rand();
         let destination_ciphertext = destination_pubkey.encrypt_with(amount, &destination_opening);
 
@@ -193,10 +194,24 @@ mod test {
     use super::*;
 
     #[test]
-    fn test_close_account_correctness() {
+    fn test_withdraw_withheld() {
         let withdraw_withheld_authority_keypair = ElGamalKeypair::new_rand();
         let dest_keypair = ElGamalKeypair::new_rand();
 
+        let amount: u64 = 0;
+        let withdraw_withheld_authority_ciphertext =
+            withdraw_withheld_authority_keypair.public.encrypt(amount);
+
+        let withdraw_withheld_tokens_data = WithdrawWithheldTokensData::new(
+            &withdraw_withheld_authority_keypair,
+            &dest_keypair.public,
+            &withdraw_withheld_authority_ciphertext,
+            amount,
+        )
+        .unwrap();
+
+        assert!(withdraw_withheld_tokens_data.verify().is_ok());
+
         let amount: u64 = 55;
         let withdraw_withheld_authority_ciphertext =
             withdraw_withheld_authority_keypair.public.encrypt(amount);
@@ -210,5 +225,19 @@ mod test {
         .unwrap();
 
         assert!(withdraw_withheld_tokens_data.verify().is_ok());
+
+        let amount = u64::max_value();
+        let withdraw_withheld_authority_ciphertext =
+            withdraw_withheld_authority_keypair.public.encrypt(amount);
+
+        let withdraw_withheld_tokens_data = WithdrawWithheldTokensData::new(
+            &withdraw_withheld_authority_keypair,
+            &dest_keypair.public,
+            &withdraw_withheld_authority_ciphertext,
+            amount,
+        )
+        .unwrap();
+
+        assert!(withdraw_withheld_tokens_data.verify().is_ok());
     }
 }