add withdraw proof generation

Author: samkim-crypto

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,60 @@
+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;
+
+pub fn withdraw_proof_data(
+    current_available_balance: &ElGamalCiphertext,
+    current_balance: u64,
+    withdraw_amount: u64,
+    elgamal_keypair: &ElGamalKeypair,
+) -> Result<
+    (
+        CiphertextCommitmentEqualityProofData,
+        BatchedRangeProofU64Data,
+    ),
+    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((equality_proof_data, range_proof_data))
+}