diff --git a/token/client/src/token.rs b/token/client/src/token.rs index 3ecd99757f9..64a8b01a5c0 100644 --- a/token/client/src/token.rs +++ b/token/client/src/token.rs @@ -1386,14 +1386,9 @@ where withdraw_withheld_authority_elgamal_keypair: &ElGamalKeypair, destination_token_account: &Pubkey, amount: u64, + encrypted_aggregate_withheld_amount: &ElGamalCiphertext, sources: &[&Pubkey], ) -> TokenResult { - let mint_state = self.get_mint_info().await.unwrap(); - - let ct_mint = mint_state - .get_extension::() - .unwrap(); - let destination_state = self .get_account_info(destination_token_account) .await @@ -1405,7 +1400,7 @@ where let proof_data = confidential_transfer::instruction::WithdrawWithheldTokensData::new( withdraw_withheld_authority_elgamal_keypair, &destination_extension.encryption_pubkey.try_into().unwrap(), - &ct_mint.withheld_amount.try_into().unwrap(), + encrypted_aggregate_withheld_amount, amount, ) .map_err(TokenError::Proof)?; diff --git a/token/program-2022-test/tests/confidential_transfer.rs b/token/program-2022-test/tests/confidential_transfer.rs index ab6708bc60b..2f0f035a0ff 100644 --- a/token/program-2022-test/tests/confidential_transfer.rs +++ b/token/program-2022-test/tests/confidential_transfer.rs @@ -1254,3 +1254,110 @@ async fn ct_withdraw_withheld_tokens_from_mint() { ) .await; } + +#[tokio::test] +async fn ct_withdraw_withheld_tokens_from_accounts() { + let ConfidentialTransferMintWithKeypairs { + ct_mint, + ct_mint_withdraw_withheld_authority_encryption_keypair, + .. + } = ConfidentialTransferMintWithKeypairs::new(); + + let ct_mint_withdraw_withheld_authority = Keypair::new(); + + let mut context = TestContext::new().await; + context + .init_token_with_mint(vec![ + ExtensionInitializationParams::TransferFeeConfig { + transfer_fee_config_authority: Some(Pubkey::new_unique()), + withdraw_withheld_authority: Some(ct_mint_withdraw_withheld_authority.pubkey()), + transfer_fee_basis_points: TEST_FEE_BASIS_POINTS, + maximum_fee: TEST_MAXIMUM_FEE, + }, + ExtensionInitializationParams::ConfidentialTransferMint { ct_mint }, + ]) + .await + .unwrap(); + + let TokenContext { + token, + alice, + bob, + mint_authority, + decimals, + .. + } = context.token_context.unwrap(); + + let epoch_info = test_epoch_info(); + + let alice_meta = + ConfidentialTokenAccountMeta::with_tokens(&token, &alice, &mint_authority, 100, decimals) + .await; + let bob_meta = ConfidentialTokenAccountMeta::new(&token, &bob).await; + + // Test fee is 2.5% so the withheld fees should be 3 + token + .confidential_transfer_transfer_with_fee( + &alice_meta.token_account, + &bob_meta.token_account, + &alice, + 100, + 100, + &alice_meta.elgamal_keypair, + alice_meta.ae_key.encrypt(0_u64), + &epoch_info, + ) + .await + .unwrap(); + + let state = token + .get_account_info(&bob_meta.token_account) + .await + .unwrap(); + let extension = state + .get_extension::() + .unwrap(); + + assert_eq!( + extension + .withheld_amount + .decrypt(&ct_mint_withdraw_withheld_authority_encryption_keypair.secret), + Some(3), + ); + + token + .confidential_transfer_withdraw_withheld_tokens_from_accounts( + &ct_mint_withdraw_withheld_authority, + &ct_mint_withdraw_withheld_authority_encryption_keypair, + &alice_meta.token_account, + 3_u64, + &extension.withheld_amount.try_into().unwrap(), + &[&bob_meta.token_account], + ) + .await + .unwrap(); + + bob_meta + .check_balances( + &token, + ConfidentialTokenAccountBalances { + pending_balance_lo: 97, + pending_balance_hi: 0, + available_balance: 0, + decryptable_available_balance: 0, + }, + ) + .await; + + alice_meta + .check_balances( + &token, + ConfidentialTokenAccountBalances { + pending_balance_lo: 3, + pending_balance_hi: 0, + available_balance: 0, + decryptable_available_balance: 0, + }, + ) + .await; +} diff --git a/token/program-2022/src/extension/confidential_transfer/processor.rs b/token/program-2022/src/extension/confidential_transfer/processor.rs index 7300d24347d..9713b6c11fe 100644 --- a/token/program-2022/src/extension/confidential_transfer/processor.rs +++ b/token/program-2022/src/extension/confidential_transfer/processor.rs @@ -552,11 +552,11 @@ fn process_transfer( return Err(TokenError::FeeParametersMismatch.into()); } - let ciphertext_lo = EncryptedBalance::from(( + let source_ciphertext_lo = EncryptedBalance::from(( proof_data.ciphertext_lo.commitment, proof_data.ciphertext_lo.source_handle, )); - let ciphertext_hi = EncryptedBalance::from(( + let source_ciphertext_hi = EncryptedBalance::from(( proof_data.ciphertext_hi.commitment, proof_data.ciphertext_hi.source_handle, )); @@ -568,11 +568,20 @@ fn process_transfer( authority_info, account_info_iter.as_slice(), &proof_data.transfer_with_fee_pubkeys.source_pubkey, - &ciphertext_lo, - &ciphertext_hi, + &source_ciphertext_lo, + &source_ciphertext_hi, new_source_decryptable_available_balance, )?; + let destination_ciphertext_lo = EncryptedBalance::from(( + proof_data.ciphertext_lo.commitment, + proof_data.ciphertext_lo.destination_handle, + )); + let destination_ciphertext_hi = EncryptedBalance::from(( + proof_data.ciphertext_hi.commitment, + proof_data.ciphertext_hi.destination_handle, + )); + let fee_ciphertext = if token_account_info.key == destination_token_account_info.key { None } else { @@ -583,8 +592,8 @@ fn process_transfer( destination_token_account_info, mint_info, &proof_data.transfer_with_fee_pubkeys.destination_pubkey, - &ciphertext_lo, - &ciphertext_hi, + &destination_ciphertext_lo, + &destination_ciphertext_hi, fee_ciphertext, )?; } else { @@ -600,11 +609,11 @@ fn process_transfer( return Err(TokenError::ConfidentialTransferElGamalPubkeyMismatch.into()); } - let ciphertext_lo = EncryptedBalance::from(( + let source_ciphertext_lo = EncryptedBalance::from(( proof_data.ciphertext_lo.commitment, proof_data.ciphertext_lo.source_handle, )); - let ciphertext_hi = EncryptedBalance::from(( + let source_ciphertext_hi = EncryptedBalance::from(( proof_data.ciphertext_hi.commitment, proof_data.ciphertext_hi.source_handle, )); @@ -616,17 +625,26 @@ fn process_transfer( authority_info, account_info_iter.as_slice(), &proof_data.transfer_pubkeys.source_pubkey, - &ciphertext_lo, - &ciphertext_hi, + &source_ciphertext_lo, + &source_ciphertext_hi, new_source_decryptable_available_balance, )?; + let destination_ciphertext_lo = EncryptedBalance::from(( + proof_data.ciphertext_lo.commitment, + proof_data.ciphertext_lo.destination_handle, + )); + let destination_ciphertext_hi = EncryptedBalance::from(( + proof_data.ciphertext_hi.commitment, + proof_data.ciphertext_hi.destination_handle, + )); + process_destination_for_transfer( destination_token_account_info, mint_info, &proof_data.transfer_pubkeys.destination_pubkey, - &ciphertext_lo, - &ciphertext_hi, + &destination_ciphertext_lo, + &destination_ciphertext_hi, None, )?; } @@ -1069,7 +1087,7 @@ fn process_withdraw_withheld_tokens_from_accounts( // add the sum of the withheld fees to destination pending balance let new_destination_pending_balance = ops::add( &destination_confidential_transfer_account.pending_balance_lo, - &aggregate_withheld_amount, + &proof_data.destination_ciphertext, ) .ok_or(ProgramError::InvalidInstructionData)?;