[confidential-transfer] Add withdraw proof generation and extraction (#6954)

* add withdraw proof generation

* add withdraw proof extraction

* re-organize proof data using structs

Author: samkim-crypto

token/confidential-transfer/proof-extraction/src/lib.rs

@@ -2,3 +2,4 @@ pub mod encryption;
 pub mod errors;
 pub mod transfer;
 pub mod transfer_with_fee;
+pub mod withdraw;

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

@@ -0,0 +1,53 @@
+use {
+    crate::errors::TokenProofExtractionError,
+    solana_zk_sdk::{
+        encryption::pod::elgamal::{PodElGamalCiphertext, PodElGamalPubkey},
+        zk_elgamal_proof_program::proof_data::{
+            BatchedRangeProofContext, CiphertextCommitmentEqualityProofContext,
+        },
+    },
+};
+
+const REMAINING_BALANCE_BIT_LENGTH: u8 = 64;
+
+pub struct WithdrawProofContext {
+    pub source_pubkey: PodElGamalPubkey,
+    pub remaining_balance_ciphertext: PodElGamalCiphertext,
+}
+
+impl WithdrawProofContext {
+    pub fn verify_and_extract(
+        equality_proof_context: &CiphertextCommitmentEqualityProofContext,
+        range_proof_context: &BatchedRangeProofContext,
+    ) -> Result<Self, TokenProofExtractionError> {
+        let CiphertextCommitmentEqualityProofContext {
+            pubkey: source_pubkey,
+            ciphertext: remaining_balance_ciphertext,
+            commitment: remaining_balance_commitment,
+        } = equality_proof_context;
+
+        let BatchedRangeProofContext {
+            commitments: range_proof_commitments,
+            bit_lengths: range_proof_bit_lengths,
+        } = range_proof_context;
+
+        if range_proof_commitments.is_empty()
+            || range_proof_commitments[0] != *remaining_balance_commitment
+        {
+            return Err(TokenProofExtractionError::PedersenCommitmentMismatch);
+        }
+
+        if range_proof_bit_lengths.is_empty()
+            || range_proof_bit_lengths[0] != REMAINING_BALANCE_BIT_LENGTH
+        {
+            return Err(TokenProofExtractionError::RangeProofLengthMismatch);
+        }
+
+        let context_info = WithdrawProofContext {
+            source_pubkey: *source_pubkey,
+            remaining_balance_ciphertext: *remaining_balance_ciphertext,
+        };
+
+        Ok(context_info)
+    }
+}

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

@@ -10,6 +10,7 @@ pub mod encryption;
 pub mod errors;
 pub mod transfer;
 pub mod transfer_with_fee;
+pub mod withdraw;
 
 /// The low bit length of the encrypted transfer amount
 pub const TRANSFER_AMOUNT_LO_BITS: usize = 16;

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

@@ -0,0 +1,63 @@
+use {
+    crate::errors::TokenProofGenerationError,
+    solana_zk_sdk::{
+        encryption::{
+            elgamal::{ElGamal, ElGamalCiphertext, ElGamalKeypair},
+            pedersen::Pedersen,
+        },
+        zk_elgamal_proof_program::proof_data::{
+            BatchedRangeProofU64Data, CiphertextCommitmentEqualityProofData,
+        },
+    },
+};
+
+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<WithdrawProofData, TokenProofGenerationError> {
+    // Calculate the remaining balance after withdraw
+    let remaining_balance = current_balance
+        .checked_sub(withdraw_amount)
+        .ok_or(TokenProofGenerationError::NotEnoughFunds)?;
+
+    // Generate a Pedersen commitment for the remaining balance
+    let (remaining_balance_commitment, remaining_balance_opening) =
+        Pedersen::new(remaining_balance);
+
+    // Compute the remaining balance ciphertext
+    #[allow(clippy::arithmetic_side_effects)]
+    let remaining_balance_ciphertext = current_available_balance - ElGamal::encode(withdraw_amount);
+
+    // Generate proof data
+    let equality_proof_data = CiphertextCommitmentEqualityProofData::new(
+        elgamal_keypair,
+        &remaining_balance_ciphertext,
+        &remaining_balance_commitment,
+        &remaining_balance_opening,
+        remaining_balance,
+    )
+    .map_err(TokenProofGenerationError::from)?;
+
+    let range_proof_data = BatchedRangeProofU64Data::new(
+        vec![&remaining_balance_commitment],
+        vec![remaining_balance],
+        vec![REMAINING_BALANCE_BIT_LENGTH],
+        vec![&remaining_balance_opening],
+    )
+    .map_err(TokenProofGenerationError::from)?;
+
+    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();
+}