From b834e962f4919acd0c5c892166bb444289324b60 Mon Sep 17 00:00:00 2001 From: Jon C Date: Tue, 9 Jan 2024 23:11:58 +0100 Subject: [PATCH] Move confidential-transfer test types for reuse --- .../tests/confidential_transfer.rs | 185 +--------------- token/program-2022-test/tests/program_test.rs | 198 +++++++++++++++++- 2 files changed, 201 insertions(+), 182 deletions(-) diff --git a/token/program-2022-test/tests/confidential_transfer.rs b/token/program-2022-test/tests/confidential_transfer.rs index 70658ec92eb..91a5e7fafac 100644 --- a/token/program-2022-test/tests/confidential_transfer.rs +++ b/token/program-2022-test/tests/confidential_transfer.rs @@ -2,7 +2,9 @@ mod program_test; use { - program_test::{TestContext, TokenContext}, + program_test::{ + ConfidentialTokenAccountBalances, ConfidentialTokenAccountMeta, TestContext, TokenContext, + }, solana_program_test::tokio, solana_sdk::{ instruction::InstructionError, @@ -36,9 +38,8 @@ use { }, }, spl_token_client::{ - client::{SendTransaction, SimulateTransaction}, proof_generation::transfer_with_fee_split_proof_data, - token::{ExtensionInitializationParams, Token, TokenError as TokenClientError}, + token::{ExtensionInitializationParams, TokenError as TokenClientError}, }, std::{convert::TryInto, mem::size_of}, }; @@ -48,181 +49,6 @@ const TEST_MAXIMUM_FEE: u64 = 100; #[cfg(feature = "zk-ops")] const TEST_FEE_BASIS_POINTS: u16 = 250; -struct ConfidentialTokenAccountMeta { - token_account: Pubkey, - elgamal_keypair: ElGamalKeypair, - aes_key: AeKey, -} - -impl ConfidentialTokenAccountMeta { - async fn new( - token: &Token, - owner: &Keypair, - maximum_pending_balance_credit_counter: Option, - require_memo: bool, - require_fee: bool, - ) -> Self - where - T: SendTransaction + SimulateTransaction, - { - let token_account_keypair = Keypair::new(); - - let mut extensions = vec![ExtensionType::ConfidentialTransferAccount]; - if require_memo { - extensions.push(ExtensionType::MemoTransfer); - } - if require_fee { - extensions.push(ExtensionType::ConfidentialTransferFeeAmount); - } - - token - .create_auxiliary_token_account_with_extension_space( - &token_account_keypair, - &owner.pubkey(), - extensions, - ) - .await - .unwrap(); - let token_account = token_account_keypair.pubkey(); - - let elgamal_keypair = - ElGamalKeypair::new_from_signer(owner, &token_account.to_bytes()).unwrap(); - let aes_key = AeKey::new_from_signer(owner, &token_account.to_bytes()).unwrap(); - - token - .confidential_transfer_configure_token_account( - &token_account, - &owner.pubkey(), - None, - maximum_pending_balance_credit_counter, - &elgamal_keypair, - &aes_key, - &[owner], - ) - .await - .unwrap(); - - if require_memo { - token - .enable_required_transfer_memos(&token_account, &owner.pubkey(), &[owner]) - .await - .unwrap(); - } - - Self { - token_account, - elgamal_keypair, - aes_key, - } - } - - #[allow(clippy::too_many_arguments)] - #[cfg(feature = "zk-ops")] - async fn new_with_tokens( - token: &Token, - owner: &Keypair, - maximum_pending_balance_credit_counter: Option, - require_memo: bool, - require_fee: bool, - mint_authority: &Keypair, - amount: u64, - decimals: u8, - ) -> Self - where - T: SendTransaction + SimulateTransaction, - { - let meta = Self::new( - token, - owner, - maximum_pending_balance_credit_counter, - require_memo, - require_fee, - ) - .await; - - token - .mint_to( - &meta.token_account, - &mint_authority.pubkey(), - amount, - &[mint_authority], - ) - .await - .unwrap(); - - token - .confidential_transfer_deposit( - &meta.token_account, - &owner.pubkey(), - amount, - decimals, - &[owner], - ) - .await - .unwrap(); - - token - .confidential_transfer_apply_pending_balance( - &meta.token_account, - &owner.pubkey(), - None, - meta.elgamal_keypair.secret(), - &meta.aes_key, - &[owner], - ) - .await - .unwrap(); - meta - } - - #[cfg(feature = "zk-ops")] - async fn check_balances(&self, token: &Token, expected: ConfidentialTokenAccountBalances) - where - T: SendTransaction + SimulateTransaction, - { - let state = token.get_account_info(&self.token_account).await.unwrap(); - let extension = state - .get_extension::() - .unwrap(); - - assert_eq!( - extension - .pending_balance_lo - .decrypt(self.elgamal_keypair.secret()) - .unwrap(), - expected.pending_balance_lo, - ); - assert_eq!( - extension - .pending_balance_hi - .decrypt(self.elgamal_keypair.secret()) - .unwrap(), - expected.pending_balance_hi, - ); - assert_eq!( - extension - .available_balance - .decrypt(self.elgamal_keypair.secret()) - .unwrap(), - expected.available_balance, - ); - assert_eq!( - self.aes_key - .decrypt(&extension.decryptable_available_balance.try_into().unwrap()) - .unwrap(), - expected.decryptable_available_balance, - ); - } -} - -#[cfg(feature = "zk-ops")] -struct ConfidentialTokenAccountBalances { - pending_balance_lo: u64, - pending_balance_hi: u64, - available_balance: u64, - decryptable_available_balance: u64, -} - #[tokio::test] async fn confidential_transfer_configure_token_account() { let authority = Keypair::new(); @@ -568,12 +394,13 @@ async fn confidential_transfer_enable_disable_non_confidential_credits() { .unwrap(); assert!(bool::from(&extension.allow_non_confidential_credits)); + // transfer a different number to change the signature token .transfer( &alice_meta.token_account, &bob_meta.token_account, &alice.pubkey(), - 10, + 9, &[&alice], ) .await diff --git a/token/program-2022-test/tests/program_test.rs b/token/program-2022-test/tests/program_test.rs index 1400285a0c0..ef106aedfb1 100644 --- a/token/program-2022-test/tests/program_test.rs +++ b/token/program-2022-test/tests/program_test.rs @@ -2,10 +2,24 @@ use { solana_program_test::{processor, tokio::sync::Mutex, ProgramTest, ProgramTestContext}, - solana_sdk::signer::{keypair::Keypair, Signer}, - spl_token_2022::{id, native_mint, processor::Processor}, + solana_sdk::{ + pubkey::Pubkey, + signer::{keypair::Keypair, Signer}, + }, + spl_token_2022::{ + extension::{ + confidential_transfer::ConfidentialTransferAccount, BaseStateWithExtensions, + ExtensionType, + }, + id, native_mint, + processor::Processor, + solana_zk_token_sdk::encryption::{auth_encryption::*, elgamal::*}, + }, spl_token_client::{ - client::{ProgramBanksClient, ProgramBanksClientProcessTransaction, ProgramClient}, + client::{ + ProgramBanksClient, ProgramBanksClientProcessTransaction, ProgramClient, + SendTransaction, SimulateTransaction, + }, token::{ExtensionInitializationParams, Token, TokenResult}, }, std::sync::Arc, @@ -160,3 +174,181 @@ impl TestContext { pub(crate) fn keypair_clone(kp: &Keypair) -> Keypair { Keypair::from_bytes(&kp.to_bytes()).expect("failed to copy keypair") } + +pub(crate) struct ConfidentialTokenAccountMeta { + pub(crate) token_account: Pubkey, + pub(crate) elgamal_keypair: ElGamalKeypair, + pub(crate) aes_key: AeKey, +} + +impl ConfidentialTokenAccountMeta { + pub(crate) async fn new( + token: &Token, + owner: &Keypair, + maximum_pending_balance_credit_counter: Option, + require_memo: bool, + require_fee: bool, + ) -> Self + where + T: SendTransaction + SimulateTransaction, + { + let token_account_keypair = Keypair::new(); + + let mut extensions = vec![ExtensionType::ConfidentialTransferAccount]; + if require_memo { + extensions.push(ExtensionType::MemoTransfer); + } + if require_fee { + extensions.push(ExtensionType::ConfidentialTransferFeeAmount); + } + + token + .create_auxiliary_token_account_with_extension_space( + &token_account_keypair, + &owner.pubkey(), + extensions, + ) + .await + .unwrap(); + let token_account = token_account_keypair.pubkey(); + + let elgamal_keypair = + ElGamalKeypair::new_from_signer(owner, &token_account.to_bytes()).unwrap(); + let aes_key = AeKey::new_from_signer(owner, &token_account.to_bytes()).unwrap(); + + token + .confidential_transfer_configure_token_account( + &token_account, + &owner.pubkey(), + None, + maximum_pending_balance_credit_counter, + &elgamal_keypair, + &aes_key, + &[owner], + ) + .await + .unwrap(); + + if require_memo { + token + .enable_required_transfer_memos(&token_account, &owner.pubkey(), &[owner]) + .await + .unwrap(); + } + + Self { + token_account, + elgamal_keypair, + aes_key, + } + } + + #[allow(clippy::too_many_arguments)] + #[cfg(feature = "zk-ops")] + pub(crate) async fn new_with_tokens( + token: &Token, + owner: &Keypair, + maximum_pending_balance_credit_counter: Option, + require_memo: bool, + require_fee: bool, + mint_authority: &Keypair, + amount: u64, + decimals: u8, + ) -> Self + where + T: SendTransaction + SimulateTransaction, + { + let meta = Self::new( + token, + owner, + maximum_pending_balance_credit_counter, + require_memo, + require_fee, + ) + .await; + + token + .mint_to( + &meta.token_account, + &mint_authority.pubkey(), + amount, + &[mint_authority], + ) + .await + .unwrap(); + + token + .confidential_transfer_deposit( + &meta.token_account, + &owner.pubkey(), + amount, + decimals, + &[owner], + ) + .await + .unwrap(); + + token + .confidential_transfer_apply_pending_balance( + &meta.token_account, + &owner.pubkey(), + None, + meta.elgamal_keypair.secret(), + &meta.aes_key, + &[owner], + ) + .await + .unwrap(); + meta + } + + #[cfg(feature = "zk-ops")] + pub(crate) async fn check_balances( + &self, + token: &Token, + expected: ConfidentialTokenAccountBalances, + ) where + T: SendTransaction + SimulateTransaction, + { + let state = token.get_account_info(&self.token_account).await.unwrap(); + let extension = state + .get_extension::() + .unwrap(); + + assert_eq!( + extension + .pending_balance_lo + .decrypt(self.elgamal_keypair.secret()) + .unwrap(), + expected.pending_balance_lo, + ); + assert_eq!( + extension + .pending_balance_hi + .decrypt(self.elgamal_keypair.secret()) + .unwrap(), + expected.pending_balance_hi, + ); + assert_eq!( + extension + .available_balance + .decrypt(self.elgamal_keypair.secret()) + .unwrap(), + expected.available_balance, + ); + assert_eq!( + self.aes_key + .decrypt(&extension.decryptable_available_balance.try_into().unwrap()) + .unwrap(), + expected.decryptable_available_balance, + ); + } +} + +#[cfg(feature = "zk-ops")] +pub(crate) struct ConfidentialTokenAccountBalances { + pub(crate) pending_balance_lo: u64, + pub(crate) pending_balance_hi: u64, + pub(crate) available_balance: u64, + pub(crate) decryptable_available_balance: u64, +}