diff --git a/token/client/src/token.rs b/token/client/src/token.rs index 0a37488551f..94f3fe0a2b1 100644 --- a/token/client/src/token.rs +++ b/token/client/src/token.rs @@ -966,4 +966,40 @@ where ) .await } + + /// Enable confidential transfer `Deposit` and `Transfer` instructions for a token account + pub async fn confidential_transfer_enable_balance_credits( + &self, + token_account: &Pubkey, + authority: &S2, + ) -> TokenResult { + self.process_ixs( + &[confidential_transfer::instruction::enable_balance_credits( + &self.program_id, + token_account, + &authority.pubkey(), + &[], + )?], + &[authority], + ) + .await + } + + /// Disable confidential transfer `Deposit` and `Transfer` instructions for a token account + pub async fn confidential_transfer_disable_balance_credits( + &self, + token_account: &Pubkey, + authority: &S2, + ) -> TokenResult { + self.process_ixs( + &[confidential_transfer::instruction::disable_balance_credits( + &self.program_id, + token_account, + &authority.pubkey(), + &[], + )?], + &[authority], + ) + .await + } } diff --git a/token/program-2022-test/tests/confidential_transfer.rs b/token/program-2022-test/tests/confidential_transfer.rs index 2ab8cb3348c..69320f6e0c3 100644 --- a/token/program-2022-test/tests/confidential_transfer.rs +++ b/token/program-2022-test/tests/confidential_transfer.rs @@ -121,7 +121,7 @@ async fn ct_configure_token_account() { token .confidential_transfer_configure_token_account( - dbg!(&alice_token_account), + &alice_token_account, &alice, alice_elgamal_keypair.public, alice_ae_key.encrypt(0_u64), @@ -134,6 +134,7 @@ async fn ct_configure_token_account() { .get_extension::() .unwrap(); assert!(!bool::from(&extension.approved)); + assert!(bool::from(&extension.allow_balance_credits)); assert_eq!( extension.elgamal_pubkey, alice_elgamal_keypair.public.into() @@ -156,3 +157,59 @@ async fn ct_configure_token_account() { .unwrap(); assert!(bool::from(&extension.approved)); } + +#[tokio::test] +async fn ct_enable_disable_balance_credits() { + let ct_mint = ConfidentialTransferMint::default(); + let mut context = TestContext::new().await; + context + .init_token_with_mint(vec![ + ExtensionInitializationParams::ConfidentialTransferMint { ct_mint }, + ]) + .await + .unwrap(); + + let TokenContext { token, alice, .. } = context.token_context.unwrap(); + + let alice_token_account = token + .create_auxiliary_token_account_with_extension_space( + &alice, + &alice.pubkey(), + vec![ExtensionType::ConfidentialTransferAccount], + ) + .await + .unwrap(); + + let alice_elgamal_keypair = ElGamalKeypair::new_rand(); + let alice_ae_key = AeKey::new(&alice, &alice_token_account).unwrap(); + + token + .confidential_transfer_configure_token_account( + &alice_token_account, + &alice, + alice_elgamal_keypair.public, + alice_ae_key.encrypt(0_u64), + ) + .await + .unwrap(); + + token + .confidential_transfer_disable_balance_credits(&alice_token_account, &alice) + .await + .unwrap(); + let state = token.get_account_info(&alice_token_account).await.unwrap(); + let extension = state + .get_extension::() + .unwrap(); + assert!(!bool::from(&extension.allow_balance_credits)); + + token + .confidential_transfer_enable_balance_credits(&alice_token_account, &alice) + .await + .unwrap(); + let state = token.get_account_info(&alice_token_account).await.unwrap(); + let extension = state + .get_extension::() + .unwrap(); + assert!(bool::from(&extension.allow_balance_credits)); +} diff --git a/token/program-2022/src/extension/confidential_transfer/instruction.rs b/token/program-2022/src/extension/confidential_transfer/instruction.rs index 1966652cee2..580afb80bc2 100644 --- a/token/program-2022/src/extension/confidential_transfer/instruction.rs +++ b/token/program-2022/src/extension/confidential_transfer/instruction.rs @@ -653,7 +653,6 @@ pub fn transfer( /// Create a inner `ApplyPendingBalance` instruction /// /// This instruction is suitable for use with a cross-program `invoke` -#[allow(clippy::too_many_arguments)] pub fn inner_apply_pending_balance( token_program_id: &Pubkey, token_account: &Pubkey, @@ -705,13 +704,13 @@ pub fn apply_pending_balance( ]) } -/// Create a `EnableBalanceCredits` instruction -pub fn enable_balance_credits( +fn enable_or_disable_balance_credits( + instruction: ConfidentialTransferInstruction, token_program_id: &Pubkey, token_account: &Pubkey, authority: &Pubkey, multisig_signers: &[&Pubkey], -) -> Result, ProgramError> { +) -> Result { check_program_account(token_program_id)?; let mut accounts = vec![ AccountMeta::new(*token_account, false), @@ -722,36 +721,42 @@ pub fn enable_balance_credits( accounts.push(AccountMeta::new_readonly(**multisig_signer, true)); } - Ok(vec![encode_instruction( + Ok(encode_instruction( token_program_id, accounts, - ConfidentialTransferInstruction::EnableBalanceCredits, + instruction, &(), - )]) + )) +} + +/// Create a `EnableBalanceCredits` instruction +pub fn enable_balance_credits( + token_program_id: &Pubkey, + token_account: &Pubkey, + authority: &Pubkey, + multisig_signers: &[&Pubkey], +) -> Result { + enable_or_disable_balance_credits( + ConfidentialTransferInstruction::EnableBalanceCredits, + token_program_id, + token_account, + authority, + multisig_signers, + ) } /// Create a `DisableBalanceCredits` instruction -#[cfg(not(target_arch = "bpf"))] pub fn disable_balance_credits( token_program_id: &Pubkey, token_account: &Pubkey, authority: &Pubkey, multisig_signers: &[&Pubkey], -) -> Result, ProgramError> { - check_program_account(token_program_id)?; - let mut accounts = vec![ - AccountMeta::new(*token_account, false), - AccountMeta::new_readonly(*authority, multisig_signers.is_empty()), - ]; - - for multisig_signer in multisig_signers.iter() { - accounts.push(AccountMeta::new_readonly(**multisig_signer, true)); - } - - Ok(vec![encode_instruction( - token_program_id, - accounts, +) -> Result { + enable_or_disable_balance_credits( ConfidentialTransferInstruction::DisableBalanceCredits, - &(), - )]) + token_program_id, + token_account, + authority, + multisig_signers, + ) }