diff --git a/token/confidential-transfer/proof-extraction/src/lib.rs b/token/confidential-transfer/proof-extraction/src/lib.rs index 2933275b062..15cee9bdc5d 100644 --- a/token/confidential-transfer/proof-extraction/src/lib.rs +++ b/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; diff --git a/token/confidential-transfer/proof-extraction/src/withdraw.rs b/token/confidential-transfer/proof-extraction/src/withdraw.rs new file mode 100644 index 00000000000..7fc5fb17dc5 --- /dev/null +++ b/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 { + 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) + } +} diff --git a/token/confidential-transfer/proof-generation/src/lib.rs b/token/confidential-transfer/proof-generation/src/lib.rs index ce8b0ea1cb7..ba769a06fb4 100644 --- a/token/confidential-transfer/proof-generation/src/lib.rs +++ b/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; diff --git a/token/confidential-transfer/proof-generation/src/withdraw.rs b/token/confidential-transfer/proof-generation/src/withdraw.rs new file mode 100644 index 00000000000..ece1fedc58e --- /dev/null +++ b/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 { + // 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, + }) +} diff --git a/token/confidential-transfer/proof-tests/tests/proof_test.rs b/token/confidential-transfer/proof-tests/tests/proof_test.rs index f1d52ff1ad6..da87a3ab8cb 100644 --- a/token/confidential-transfer/proof-tests/tests/proof_test.rs +++ b/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(); +}