Skip to content

Commit

Permalink
token-2022: add tests for withdraw of confidential withheld tokens fr…
Browse files Browse the repository at this point in the history
…om accounts (#3280)

* token-2022: add tests for withdraw of confidential withheld tokens from accounts

* token-2022: cargo fmt

* token-2022: cargo clippy
  • Loading branch information
samkim-crypto authored Jun 21, 2022
1 parent 0f316fb commit c3137af
Show file tree
Hide file tree
Showing 3 changed files with 140 additions and 20 deletions.
9 changes: 2 additions & 7 deletions token/client/src/token.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<T::Output> {
let mint_state = self.get_mint_info().await.unwrap();

let ct_mint = mint_state
.get_extension::<confidential_transfer::ConfidentialTransferMint>()
.unwrap();

let destination_state = self
.get_account_info(destination_token_account)
.await
Expand All @@ -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)?;
Expand Down
107 changes: 107 additions & 0 deletions token/program-2022-test/tests/confidential_transfer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::<ConfidentialTransferAccount>()
.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;
}
44 changes: 31 additions & 13 deletions token/program-2022/src/extension/confidential_transfer/processor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
));
Expand All @@ -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 {
Expand All @@ -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 {
Expand All @@ -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,
));
Expand All @@ -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,
)?;
}
Expand Down Expand Up @@ -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)?;

Expand Down

0 comments on commit c3137af

Please sign in to comment.