re-organize proof data using structs

Author: samkim-crypto

token/confidential-transfer/proof-generation/src/withdraw.rs

@@ -13,18 +13,18 @@ use {
 
 const REMAINING_BALANCE_BIT_LENGTH: usize = 64;
 
+/// Proof data required for a withdraw instruction
+pub struct WithdrawProofData {
+    pub equality_proof_data: CiphertextCommitmentEqualityProofData,
+    pub range_proof_data: BatchedRangeProofU64Data,
+}
+
 pub fn withdraw_proof_data(
     current_available_balance: &ElGamalCiphertext,
     current_balance: u64,
     withdraw_amount: u64,
     elgamal_keypair: &ElGamalKeypair,
-) -> Result<
-    (
-        CiphertextCommitmentEqualityProofData,
-        BatchedRangeProofU64Data,
-    ),
-    TokenProofGenerationError,
-> {
+) -> Result<WithdrawProofData, TokenProofGenerationError> {
     // Calculate the remaining balance after withdraw
     let remaining_balance = current_balance
         .checked_sub(withdraw_amount)
@@ -56,5 +56,8 @@ pub fn withdraw_proof_data(
     )
     .map_err(TokenProofGenerationError::from)?;
 
-    Ok((equality_proof_data, range_proof_data))
+    Ok(WithdrawProofData {
+        equality_proof_data,
+        range_proof_data,
+    })
 }

token/confidential-transfer/proof-tests/tests/proof_test.rs

@@ -5,9 +5,12 @@ use {
     },
     spl_token_confidential_transfer_proof_extraction::{
         transfer::TransferProofContext, transfer_with_fee::TransferWithFeeProofContext,
+        withdraw::WithdrawProofContext,
     },
     spl_token_confidential_transfer_proof_generation::{
-        transfer::transfer_split_proof_data, transfer_with_fee::transfer_with_fee_split_proof_data,
+        transfer::transfer_split_proof_data,
+        transfer_with_fee::transfer_with_fee_split_proof_data,
+        withdraw::{withdraw_proof_data, WithdrawProofData},
     },
 };
 
@@ -140,3 +143,38 @@ fn test_transfer_with_fee_proof_validity(
     )
     .unwrap();
 }
+
+#[test]
+fn test_withdraw_proof_correctness() {
+    test_withdraw_validity(0, 0);
+    test_withdraw_validity(77, 55);
+    test_withdraw_validity(65535, 65535);
+    test_withdraw_validity(65536, 65536);
+    test_withdraw_validity(281474976710655, 281474976710655);
+}
+
+fn test_withdraw_validity(spendable_balance: u64, withdraw_amount: u64) {
+    let keypair = ElGamalKeypair::new_rand();
+
+    let spendable_ciphertext = keypair.pubkey().encrypt(spendable_balance);
+
+    let WithdrawProofData {
+        equality_proof_data,
+        range_proof_data,
+    } = withdraw_proof_data(
+        &spendable_ciphertext,
+        spendable_balance,
+        withdraw_amount,
+        &keypair,
+    )
+    .unwrap();
+
+    equality_proof_data.verify_proof().unwrap();
+    range_proof_data.verify_proof().unwrap();
+
+    WithdrawProofContext::verify_and_extract(
+        equality_proof_data.context_data(),
+        range_proof_data.context_data(),
+    )
+    .unwrap();
+}