gorb_ctpg/
withdraw.rs

1use {
2    crate::errors::TokenProofGenerationError,
3    solana_zk_sdk::{
4        encryption::{
5            elgamal::{ElGamal, ElGamalCiphertext, ElGamalKeypair},
6            pedersen::Pedersen,
7        },
8        zk_elgamal_proof_program::proof_data::{
9            BatchedRangeProofU64Data, CiphertextCommitmentEqualityProofData,
10        },
11    },
12};
13
14const REMAINING_BALANCE_BIT_LENGTH: usize = 64;
15
16/// Proof data required for a withdraw instruction
17pub struct WithdrawProofData {
18    pub equality_proof_data: CiphertextCommitmentEqualityProofData,
19    pub range_proof_data: BatchedRangeProofU64Data,
20}
21
22pub fn withdraw_proof_data(
23    current_available_balance: &ElGamalCiphertext,
24    current_balance: u64,
25    withdraw_amount: u64,
26    elgamal_keypair: &ElGamalKeypair,
27) -> Result<WithdrawProofData, TokenProofGenerationError> {
28    // Calculate the remaining balance after withdraw
29    let remaining_balance = current_balance
30        .checked_sub(withdraw_amount)
31        .ok_or(TokenProofGenerationError::NotEnoughFunds)?;
32
33    // Generate a Pedersen commitment for the remaining balance
34    let (remaining_balance_commitment, remaining_balance_opening) =
35        Pedersen::new(remaining_balance);
36
37    // Compute the remaining balance ciphertext
38    #[allow(clippy::arithmetic_side_effects)]
39    let remaining_balance_ciphertext = current_available_balance - ElGamal::encode(withdraw_amount);
40
41    // Generate proof data
42    let equality_proof_data = CiphertextCommitmentEqualityProofData::new(
43        elgamal_keypair,
44        &remaining_balance_ciphertext,
45        &remaining_balance_commitment,
46        &remaining_balance_opening,
47        remaining_balance,
48    )
49    .map_err(TokenProofGenerationError::from)?;
50
51    let range_proof_data = BatchedRangeProofU64Data::new(
52        vec![&remaining_balance_commitment],
53        vec![remaining_balance],
54        vec![REMAINING_BALANCE_BIT_LENGTH],
55        vec![&remaining_balance_opening],
56    )
57    .map_err(TokenProofGenerationError::from)?;
58
59    Ok(WithdrawProofData {
60        equality_proof_data,
61        range_proof_data,
62    })
63}