From c51ccf05c27cbe9d72cf005c49b37a11c99807b8 Mon Sep 17 00:00:00 2001 From: samkim-crypto Date: Fri, 16 Dec 2022 12:42:55 +0900 Subject: [PATCH] [confidential-extension] Refactor and add comments for readability (#3921) * pass over `InitializeAccount` to `Withdraw` instructions * add `valid_as_source` and `valid_as_destination` functions * pass over `Transfer` instruction * pass over `ApplyPendingBalance` to `HarvestWithheldTokensToMint` instructions * refactor pending balance credit counter increment --- .../tests/confidential_transfer.rs | 10 +- .../extension/confidential_transfer/mod.rs | 46 +- .../confidential_transfer/processor.rs | 475 ++++++++---------- 3 files changed, 257 insertions(+), 274 deletions(-) diff --git a/token/program-2022-test/tests/confidential_transfer.rs b/token/program-2022-test/tests/confidential_transfer.rs index 3c482005..d2e886bf 100644 --- a/token/program-2022-test/tests/confidential_transfer.rs +++ b/token/program-2022-test/tests/confidential_transfer.rs @@ -992,7 +992,10 @@ async fn ct_transfer() { assert_eq!( err, TokenClientError::Client(Box::new(TransportError::TransactionError( - TransactionError::InstructionError(0, InstructionError::InvalidAccountData) + TransactionError::InstructionError( + 0, + InstructionError::Custom(TokenError::ConfidentialTransferAccountHasBalance as u32) + ) ))) ); @@ -1202,7 +1205,10 @@ async fn ct_transfer_with_fee() { assert_eq!( err, TokenClientError::Client(Box::new(TransportError::TransactionError( - TransactionError::InstructionError(0, InstructionError::InvalidAccountData) + TransactionError::InstructionError( + 0, + InstructionError::Custom(TokenError::ConfidentialTransferAccountHasBalance as u32) + ) ))) ); diff --git a/token/program-2022/src/extension/confidential_transfer/mod.rs b/token/program-2022/src/extension/confidential_transfer/mod.rs index add02a2b..20861483 100644 --- a/token/program-2022/src/extension/confidential_transfer/mod.rs +++ b/token/program-2022/src/extension/confidential_transfer/mod.rs @@ -141,7 +141,7 @@ impl Extension for ConfidentialTransferAccount { } impl ConfidentialTransferAccount { - /// Check if a `ConfidentialTransferAccount` has been approved for use + /// Check if a `ConfidentialTransferAccount` has been approved for use. pub fn approved(&self) -> ProgramResult { if bool::from(&self.approved) { Ok(()) @@ -150,7 +150,7 @@ impl ConfidentialTransferAccount { } } - /// Check if a `ConfidentialTransferAccount` is in a closable state + /// Check if a `ConfidentialTransferAccount` is in a closable state. pub fn closable(&self) -> ProgramResult { if self.pending_balance_lo == EncryptedBalance::zeroed() && self.pending_balance_hi == EncryptedBalance::zeroed() @@ -164,7 +164,7 @@ impl ConfidentialTransferAccount { } /// Check if a base account of a `ConfidentialTransferAccount` accepts non-confidential - /// transfers + /// transfers. pub fn non_confidential_transfer_allowed(&self) -> ProgramResult { if bool::from(&self.allow_non_confidential_credits) { Ok(()) @@ -172,4 +172,44 @@ impl ConfidentialTransferAccount { Err(TokenError::NonConfidentialTransfersDisabled.into()) } } + + /// Checks if a `ConfidentialTransferAccount` is configured to send funds. + pub fn valid_as_source(&self) -> ProgramResult { + self.approved() + } + + /// Checks if a confidential extension is configured to receive funds. + /// + /// A destination account can receive funds if the following conditions are satisfied: + /// 1. The account is approved by the confidential transfer mint authority + /// 2. The account is not disabled by the account owner + /// 3. The number of credits into the account has reached the maximum credit counter + pub fn valid_as_destination(&self) -> ProgramResult { + self.approved()?; + + if !bool::from(self.allow_confidential_credits) { + return Err(TokenError::ConfidentialTransferDepositsAndTransfersDisabled.into()); + } + + let new_destination_pending_balance_credit_counter = + u64::from(self.pending_balance_credit_counter) + .checked_add(1) + .ok_or(TokenError::Overflow)?; + if new_destination_pending_balance_credit_counter + > u64::from(self.maximum_pending_balance_credit_counter) + { + return Err(TokenError::MaximumPendingBalanceCreditCounterExceeded.into()); + } + + Ok(()) + } + + /// Increments a confidential extension pending balance credit counter. + pub fn increment_pending_balance_credit_counter(&mut self) -> ProgramResult { + self.pending_balance_credit_counter = (u64::from(self.pending_balance_credit_counter) + .checked_add(1) + .ok_or(TokenError::Overflow)?) + .into(); + Ok(()) + } } diff --git a/token/program-2022/src/extension/confidential_transfer/processor.rs b/token/program-2022/src/extension/confidential_transfer/processor.rs index 583442f1..4ddb5857 100644 --- a/token/program-2022/src/extension/confidential_transfer/processor.rs +++ b/token/program-2022/src/extension/confidential_transfer/processor.rs @@ -33,6 +33,11 @@ use { solana_zk_token_sdk::zk_token_elgamal::ops, }; +/// Decodes the zero-knowledge proof instruction associated with the token instruction. +/// +/// `ConfigureAccount`, `EmptyAccount`, `Withdraw`, `Transfer`, `WithdrawWithheldTokensFromMint`, +/// and `WithdrawWithheldTokensFromAccounts` instructions require corresponding zero-knowledge +/// proof instructions. fn decode_proof_instruction( expected: ProofInstruction, instruction: &Instruction, @@ -47,40 +52,6 @@ fn decode_proof_instruction( ProofInstruction::decode_data(&instruction.data).ok_or(ProgramError::InvalidInstructionData) } -/// Checks if a confidential extension is configured to send funds -#[cfg(feature = "zk-ops")] -fn check_source_confidential_account( - source_confidential_transfer_account: &ConfidentialTransferAccount, -) -> ProgramResult { - source_confidential_transfer_account.approved() -} - -/// Checks if a confidential extension is configured to receive funds -#[cfg(feature = "zk-ops")] -fn check_destination_confidential_account( - destination_confidential_transfer_account: &ConfidentialTransferAccount, -) -> ProgramResult { - destination_confidential_transfer_account.approved()?; - - if !bool::from(&destination_confidential_transfer_account.allow_confidential_credits) { - return Err(TokenError::ConfidentialTransferDepositsAndTransfersDisabled.into()); - } - - let new_destination_pending_balance_credit_counter = - u64::from(destination_confidential_transfer_account.pending_balance_credit_counter) - .checked_add(1) - .ok_or(ProgramError::InvalidInstructionData)?; - if new_destination_pending_balance_credit_counter - > u64::from( - destination_confidential_transfer_account.maximum_pending_balance_credit_counter, - ) - { - return Err(TokenError::MaximumPendingBalanceCreditCounterExceeded.into()); - } - - Ok(()) -} - /// Processes an [InitializeMint] instruction. fn process_initialize_mint( accounts: &[AccountInfo], @@ -161,13 +132,15 @@ fn process_configure_account( let mint = StateWithExtensions::::unpack(mint_data)?; let confidential_transfer_mint = mint.get_extension::()?; - let previous_instruction = + // zero-knowledge proof certifies that the supplied encryption (ElGamal) public key is valid + let zkp_instruction = get_instruction_relative(proof_instruction_offset, instructions_sysvar_info)?; let proof_data = decode_proof_instruction::( ProofInstruction::VerifyPubkeyValidity, - &previous_instruction, + &zkp_instruction, )?; - + // Check that the encryption public key associated with the confidential extension account is + // consistent with what was actually used to generate the zkp. if proof_data.pubkey != *encryption_pubkey { return Err(TokenError::ConfidentialTransferElGamalPubkeyMismatch.into()); } @@ -181,35 +154,7 @@ fn process_configure_account( confidential_transfer_account.maximum_pending_balance_credit_counter = *maximum_pending_balance_credit_counter; - /* - An ElGamal ciphertext is of the form - ElGamalCiphertext { - msg_comm: r * H + x * G - decrypt_handle: r * P - } - - where - - G, H: constants for the system (RistrettoPoint) - - P: ElGamal public key component (RistrettoPoint) - - r: encryption randomness (Scalar) - - x: message (Scalar) - - Upon receiving a `ConfigureAccount` instruction, the ZK Token program should encrypt x=0 - (i.e. Scalar::zero()) and store it as `pending_balance_lo`, `pending_balance_hi`, and - `available_balance`. - - For regular encryption, it is important that r is generated from a proper randomness source. But - for the `ConfigureAccount` instruction, it is already known that x is always 0. So r can just be - set Scalar::zero(). - - This means that the ElGamalCiphertext should simply be - ElGamalCiphertext { - msg_comm: 0 * H + 0 * G = 0 - decrypt_handle: 0 * P = 0 - } - - This should just be encoded as [0; 64] - */ + // The all-zero ciphertext [0; 64] is a valid encryption of zero confidential_transfer_account.pending_balance_lo = EncryptedBalance::zeroed(); confidential_transfer_account.pending_balance_hi = EncryptedBalance::zeroed(); confidential_transfer_account.available_balance = EncryptedBalance::zeroed(); @@ -278,33 +223,40 @@ fn process_empty_account( let mut confidential_transfer_account = token_account.get_extension_mut::()?; - let previous_instruction = + // An account can be closed only if the remaining balance is zero. This means that for the + // confidential extension account, the ciphertexts associated with the following components + // must be an encryption of zero: + // 1. The pending balance + // 2. The available balance + // 3. The withheld balance + // + // For the pending and withheld balance ciphertexts, it suffices to check that they are + // all-zero ciphertexts (i.e. [0; 64]). If any of these ciphertexts are valid encryption of + // zero but not an all-zero ciphertext, then an `ApplyPendingBalance` or + // `HarvestWithheldTokensToMint` instructions can be used to flush-out these balances first. + // + // For the available balance, it is not possible to deduce whether the ciphertext encrypts zero + // or not by simply inspecting the ciphertext bytes (otherwise, this would violate + // confidentiality). The available balance is verified using a zero-knowledge proof. + let zkp_instruction = get_instruction_relative(proof_instruction_offset, instructions_sysvar_info)?; let proof_data = decode_proof_instruction::( ProofInstruction::VerifyCloseAccount, - &previous_instruction, + &zkp_instruction, )?; - - if confidential_transfer_account.pending_balance_lo != EncryptedBalance::zeroed() { - msg!("Pending balance is not zero"); - return Err(ProgramError::InvalidAccountData); + // Check that the encryption public key and ciphertext associated with the confidential + // extension account are consistent with those that were actually used to generate the zkp. + if confidential_transfer_account.encryption_pubkey != proof_data.pubkey { + msg!("Encryption public-key mismatch"); + return Err(TokenError::ConfidentialTransferElGamalPubkeyMismatch.into()); } - - if confidential_transfer_account.pending_balance_hi != EncryptedBalance::zeroed() { - msg!("Pending balance is not zero"); - return Err(ProgramError::InvalidAccountData); - } - if confidential_transfer_account.available_balance != proof_data.ciphertext { msg!("Available balance mismatch"); return Err(ProgramError::InvalidInstructionData); } - - if confidential_transfer_account.encryption_pubkey != proof_data.pubkey { - return Err(TokenError::ConfidentialTransferElGamalPubkeyMismatch.into()); - } - confidential_transfer_account.available_balance = EncryptedBalance::zeroed(); + + // check that all balances are all-zero ciphertexts confidential_transfer_account.closable()?; Ok(()) @@ -359,46 +311,51 @@ fn process_deposit( // Wrapped SOL deposits are not supported because lamports cannot be vanished. assert!(!token_account.base.is_native()); - // A deposit amount must be a 48-bit number - if amount >> MAXIMUM_DEPOSIT_TRANSFER_AMOUNT_BIT_LENGTH > 0 { - return Err(TokenError::MaximumDepositAmountExceeded.into()); - } - token_account.base.amount = token_account .base .amount .checked_sub(amount) .ok_or(TokenError::Overflow)?; - token_account.pack_base(); let mut confidential_transfer_account = token_account.get_extension_mut::()?; - check_destination_confidential_account(confidential_transfer_account)?; + confidential_transfer_account.valid_as_destination()?; - // Divide the deposit amount into low 16 and high 48-bit numbers and then add to the - // appropriate pending ciphertexts - confidential_transfer_account.pending_balance_lo = ops::add_to( - &confidential_transfer_account.pending_balance_lo, - amount << PENDING_BALANCE_HI_BIT_LENGTH >> PENDING_BALANCE_HI_BIT_LENGTH, - ) - .ok_or(ProgramError::InvalidInstructionData)?; + // A deposit amount must be a 48-bit number + let (amount_lo, amount_hi) = verify_and_split_deposit_amount(amount)?; - confidential_transfer_account.pending_balance_hi = ops::add_to( - &confidential_transfer_account.pending_balance_hi, - amount >> PENDING_BALANCE_LO_BIT_LENGTH, - ) - .ok_or(ProgramError::InvalidInstructionData)?; + // Prevent unnecessary ciphertext arithmetic syscalls if `amount_lo` or `amount_hi` is zero + if amount_lo > 0 { + confidential_transfer_account.pending_balance_lo = + ops::add_to(&confidential_transfer_account.pending_balance_lo, amount_lo) + .ok_or(ProgramError::InvalidInstructionData)?; + } + if amount_hi > 0 { + confidential_transfer_account.pending_balance_hi = + ops::add_to(&confidential_transfer_account.pending_balance_hi, amount_hi) + .ok_or(ProgramError::InvalidInstructionData)?; + } - confidential_transfer_account.pending_balance_credit_counter = - (u64::from(confidential_transfer_account.pending_balance_credit_counter) - .checked_add(1) - .ok_or(ProgramError::InvalidInstructionData)?) - .into(); + confidential_transfer_account.increment_pending_balance_credit_counter()?; Ok(()) } +/// Verifies that a deposit amount is a 48-bit number and returns the least significant 16 bits and +/// most significant 32 bits of the amount. +#[cfg(feature = "zk-ops")] +fn verify_and_split_deposit_amount(amount: u64) -> Result<(u64, u64), TokenError> { + if amount >> MAXIMUM_DEPOSIT_TRANSFER_AMOUNT_BIT_LENGTH > 0 { + return Err(TokenError::MaximumDepositAmountExceeded); + } + let deposit_amount_lo = + amount << (64 - PENDING_BALANCE_LO_BIT_LENGTH) >> PENDING_BALANCE_HI_BIT_LENGTH; + let deposit_amount_hi = amount >> PENDING_BALANCE_LO_BIT_LENGTH; + + Ok((deposit_amount_lo, deposit_amount_hi)) +} + /// Processes a [Withdraw] instruction. #[cfg(feature = "zk-ops")] fn process_withdraw( @@ -428,14 +385,6 @@ fn process_withdraw( return Err(TokenError::NonTransferable.into()); } - let previous_instruction = - get_instruction_relative(proof_instruction_offset, instructions_sysvar_info)?; - - let proof_data = decode_proof_instruction::( - ProofInstruction::VerifyWithdraw, - &previous_instruction, - )?; - check_program_account(token_account_info.owner)?; let token_account_data = &mut token_account_info.data.borrow_mut(); let mut token_account = StateWithExtensionsMut::::unpack(token_account_data)?; @@ -461,28 +410,40 @@ fn process_withdraw( let mut confidential_transfer_account = token_account.get_extension_mut::()?; - check_source_confidential_account(confidential_transfer_account)?; + confidential_transfer_account.valid_as_source()?; + // Zero-knowledge proof certifies that the account has enough available balance to withdraw the + // amount. + let zkp_instruction = + get_instruction_relative(proof_instruction_offset, instructions_sysvar_info)?; + let proof_data = decode_proof_instruction::( + ProofInstruction::VerifyWithdraw, + &zkp_instruction, + )?; + // 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.encryption_pubkey != proof_data.pubkey { return Err(TokenError::ConfidentialTransferElGamalPubkeyMismatch.into()); } - confidential_transfer_account.available_balance = - ops::subtract_from(&confidential_transfer_account.available_balance, amount) - .ok_or(ProgramError::InvalidInstructionData)?; - + // Prevent unnecessary ciphertext arithmetic syscalls if the withdraw amount is zero + if amount > 0 { + confidential_transfer_account.available_balance = + ops::subtract_from(&confidential_transfer_account.available_balance, amount) + .ok_or(ProgramError::InvalidInstructionData)?; + } + // 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_data.final_ciphertext { return Err(TokenError::ConfidentialTransferBalanceMismatch.into()); } confidential_transfer_account.decryptable_available_balance = new_decryptable_available_balance; - token_account.base.amount = token_account .base .amount .checked_add(amount) .ok_or(TokenError::Overflow)?; - token_account.pack_base(); Ok(()) @@ -497,7 +458,7 @@ fn process_transfer( proof_instruction_offset: i64, ) -> ProgramResult { let account_info_iter = &mut accounts.iter(); - let token_account_info = next_account_info(account_info_iter)?; + let source_account_info = next_account_info(account_info_iter)?; let destination_token_account_info = next_account_info(account_info_iter)?; let mint_info = next_account_info(account_info_iter)?; let instructions_sysvar_info = next_account_info(account_info_iter)?; @@ -510,20 +471,31 @@ fn process_transfer( if mint.get_extension::().is_ok() { return Err(TokenError::NonTransferable.into()); } - let confidential_transfer_mint = mint.get_extension::()?; - let previous_instruction = - get_instruction_relative(proof_instruction_offset, instructions_sysvar_info)?; + // A `Transfer` instruction must be accompanied by a zero-knowledge proof instruction that + // certify the validity of the transfer amounts. The kind of zero-knowledge proof instruction + // depends on whether a transfer incurs a fee or not. + // - If the mint is not extended for fees or the instruction is for a self-transfer, then + // transfer fee is not required. + // - If the mint is extended for fees and the instruction is not a self-transfer, then + // transfer fee is required. if mint.get_extension::().is_err() - || token_account_info.key == destination_token_account_info.key + || source_account_info.key == destination_token_account_info.key { - // mint is not extended for fees + // Transfer fee is not required. Decode the zero-knowledge proof as `TransferData`. + // + // The zero-knowledge proof certifies that: + // 1. the transfer amount is encrypted in the correct form + // 2. the source account has enough balance to send the transfer amount + let zkp_instruction = + get_instruction_relative(proof_instruction_offset, instructions_sysvar_info)?; let proof_data = decode_proof_instruction::( ProofInstruction::VerifyTransfer, - &previous_instruction, + &zkp_instruction, )?; - + // Check that the auditor encryption public key associated wth the confidential mint is + // consistent with what was actually used to generate the zkp. if proof_data.transfer_pubkeys.auditor_pubkey != confidential_transfer_mint.auditor_encryption_pubkey { @@ -541,7 +513,7 @@ fn process_transfer( process_source_for_transfer( program_id, - token_account_info, + source_account_info, mint_info, authority_info, account_info_iter.as_slice(), @@ -570,21 +542,25 @@ fn process_transfer( None, )?; } else { - // mint is extended for fees - let transfer_fee_config = mint.get_extension::()?; - + // Transfer fee is required. Decode the zero-knowledge proof as `TransferWithFeeData`. + // + // The zero-knowledge proof certifies that: + // 1. the transfer amount is encrypted in the correct form + // 2. the source account has enough balance to send the transfer amount + // 3. the transfer fee is computed correctly and encrypted in the correct form + let zkp_instruction = + get_instruction_relative(proof_instruction_offset, instructions_sysvar_info)?; let proof_data = decode_proof_instruction::( ProofInstruction::VerifyTransferWithFee, - &previous_instruction, + &zkp_instruction, )?; - + // Check that the encryption public keys associated with the confidential extension mint + // are consistent with the keys that were used to generate the zkp. if proof_data.transfer_with_fee_pubkeys.auditor_pubkey != confidential_transfer_mint.auditor_encryption_pubkey { return Err(TokenError::ConfidentialTransferElGamalPubkeyMismatch.into()); } - - // `withdraw_withheld_authority` ElGamal pubkey in proof data and mint must match if proof_data .transfer_with_fee_pubkeys .withdraw_withheld_authority_pubkey @@ -592,64 +568,48 @@ fn process_transfer( { return Err(TokenError::ConfidentialTransferElGamalPubkeyMismatch.into()); } - - // fee parameters in proof data and mint must match - let epoch = Clock::get()?.epoch; - let (maximum_fee, transfer_fee_basis_points) = - if u64::from(transfer_fee_config.newer_transfer_fee.epoch) < epoch { - ( - u64::from(transfer_fee_config.older_transfer_fee.maximum_fee), - u16::from( - transfer_fee_config - .older_transfer_fee - .transfer_fee_basis_points, - ), - ) - } else { - ( - u64::from(transfer_fee_config.newer_transfer_fee.maximum_fee), - u16::from( - transfer_fee_config - .newer_transfer_fee - .transfer_fee_basis_points, - ), - ) - }; - - if u64::from(proof_data.fee_parameters.maximum_fee) != maximum_fee - || u16::from(proof_data.fee_parameters.fee_rate_basis_points) - != transfer_fee_basis_points + // Check that the fee parameters in the mint are consistent with what were used to generate + // the zkp. + let transfer_fee_config = mint.get_extension::()?; + let fee_parameters = transfer_fee_config.get_epoch_fee(Clock::get()?.epoch); + if u64::from(fee_parameters.maximum_fee) != u64::from(proof_data.fee_parameters.maximum_fee) + || u16::from(fee_parameters.transfer_fee_basis_points) + != u16::from(proof_data.fee_parameters.fee_rate_basis_points) { return Err(TokenError::FeeParametersMismatch.into()); } - let source_ciphertext_lo = EncryptedBalance::from(( + // From the proof data, decode lo and hi transfer amounts encrypted under the source + // encryption public key + let source_transfer_amount_lo = EncryptedBalance::from(( proof_data.ciphertext_lo.commitment, proof_data.ciphertext_lo.source_handle, )); - let source_ciphertext_hi = EncryptedBalance::from(( + let source_transfer_amount_hi = EncryptedBalance::from(( proof_data.ciphertext_hi.commitment, proof_data.ciphertext_hi.source_handle, )); process_source_for_transfer( program_id, - token_account_info, + source_account_info, mint_info, authority_info, account_info_iter.as_slice(), &proof_data.transfer_with_fee_pubkeys.source_pubkey, - &source_ciphertext_lo, - &source_ciphertext_hi, + &source_transfer_amount_lo, + &source_transfer_amount_hi, &proof_data.new_source_ciphertext, new_source_decryptable_available_balance, )?; - let destination_ciphertext_lo = EncryptedBalance::from(( + // From the proof datay, decode lo and hi transfer amounts encrypted under the destination + // encryption public key + let destination_transfer_amount_lo = EncryptedBalance::from(( proof_data.ciphertext_lo.commitment, proof_data.ciphertext_lo.destination_handle, )); - let destination_ciphertext_hi = EncryptedBalance::from(( + let destination_transfer_amount_hi = EncryptedBalance::from(( proof_data.ciphertext_hi.commitment, proof_data.ciphertext_hi.destination_handle, )); @@ -658,8 +618,8 @@ fn process_transfer( destination_token_account_info, mint_info, &proof_data.transfer_with_fee_pubkeys.destination_pubkey, - &destination_ciphertext_lo, - &destination_ciphertext_hi, + &destination_transfer_amount_lo, + &destination_transfer_amount_hi, Some((&proof_data.fee_ciphertext_lo, &proof_data.fee_ciphertext_hi)), )?; } @@ -671,19 +631,19 @@ fn process_transfer( #[cfg(feature = "zk-ops")] fn process_source_for_transfer( program_id: &Pubkey, - token_account_info: &AccountInfo, + source_account_info: &AccountInfo, mint_info: &AccountInfo, authority_info: &AccountInfo, signers: &[AccountInfo], source_encryption_pubkey: &EncryptionPubkey, - source_ciphertext_lo: &EncryptedBalance, - source_ciphertext_hi: &EncryptedBalance, + source_transfer_amount_lo: &EncryptedBalance, + source_transfer_amount_hi: &EncryptedBalance, expected_new_source_available_balance: &EncryptedBalance, new_source_decryptable_available_balance: DecryptableBalance, ) -> ProgramResult { - check_program_account(token_account_info.owner)?; + check_program_account(source_account_info.owner)?; let authority_info_data_len = authority_info.data_len(); - let token_account_data = &mut token_account_info.data.borrow_mut(); + let token_account_data = &mut source_account_info.data.borrow_mut(); let mut token_account = StateWithExtensionsMut::::unpack(token_account_data)?; Processor::validate_owner( @@ -704,21 +664,23 @@ fn process_source_for_transfer( let mut confidential_transfer_account = token_account.get_extension_mut::()?; - check_source_confidential_account(confidential_transfer_account)?; + confidential_transfer_account.valid_as_source()?; + // Check that the source encryption public key is consistent with what was actually used to + // generate the zkp. if *source_encryption_pubkey != confidential_transfer_account.encryption_pubkey { return Err(TokenError::ConfidentialTransferElGamalPubkeyMismatch.into()); } - let new_source_available_balance = { - ops::subtract_with_lo_hi( - &confidential_transfer_account.available_balance, - source_ciphertext_lo, - source_ciphertext_hi, - ) - .ok_or(ProgramError::InvalidInstructionData)? - }; + let new_source_available_balance = ops::subtract_with_lo_hi( + &confidential_transfer_account.available_balance, + source_transfer_amount_lo, + source_transfer_amount_hi, + ) + .ok_or(ProgramError::InvalidInstructionData)?; + // Check that the computed available balance is consistent with what was actually used to + // generate the zkp on the client side. if new_source_available_balance != *expected_new_source_available_balance { return Err(TokenError::ConfidentialTransferBalanceMismatch.into()); } @@ -735,8 +697,8 @@ fn process_destination_for_transfer( destination_token_account_info: &AccountInfo, mint_info: &AccountInfo, destination_encryption_pubkey: &EncryptionPubkey, - destination_ciphertext_lo: &EncryptedBalance, - destination_ciphertext_hi: &EncryptedBalance, + destination_transfer_amount_lo: &EncryptedBalance, + destination_transfer_amount_hi: &EncryptedBalance, encrypted_fee: Option<(&EncryptedFee, &EncryptedFee)>, ) -> ProgramResult { check_program_account(destination_token_account_info.owner)?; @@ -758,89 +720,73 @@ fn process_destination_for_transfer( let mut destination_confidential_transfer_account = destination_token_account.get_extension_mut::()?; - check_destination_confidential_account(destination_confidential_transfer_account)?; + destination_confidential_transfer_account.valid_as_destination()?; if *destination_encryption_pubkey != destination_confidential_transfer_account.encryption_pubkey { return Err(TokenError::ConfidentialTransferElGamalPubkeyMismatch.into()); } - let new_destination_pending_balance_lo = ops::add( + destination_confidential_transfer_account.pending_balance_lo = ops::add( &destination_confidential_transfer_account.pending_balance_lo, - destination_ciphertext_lo, + destination_transfer_amount_lo, ) .ok_or(ProgramError::InvalidInstructionData)?; - let new_destination_pending_balance_hi = ops::add( + destination_confidential_transfer_account.pending_balance_hi = ops::add( &destination_confidential_transfer_account.pending_balance_hi, - destination_ciphertext_hi, + destination_transfer_amount_hi, ) .ok_or(ProgramError::InvalidInstructionData)?; - let new_destination_pending_balance_credit_counter = - u64::from(destination_confidential_transfer_account.pending_balance_credit_counter) - .checked_add(1) - .ok_or(ProgramError::InvalidInstructionData)?; + destination_confidential_transfer_account.increment_pending_balance_credit_counter()?; - destination_confidential_transfer_account.pending_balance_lo = - new_destination_pending_balance_lo; - destination_confidential_transfer_account.pending_balance_hi = - new_destination_pending_balance_hi; - destination_confidential_transfer_account.pending_balance_credit_counter = - new_destination_pending_balance_credit_counter.into(); - - // update destination account withheld fees + // Process transfer fee if let Some((ciphertext_fee_lo, ciphertext_fee_hi)) = encrypted_fee { - // subtract fee from destination pending balance - let ciphertext_fee_lo_destination: EncryptedWithheldAmount = ( + // Decode lo and hi fee amounts encrypted under the destination encryption public key + let destination_fee_lo: EncryptedWithheldAmount = ( ciphertext_fee_lo.commitment, ciphertext_fee_lo.destination_handle, ) .into(); - - let ciphertext_fee_hi_destination: EncryptedWithheldAmount = ( + let destination_fee_hi: EncryptedWithheldAmount = ( ciphertext_fee_hi.commitment, ciphertext_fee_hi.destination_handle, ) .into(); - let new_destination_pending_balance_lo = ops::subtract( + // Subtract the fee amount from the destination pending balance + destination_confidential_transfer_account.pending_balance_lo = ops::subtract( &destination_confidential_transfer_account.pending_balance_lo, - &ciphertext_fee_lo_destination, + &destination_fee_lo, ) .ok_or(ProgramError::InvalidInstructionData)?; - - let new_destination_pending_balance_hi = ops::subtract( + destination_confidential_transfer_account.pending_balance_hi = ops::subtract( &destination_confidential_transfer_account.pending_balance_hi, - &ciphertext_fee_hi_destination, + &destination_fee_hi, ) .ok_or(ProgramError::InvalidInstructionData)?; - // add encrypted fee to current withheld fee - let ciphertext_fee_lo_withheld_authority: EncryptedWithheldAmount = ( + // Decode lo and hi fee amounts encrypted under the withdraw authority encryption public + // key + let withdraw_withheld_authority_fee_lo: EncryptedWithheldAmount = ( ciphertext_fee_lo.commitment, ciphertext_fee_lo.withdraw_withheld_authority_handle, ) .into(); - - let ciphertext_fee_hi_withheld_authority: EncryptedWithheldAmount = ( + let withdraw_withheld_authority_fee_hi: EncryptedWithheldAmount = ( ciphertext_fee_hi.commitment, ciphertext_fee_hi.withdraw_withheld_authority_handle, ) .into(); - let new_withheld_amount = ops::add_with_lo_hi( + // Add the fee amount to the destination withheld fee + destination_confidential_transfer_account.withheld_amount = ops::add_with_lo_hi( &destination_confidential_transfer_account.withheld_amount, - &ciphertext_fee_lo_withheld_authority, - &ciphertext_fee_hi_withheld_authority, + &withdraw_withheld_authority_fee_lo, + &withdraw_withheld_authority_fee_hi, ) .ok_or(ProgramError::InvalidInstructionData)?; - - destination_confidential_transfer_account.pending_balance_lo = - new_destination_pending_balance_lo; - destination_confidential_transfer_account.pending_balance_hi = - new_destination_pending_balance_hi; - destination_confidential_transfer_account.withheld_amount = new_withheld_amount; } Ok(()) @@ -1006,30 +952,31 @@ fn process_withdraw_withheld_tokens_from_mint( } let mut destination_confidential_transfer_account = destination_account.get_extension_mut::()?; - check_destination_confidential_account(destination_confidential_transfer_account)?; + destination_confidential_transfer_account.valid_as_destination()?; - // verify consistency of proof data - let previous_instruction = + // Zero-knowledge proof certifies that the exact withheld amount is credited to the source + // account. + let zkp_instruction = get_instruction_relative(proof_instruction_offset, instructions_sysvar_info)?; let proof_data = decode_proof_instruction::( ProofInstruction::VerifyWithdrawWithheldTokens, - &previous_instruction, + &zkp_instruction, )?; - - // withdraw withheld authority ElGamal pubkey should match in the proof data and mint + // Checks that the withdraw authority encryption public key associated with the mint is + // consistent with what was actually used to generate the zkp. if proof_data.withdraw_withheld_authority_pubkey != confidential_transfer_mint.withdraw_withheld_authority_encryption_pubkey { return Err(TokenError::ConfidentialTransferElGamalPubkeyMismatch.into()); } - - // destination ElGamal pubkey should match in the proof data and destination account + // Checks that the encryption public key associated with the destination account is consistent + // with what was actually used to generate the zkp. if proof_data.destination_pubkey != destination_confidential_transfer_account.encryption_pubkey { return Err(TokenError::ConfidentialTransferElGamalPubkeyMismatch.into()); } - - // withheld amount ciphertext must match in the proof data and mint + // Checks that the withheld amount ciphertext is consistent with the ciphertext data that was + // actually used to generate the zkp. if proof_data.withdraw_withheld_authority_ciphertext != confidential_transfer_mint.withheld_amount { @@ -1037,22 +984,16 @@ fn process_withdraw_withheld_tokens_from_mint( } // The proof data contains the mint withheld amount encrypted under the destination ElGamal pubkey. - // This amount should be added to the destination pending balance. - let new_destination_pending_balance = ops::add( + // This amount is added to the destination pending balance. + destination_confidential_transfer_account.pending_balance_lo = ops::add( &destination_confidential_transfer_account.pending_balance_lo, &proof_data.destination_ciphertext, ) .ok_or(ProgramError::InvalidInstructionData)?; - destination_confidential_transfer_account.pending_balance_lo = new_destination_pending_balance; + destination_confidential_transfer_account.increment_pending_balance_credit_counter()?; - destination_confidential_transfer_account.pending_balance_credit_counter = - (u64::from(destination_confidential_transfer_account.pending_balance_credit_counter) - .checked_add(1) - .ok_or(ProgramError::InvalidInstructionData)?) - .into(); - - // fee is now withdrawn, so zero out mint withheld amount + // Fee is now withdrawn, so zero out the mint withheld amount. confidential_transfer_mint.withheld_amount = EncryptedWithheldAmount::zeroed(); Ok(()) @@ -1104,7 +1045,7 @@ fn process_withdraw_withheld_tokens_from_accounts( return Err(TokenError::AccountFrozen.into()); } - // sum up the withheld amounts in all the accounts + // Sum up the withheld amounts in all the accounts. let mut aggregate_withheld_amount = EncryptedWithheldAmount::zeroed(); for account_info in &account_infos[num_signers..] { // self-harvest, can't double-borrow the underlying data @@ -1137,49 +1078,45 @@ fn process_withdraw_withheld_tokens_from_accounts( let mut destination_confidential_transfer_account = destination_account.get_extension_mut::()?; - check_destination_confidential_account(destination_confidential_transfer_account)?; + destination_confidential_transfer_account.valid_as_destination()?; - // verify consistency of proof data - let previous_instruction = + // Zero-knowledge proof certifies that the exact aggregate withheld amount is credited to the + // source account. + let zkp_instruction = get_instruction_relative(proof_instruction_offset, instructions_sysvar_info)?; let proof_data = decode_proof_instruction::( ProofInstruction::VerifyWithdrawWithheldTokens, - &previous_instruction, + &zkp_instruction, )?; - - // withdraw withheld authority ElGamal pubkey should match in the proof data and mint + // Checks that the withdraw authority encryption public key associated with the mint is + // consistent with what was actually used to generate the zkp. let confidential_transfer_mint = mint.get_extension_mut::()?; if proof_data.withdraw_withheld_authority_pubkey != confidential_transfer_mint.withdraw_withheld_authority_encryption_pubkey { return Err(TokenError::ConfidentialTransferElGamalPubkeyMismatch.into()); } - - // destination ElGamal pubkey should match in the proof data and destination account + // Checks that the encryption public key associated with the destination account is consistent + // with what was actually used to generate the zkp. if proof_data.destination_pubkey != destination_confidential_transfer_account.encryption_pubkey { return Err(TokenError::ConfidentialTransferElGamalPubkeyMismatch.into()); } - - // withheld amount ciphertext must match in the proof data and mint + // Checks that the withheld amount ciphertext is consistent with the ciphertext data that was + // actually used to generate the zkp. if proof_data.withdraw_withheld_authority_ciphertext != aggregate_withheld_amount { return Err(TokenError::ConfidentialTransferBalanceMismatch.into()); } - // add the sum of the withheld fees to destination pending balance - let new_destination_pending_balance = ops::add( + // The proof data contains the mint withheld amount encrypted under the destination ElGamal pubkey. + // This amount is added to the destination pending balance. + destination_confidential_transfer_account.pending_balance_lo = ops::add( &destination_confidential_transfer_account.pending_balance_lo, &proof_data.destination_ciphertext, ) .ok_or(ProgramError::InvalidInstructionData)?; - destination_confidential_transfer_account.pending_balance_credit_counter = - (u64::from(destination_confidential_transfer_account.pending_balance_credit_counter) - .checked_add(1) - .ok_or(ProgramError::InvalidInstructionData)?) - .into(); - - destination_confidential_transfer_account.pending_balance_lo = new_destination_pending_balance; + destination_confidential_transfer_account.increment_pending_balance_credit_counter()?; Ok(()) }