Skip to content
This repository was archived by the owner on Mar 11, 2025. It is now read-only.

Commit c3137af

Browse files
samkim-cryptosamkim-crypto
authored
token-2022: add tests for withdraw of confidential withheld tokens from accounts (#3280)
* token-2022: add tests for withdraw of confidential withheld tokens from accounts * token-2022: cargo fmt * token-2022: cargo clippy
1 parent 0f316fb commit c3137af

File tree

3 files changed

+140
-20
lines changed

3 files changed

+140
-20
lines changed

‎token/client/src/token.rs‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1386,14 +1386,9 @@ where
13861386
withdraw_withheld_authority_elgamal_keypair: &ElGamalKeypair,
13871387
destination_token_account: &Pubkey,
13881388
amount: u64,
1389+
encrypted_aggregate_withheld_amount: &ElGamalCiphertext,
13891390
sources: &[&Pubkey],
13901391
) -> TokenResult<T::Output> {
1391-
let mint_state = self.get_mint_info().await.unwrap();
1392-
1393-
let ct_mint = mint_state
1394-
.get_extension::<confidential_transfer::ConfidentialTransferMint>()
1395-
.unwrap();
1396-
13971392
let destination_state = self
13981393
.get_account_info(destination_token_account)
13991394
.await
@@ -1405,7 +1400,7 @@ where
14051400
let proof_data = confidential_transfer::instruction::WithdrawWithheldTokensData::new(
14061401
withdraw_withheld_authority_elgamal_keypair,
14071402
&destination_extension.encryption_pubkey.try_into().unwrap(),
1408-
&ct_mint.withheld_amount.try_into().unwrap(),
1403+
encrypted_aggregate_withheld_amount,
14091404
amount,
14101405
)
14111406
.map_err(TokenError::Proof)?;

‎token/program-2022-test/tests/confidential_transfer.rs‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1254,3 +1254,110 @@ async fn ct_withdraw_withheld_tokens_from_mint() {
12541254
)
12551255
.await;
12561256
}
1257+
1258+
#[tokio::test]
1259+
async fn ct_withdraw_withheld_tokens_from_accounts() {
1260+
let ConfidentialTransferMintWithKeypairs {
1261+
ct_mint,
1262+
ct_mint_withdraw_withheld_authority_encryption_keypair,
1263+
..
1264+
} = ConfidentialTransferMintWithKeypairs::new();
1265+
1266+
let ct_mint_withdraw_withheld_authority = Keypair::new();
1267+
1268+
let mut context = TestContext::new().await;
1269+
context
1270+
.init_token_with_mint(vec![
1271+
ExtensionInitializationParams::TransferFeeConfig {
1272+
transfer_fee_config_authority: Some(Pubkey::new_unique()),
1273+
withdraw_withheld_authority: Some(ct_mint_withdraw_withheld_authority.pubkey()),
1274+
transfer_fee_basis_points: TEST_FEE_BASIS_POINTS,
1275+
maximum_fee: TEST_MAXIMUM_FEE,
1276+
},
1277+
ExtensionInitializationParams::ConfidentialTransferMint { ct_mint },
1278+
])
1279+
.await
1280+
.unwrap();
1281+
1282+
let TokenContext {
1283+
token,
1284+
alice,
1285+
bob,
1286+
mint_authority,
1287+
decimals,
1288+
..
1289+
} = context.token_context.unwrap();
1290+
1291+
let epoch_info = test_epoch_info();
1292+
1293+
let alice_meta =
1294+
ConfidentialTokenAccountMeta::with_tokens(&token, &alice, &mint_authority, 100, decimals)
1295+
.await;
1296+
let bob_meta = ConfidentialTokenAccountMeta::new(&token, &bob).await;
1297+
1298+
// Test fee is 2.5% so the withheld fees should be 3
1299+
token
1300+
.confidential_transfer_transfer_with_fee(
1301+
&alice_meta.token_account,
1302+
&bob_meta.token_account,
1303+
&alice,
1304+
100,
1305+
100,
1306+
&alice_meta.elgamal_keypair,
1307+
alice_meta.ae_key.encrypt(0_u64),
1308+
&epoch_info,
1309+
)
1310+
.await
1311+
.unwrap();
1312+
1313+
let state = token
1314+
.get_account_info(&bob_meta.token_account)
1315+
.await
1316+
.unwrap();
1317+
let extension = state
1318+
.get_extension::<ConfidentialTransferAccount>()
1319+
.unwrap();
1320+
1321+
assert_eq!(
1322+
extension
1323+
.withheld_amount
1324+
.decrypt(&ct_mint_withdraw_withheld_authority_encryption_keypair.secret),
1325+
Some(3),
1326+
);
1327+
1328+
token
1329+
.confidential_transfer_withdraw_withheld_tokens_from_accounts(
1330+
&ct_mint_withdraw_withheld_authority,
1331+
&ct_mint_withdraw_withheld_authority_encryption_keypair,
1332+
&alice_meta.token_account,
1333+
3_u64,
1334+
&extension.withheld_amount.try_into().unwrap(),
1335+
&[&bob_meta.token_account],
1336+
)
1337+
.await
1338+
.unwrap();
1339+
1340+
bob_meta
1341+
.check_balances(
1342+
&token,
1343+
ConfidentialTokenAccountBalances {
1344+
pending_balance_lo: 97,
1345+
pending_balance_hi: 0,
1346+
available_balance: 0,
1347+
decryptable_available_balance: 0,
1348+
},
1349+
)
1350+
.await;
1351+
1352+
alice_meta
1353+
.check_balances(
1354+
&token,
1355+
ConfidentialTokenAccountBalances {
1356+
pending_balance_lo: 3,
1357+
pending_balance_hi: 0,
1358+
available_balance: 0,
1359+
decryptable_available_balance: 0,
1360+
},
1361+
)
1362+
.await;
1363+
}

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

Lines changed: 31 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -552,11 +552,11 @@ fn process_transfer(
552552
return Err(TokenError::FeeParametersMismatch.into());
553553
}
554554

555-
let ciphertext_lo = EncryptedBalance::from((
555+
let source_ciphertext_lo = EncryptedBalance::from((
556556
proof_data.ciphertext_lo.commitment,
557557
proof_data.ciphertext_lo.source_handle,
558558
));
559-
let ciphertext_hi = EncryptedBalance::from((
559+
let source_ciphertext_hi = EncryptedBalance::from((
560560
proof_data.ciphertext_hi.commitment,
561561
proof_data.ciphertext_hi.source_handle,
562562
));
@@ -568,11 +568,20 @@ fn process_transfer(
568568
authority_info,
569569
account_info_iter.as_slice(),
570570
&proof_data.transfer_with_fee_pubkeys.source_pubkey,
571-
&ciphertext_lo,
572-
&ciphertext_hi,
571+
&source_ciphertext_lo,
572+
&source_ciphertext_hi,
573573
new_source_decryptable_available_balance,
574574
)?;
575575

576+
let destination_ciphertext_lo = EncryptedBalance::from((
577+
proof_data.ciphertext_lo.commitment,
578+
proof_data.ciphertext_lo.destination_handle,
579+
));
580+
let destination_ciphertext_hi = EncryptedBalance::from((
581+
proof_data.ciphertext_hi.commitment,
582+
proof_data.ciphertext_hi.destination_handle,
583+
));
584+
576585
let fee_ciphertext = if token_account_info.key == destination_token_account_info.key {
577586
None
578587
} else {
@@ -583,8 +592,8 @@ fn process_transfer(
583592
destination_token_account_info,
584593
mint_info,
585594
&proof_data.transfer_with_fee_pubkeys.destination_pubkey,
586-
&ciphertext_lo,
587-
&ciphertext_hi,
595+
&destination_ciphertext_lo,
596+
&destination_ciphertext_hi,
588597
fee_ciphertext,
589598
)?;
590599
} else {
@@ -600,11 +609,11 @@ fn process_transfer(
600609
return Err(TokenError::ConfidentialTransferElGamalPubkeyMismatch.into());
601610
}
602611

603-
let ciphertext_lo = EncryptedBalance::from((
612+
let source_ciphertext_lo = EncryptedBalance::from((
604613
proof_data.ciphertext_lo.commitment,
605614
proof_data.ciphertext_lo.source_handle,
606615
));
607-
let ciphertext_hi = EncryptedBalance::from((
616+
let source_ciphertext_hi = EncryptedBalance::from((
608617
proof_data.ciphertext_hi.commitment,
609618
proof_data.ciphertext_hi.source_handle,
610619
));
@@ -616,17 +625,26 @@ fn process_transfer(
616625
authority_info,
617626
account_info_iter.as_slice(),
618627
&proof_data.transfer_pubkeys.source_pubkey,
619-
&ciphertext_lo,
620-
&ciphertext_hi,
628+
&source_ciphertext_lo,
629+
&source_ciphertext_hi,
621630
new_source_decryptable_available_balance,
622631
)?;
623632

633+
let destination_ciphertext_lo = EncryptedBalance::from((
634+
proof_data.ciphertext_lo.commitment,
635+
proof_data.ciphertext_lo.destination_handle,
636+
));
637+
let destination_ciphertext_hi = EncryptedBalance::from((
638+
proof_data.ciphertext_hi.commitment,
639+
proof_data.ciphertext_hi.destination_handle,
640+
));
641+
624642
process_destination_for_transfer(
625643
destination_token_account_info,
626644
mint_info,
627645
&proof_data.transfer_pubkeys.destination_pubkey,
628-
&ciphertext_lo,
629-
&ciphertext_hi,
646+
&destination_ciphertext_lo,
647+
&destination_ciphertext_hi,
630648
None,
631649
)?;
632650
}
@@ -1069,7 +1087,7 @@ fn process_withdraw_withheld_tokens_from_accounts(
10691087
// add the sum of the withheld fees to destination pending balance
10701088
let new_destination_pending_balance = ops::add(
10711089
&destination_confidential_transfer_account.pending_balance_lo,
1072-
&aggregate_withheld_amount,
1090+
&proof_data.destination_ciphertext,
10731091
)
10741092
.ok_or(ProgramError::InvalidInstructionData)?;
10751093

0 commit comments

Comments
 (0)