update withdraw instruction

Author: samkim-crypto

token/program-2022/src/extension/confidential_transfer/account_info.rs

@@ -15,8 +15,10 @@ use {
         zk_elgamal_proof_program::proof_data::ZeroCiphertextProofData,
     },
     spl_pod::primitives::PodU64,
-    spl_token_confidential_transfer_proof_generation::transfer::{
-        transfer_split_proof_data, TransferProofData,
+    spl_token_confidential_transfer_proof_generation::{
+        transfer::{transfer_split_proof_data, TransferProofData},
+        transfer_with_fee::{transfer_with_fee_split_proof_data, TransferWithFeeProofData},
+        withdraw::{withdraw_proof_data, WithdrawProofData},
     },
 };
 
@@ -176,20 +178,20 @@ impl WithdrawAccountInfo {
         withdraw_amount: u64,
         elgamal_keypair: &ElGamalKeypair,
         aes_key: &AeKey,
-    ) -> Result<WithdrawData, TokenError> {
+    ) -> Result<WithdrawProofData, TokenError> {
         let current_available_balance = self
             .available_balance
             .try_into()
             .map_err(|_| TokenError::MalformedCiphertext)?;
         let current_decrypted_available_balance = self.decrypted_available_balance(aes_key)?;
 
-        WithdrawData::new(
+        withdraw_proof_data(
+            &current_available_balance,
+            current_decrypted_available_balance,
             withdraw_amount,
             elgamal_keypair,
-            current_decrypted_available_balance,
-            &current_available_balance,
         )
-        .map_err(|_| TokenError::ProofGeneration)
+        .map_err(|e| -> TokenError { e.into() })
     }
 
     /// Update the decryptable available balance.
@@ -268,6 +270,43 @@ impl TransferAccountInfo {
         .map_err(|e| -> TokenError { e.into() })
     }
 
+    /// Create a transfer proof data that is split into equality, ciphertext validity (transfer
+    /// amount), percentage-with-cap, ciphertext validity (fee), and range proofs.
+    pub fn generate_split_transfer_with_fee_proof_data(
+        &self,
+        transfer_amount: u64,
+        source_elgamal_keypair: &ElGamalKeypair,
+        aes_key: &AeKey,
+        destination_elgamal_pubkey: &ElGamalPubkey,
+        auditor_elgamal_pubkey: Option<&ElGamalPubkey>,
+        withdraw_withheld_authority_elgamal_pubkey: &ElGamalPubkey,
+        fee_rate_basis_points: u16,
+        maximum_fee: u64,
+    ) -> Result<TransferWithFeeProofData, TokenError> {
+        let current_available_balance = self
+            .available_balance
+            .try_into()
+            .map_err(|_| TokenError::MalformedCiphertext)?;
+        let current_decryptable_available_balance = self
+            .decryptable_available_balance
+            .try_into()
+            .map_err(|_| TokenError::MalformedCiphertext)?;
+
+        transfer_with_fee_split_proof_data(
+            &current_available_balance,
+            &current_decryptable_available_balance,
+            transfer_amount,
+            source_elgamal_keypair,
+            aes_key,
+            destination_elgamal_pubkey,
+            auditor_elgamal_pubkey,
+            withdraw_withheld_authority_elgamal_pubkey,
+            fee_rate_basis_points,
+            maximum_fee,
+        )
+        .map_err(|e| -> TokenError { e.into() })
+    }
+
     /// Update the decryptable available balance.
     pub fn new_decryptable_available_balance(
         &self,

token/program-2022/src/extension/confidential_transfer/instruction.rs

@@ -205,36 +205,37 @@ pub enum ConfidentialTransferInstruction {
     /// Withdraw SPL Tokens from the available balance of a confidential token
     /// account.
     ///
+    /// In order for this instruction to be successfully processed, it must be accompanied by the
+    /// following list of `zk_elgamal_proof` program instructions:
+    /// - `VerifyCiphertextCommitmentEquality`
+    /// - `VerifyBatchedRangeProofU64`
+    /// These instructions can be accompanied in the same transaction or can be pre-verified into a
+    /// context state account, in which case, only their context state account address need to be
+    /// provided.
+    ///
     /// Fails if the source or destination accounts are frozen.
     /// Fails if the associated mint is extended as `NonTransferable`.
     ///
-    /// In order for this instruction to be successfully processed, it must be
-    /// accompanied by the `VerifyWithdraw` instruction of the
-    /// `zk_token_proof` program in the same transaction or the address of a
-    /// context state account for the proof must be provided.
-    ///
     /// Accounts expected by this instruction:
     ///
     ///   * Single owner/delegate
     ///   0. `[writable]` The SPL Token account.
     ///   1. `[]` The token mint.
-    ///   2. `[]` Instructions sysvar if `VerifyWithdraw` is included in the
-    ///      same transaction or context state account if `VerifyWithdraw` is
-    ///      pre-verified into a context state account.
-    ///   3. `[]` (Optional) Record account if the accompanying proof is to be
-    ///      read from a record account.
-    ///   4. `[signer]` The single source account owner.
+    ///   2. `[]` (Optional) Instructions sysvar if at least one of the `zk_elgamal_proof`
+    ///      instructions are included in the same transaction.
+    ///   3. `[]` (Optional) Equality proof record account or context state account.
+    ///   4. `[]` (Optional) Range proof record account or context state account.
+    ///   5. `[signer]` The single source account owner.
     ///
     ///   * Multisignature owner/delegate
     ///   0. `[writable]` The SPL Token account.
     ///   1. `[]` The token mint.
-    ///   2. `[]` Instructions sysvar if `VerifyWithdraw` is included in the
-    ///      same transaction or context state account if `VerifyWithdraw` is
-    ///      pre-verified into a context state account.
-    ///   3. `[]` (Optional) Record account if the accompanying proof is to be
-    ///      read from a record account.
-    ///   4. `[]` The multisig  source account owner.
-    ///   5.. `[signer]` Required M signer accounts for the SPL Token Multisig
+    ///   2. `[]` (Optional) Instructions sysvar if at least one of the `zk_elgamal_proof`
+    ///      instructions are included in the same transaction.
+    ///   3. `[]` (Optional) Equality proof record account or context state account.
+    ///   4. `[]` (Optional) Range proof record account or context state account.
+    ///   5. `[]` The multisig  source account owner.
+    ///   6.. `[signer]` Required M signer accounts for the SPL Token Multisig
     /// account.
     ///
     /// Data expected by this instruction:
@@ -548,10 +549,14 @@ pub struct WithdrawInstructionData {
     /// The new decryptable balance if the withdrawal succeeds
     #[cfg_attr(feature = "serde-traits", serde(with = "aeciphertext_fromstr"))]
     pub new_decryptable_available_balance: DecryptableBalance,
-    /// Relative location of the `ProofInstruction::VerifyWithdraw` instruction
+    /// Relative location of the `ProofInstruction::VerifyCiphertextCommitmentEquality` instruction
     /// to the `Withdraw` instruction in the transaction. If the offset is
     /// `0`, then use a context state account for the proof.
-    pub proof_instruction_offset: i8,
+    pub equality_proof_instruction_offset: i8,
+    /// Relative location of the `ProofInstruction::BatchedRangeProofU64` instruction
+    /// to the `Withdraw` instruction in the transaction. If the offset is
+    /// `0`, then use a context state account for the proof.
+    pub range_proof_instruction_offset: i8,
 }
 
 /// Data expected by `ConfidentialTransferInstruction::Transfer`
@@ -948,17 +953,38 @@ pub fn inner_withdraw(
     new_decryptable_available_balance: DecryptableBalance,
     authority: &Pubkey,
     multisig_signers: &[&Pubkey],
-    proof_data_location: ProofLocation<WithdrawData>,
+    equality_proof_data_location: ProofLocation<CiphertextCommitmentEqualityProofData>,
+    range_proof_data_location: ProofLocation<BatchedRangeProofU64Data>,
 ) -> Result<Instruction, ProgramError> {
     check_program_account(token_program_id)?;
     let mut accounts = vec![
         AccountMeta::new(*token_account, false),
         AccountMeta::new_readonly(*mint, false),
     ];
 
-    let proof_instruction_offset = match proof_data_location {
+    // if at least one of the proof locations is an instruction offset, sysvar
+    // account is needed
+    if equality_proof_data_location.is_instruction_offset()
+        || range_proof_data_location.is_instruction_offset()
+    {
+        accounts.push(AccountMeta::new_readonly(sysvar::instructions::id(), false));
+    }
+
+    let equality_proof_instruction_offset = match equality_proof_data_location {
+        ProofLocation::InstructionOffset(proof_instruction_offset, proof_data) => {
+            if let ProofData::RecordAccount(record_address, _) = proof_data {
+                accounts.push(AccountMeta::new_readonly(*record_address, false));
+            }
+            proof_instruction_offset.into()
+        }
+        ProofLocation::ContextStateAccount(context_state_account) => {
+            accounts.push(AccountMeta::new_readonly(*context_state_account, false));
+            0
+        }
+    };
+
+    let range_proof_instruction_offset = match range_proof_data_location {
         ProofLocation::InstructionOffset(proof_instruction_offset, proof_data) => {
-            accounts.push(AccountMeta::new_readonly(sysvar::instructions::id(), false));
             if let ProofData::RecordAccount(record_address, _) = proof_data {
                 accounts.push(AccountMeta::new_readonly(*record_address, false));
             }
@@ -988,7 +1014,8 @@ pub fn inner_withdraw(
             amount: amount.into(),
             decimals,
             new_decryptable_available_balance,
-            proof_instruction_offset,
+            equality_proof_instruction_offset,
+            range_proof_instruction_offset,
         },
     ))
 }
@@ -1004,7 +1031,8 @@ pub fn withdraw(
     new_decryptable_available_balance: PodAeCiphertext,
     authority: &Pubkey,
     multisig_signers: &[&Pubkey],
-    proof_data_location: ProofLocation<WithdrawData>,
+    equality_proof_data_location: ProofLocation<CiphertextCommitmentEqualityProofData>,
+    range_proof_data_location: ProofLocation<BatchedRangeProofU64Data>,
 ) -> Result<Vec<Instruction>, ProgramError> {
     let mut instructions = vec![inner_withdraw(
         token_program_id,
@@ -1015,24 +1043,41 @@ pub fn withdraw(
         new_decryptable_available_balance.into(),
         authority,
         multisig_signers,
-        proof_data_location,
+        equality_proof_data_location,
+        range_proof_data_location,
     )?];
 
     if let ProofLocation::InstructionOffset(proof_instruction_offset, proof_data) =
-        proof_data_location
+        equality_proof_data_location
     {
-        // This constructor appends the proof instruction right after the `Withdraw`
-        // instruction. This means that the proof instruction offset must be
-        // always be 1. To use an arbitrary proof instruction offset, use the
-        // `inner_withdraw` constructor.
         let proof_instruction_offset: i8 = proof_instruction_offset.into();
         if proof_instruction_offset != 1 {
             return Err(TokenError::InvalidProofInstructionOffset.into());
         }
         match proof_data {
-            ProofData::InstructionData(data) => instructions.push(verify_withdraw(None, data)),
+            ProofData::InstructionData(data) => instructions.push(
+                ProofInstruction::VerifyCiphertextCommitmentEquality
+                    .encode_verify_proof(None, data),
+            ),
             ProofData::RecordAccount(address, offset) => instructions.push(
-                ProofInstruction::VerifyWithdraw
+                ProofInstruction::VerifyCiphertextCommitmentEquality
+                    .encode_verify_proof_from_account(None, address, offset),
+            ),
+        };
+    };
+
+    if let ProofLocation::InstructionOffset(proof_instruction_offset, proof_data) =
+        range_proof_data_location
+    {
+        let proof_instruction_offset: i8 = proof_instruction_offset.into();
+        if proof_instruction_offset != 2 {
+            return Err(TokenError::InvalidProofInstructionOffset.into());
+        }
+        match proof_data {
+            ProofData::InstructionData(data) => instructions
+                .push(ProofInstruction::VerifyBatchedRangeProofU64.encode_verify_proof(None, data)),
+            ProofData::RecordAccount(address, offset) => instructions.push(
+                ProofInstruction::VerifyBatchedRangeProofU64
                     .encode_verify_proof_from_account(None, address, offset),
             ),
         };
@@ -1120,6 +1165,10 @@ pub fn inner_transfer(
         multisig_signers.is_empty(),
     ));
 
+    for multisig_signer in multisig_signers.iter() {
+        accounts.push(AccountMeta::new_readonly(**multisig_signer, true));
+    }
+
     Ok(encode_instruction(
         token_program_id,
         accounts,
@@ -1476,6 +1525,10 @@ pub fn inner_transfer_with_fee(
         multisig_signers.is_empty(),
     ));
 
+    for multisig_signer in multisig_signers.iter() {
+        accounts.push(AccountMeta::new_readonly(**multisig_signer, true));
+    }
+
     Ok(encode_instruction(
         token_program_id,
         accounts,

token/program-2022/src/extension/confidential_transfer/processor.rs

@@ -359,18 +359,19 @@ fn process_withdraw(
     amount: u64,
     expected_decimals: u8,
     new_decryptable_available_balance: DecryptableBalance,
-    proof_instruction_offset: i64,
+    equality_proof_instruction_offset: i64,
+    range_proof_instruction_offset: i64,
 ) -> ProgramResult {
     let account_info_iter = &mut accounts.iter();
     let token_account_info = next_account_info(account_info_iter)?;
     let mint_info = next_account_info(account_info_iter)?;
 
     // zero-knowledge proof certifies that the account has enough available balance
     // to withdraw the amount.
-    let proof_context = verify_and_extract_context::<WithdrawData, WithdrawProofContext>(
+    let proof_context = verify_withdraw_proof(
         account_info_iter,
-        proof_instruction_offset,
-        None,
+        equality_proof_instruction_offset,
+        range_proof_instruction_offset,
     )?;
 
     let authority_info = next_account_info(account_info_iter)?;
@@ -421,7 +422,7 @@ fn process_withdraw(
     // Check that the encryption public key associated with the confidential
     // extension is consistent with the public key that was actually used to
     // generate the zkp.
-    if confidential_transfer_account.elgamal_pubkey != proof_context.pubkey {
+    if confidential_transfer_account.elgamal_pubkey != proof_context.source_pubkey {
         return Err(TokenError::ConfidentialTransferElGamalPubkeyMismatch.into());
     }
 
@@ -436,7 +437,8 @@ fn process_withdraw(
     }
     // Check that the final available balance ciphertext is consistent with the
     // actual ciphertext for which the zero-knowledge proof was generated for.
-    if confidential_transfer_account.available_balance != proof_context.final_ciphertext {
+    if confidential_transfer_account.available_balance != proof_context.remaining_balance_ciphertext
+    {
         return Err(TokenError::ConfidentialTransferBalanceMismatch.into());
     }
 
@@ -1139,7 +1141,8 @@ pub(crate) fn process_instruction(
                     data.amount.into(),
                     data.decimals,
                     data.new_decryptable_available_balance,
-                    data.proof_instruction_offset as i64,
+                    data.equality_proof_instruction_offset as i64,
+                    data.range_proof_instruction_offset as i64,
                 )
             }
             #[cfg(not(feature = "zk-ops"))]

token/program-2022/src/extension/confidential_transfer/verify_proof.rs

@@ -10,10 +10,52 @@ use {
     },
     spl_token_confidential_transfer_proof_extraction::{
         transfer::TransferProofContext, transfer_with_fee::TransferWithFeeProofContext,
+        withdraw::WithdrawProofContext,
     },
     std::slice::Iter,
 };
 
+/// Verify zero-knowledge proofs needed for a [Withdraw] instruction and return the corresponding
+/// proof context.
+#[cfg(feature = "zk-ops")]
+pub fn verify_withdraw_proof(
+    account_info_iter: &mut Iter<AccountInfo>,
+    equality_proof_instruction_offset: i64,
+    range_proof_instruction_offset: i64,
+) -> Result<WithdrawProofContext, ProgramError> {
+    let sysvar_account_info =
+        if equality_proof_instruction_offset != 0 || range_proof_instruction_offset != 0 {
+            Some(next_account_info(account_info_iter)?)
+        } else {
+            None
+        };
+
+    let equality_proof_context = verify_and_extract_context::<
+        CiphertextCommitmentEqualityProofData,
+        CiphertextCommitmentEqualityProofContext,
+    >(
+        account_info_iter,
+        equality_proof_instruction_offset,
+        sysvar_account_info,
+    )?;
+
+    let range_proof_context =
+        verify_and_extract_context::<BatchedRangeProofU64Data, BatchedRangeProofContext>(
+            account_info_iter,
+            range_proof_instruction_offset,
+            sysvar_account_info,
+        )?;
+
+    // The `WithdrawProofContext` constructor verifies the consistency of the
+    // individual proof context and generates a `WithdrawProofContext` struct
+    // that is used to process the rest of the token-2022 logic.
+    let transfer_proof_context =
+        WithdrawProofContext::verify_and_extract(&equality_proof_context, &range_proof_context)
+            .map_err(|e| -> TokenError { e.into() })?;
+
+    Ok(transfer_proof_context)
+}
+
 /// Verify zero-knowledge proof needed for a [Transfer] instruction without fee
 /// and return the corresponding proof context.
 #[cfg(feature = "zk-ops")]