Add ConfidentialTransferInstruction Deposit/Transfer/Withdraw tests

This commit is contained in:
Michael Vines 2022-02-24 11:08:16 -08:00
parent 3450f173de
commit 22135e547c
4 changed files with 563 additions and 67 deletions

View File

@ -140,7 +140,7 @@ pub type TokenResult<T> = Result<T, TokenError>;
pub struct Token<T, S> {
client: Arc<dyn ProgramClient<T>>,
pubkey: Pubkey,
pubkey: Pubkey, /*token mint*/
payer: S,
program_id: Pubkey,
memo: Arc<RwLock<Option<String>>>,
@ -991,19 +991,19 @@ where
.await
}
/// Prepare a token account for closing
/// Prepare a token account with the confidential transfer extension for closing
pub async fn confidential_transfer_empty_account<S2: Signer>(
&self,
token_account: &Pubkey,
authority: &S2,
elgamal_keypair: &ElGamalKeypair,
) -> TokenResult<T::Output> {
let state = self.get_account_info(&token_account).await.unwrap();
let state = self.get_account_info(token_account).await.unwrap();
let extension =
state.get_extension::<confidential_transfer::ConfidentialTransferAccount>()?;
let proof_data = confidential_transfer::instruction::CloseAccountData::new(
&elgamal_keypair,
elgamal_keypair,
&extension.available_balance.try_into().unwrap(),
)
.map_err(TokenError::Proof)?;
@ -1021,6 +1021,155 @@ where
.await
}
/// Deposit SPL Tokens into the pending balance of a confidential token account
pub async fn confidential_transfer_deposit<S2: Signer>(
&self,
source_token_account: &Pubkey,
destination_token_account: &Pubkey,
source_token_authority: &S2,
amount: u64,
decimals: u8,
) -> TokenResult<T::Output> {
self.process_ixs(
&[confidential_transfer::instruction::deposit(
&self.program_id,
source_token_account,
&self.pubkey,
destination_token_account,
amount,
decimals,
&source_token_authority.pubkey(),
&[],
)?],
&[source_token_authority],
)
.await
}
/// Withdraw SPL Tokens from the available balance of a confidential token account
#[allow(clippy::too_many_arguments)]
pub async fn confidential_transfer_withdraw<S2: Signer>(
&self,
source_token_account: &Pubkey,
destination_token_account: &Pubkey,
source_token_authority: &S2,
amount: u64,
decimals: u8,
source_available_balance: u64,
source_elgamal_keypair: &ElGamalKeypair,
new_source_decryptable_available_balance: AeCiphertext,
) -> TokenResult<T::Output> {
let state = self.get_account_info(source_token_account).await.unwrap();
let extension =
state.get_extension::<confidential_transfer::ConfidentialTransferAccount>()?;
let proof_data = confidential_transfer::instruction::WithdrawData::new(
amount,
source_elgamal_keypair,
source_available_balance,
&extension.available_balance.try_into().unwrap(),
)
.map_err(TokenError::Proof)?;
self.process_ixs(
&confidential_transfer::instruction::withdraw(
&self.program_id,
source_token_account,
destination_token_account,
&self.pubkey,
amount,
decimals,
new_source_decryptable_available_balance,
&source_token_authority.pubkey(),
&[],
&proof_data,
)?,
&[source_token_authority],
)
.await
}
/// Transfer tokens confidentially
#[allow(clippy::too_many_arguments)]
pub async fn confidential_transfer_transfer<S2: Signer>(
&self,
source_token_account: &Pubkey,
destination_token_account: &Pubkey,
source_token_authority: &S2,
amount: u64,
source_available_balance: u64,
source_elgamal_keypair: &ElGamalKeypair,
new_source_decryptable_available_balance: AeCiphertext,
) -> TokenResult<T::Output> {
let source_state = self.get_account_info(source_token_account).await.unwrap();
let source_extension =
source_state.get_extension::<confidential_transfer::ConfidentialTransferAccount>()?;
let destination_state = self
.get_account_info(destination_token_account)
.await
.unwrap();
let destination_extension = destination_state
.get_extension::<confidential_transfer::ConfidentialTransferAccount>(
)?;
let mint_state = self.get_mint_info().await.unwrap();
let ct_mint = mint_state
.get_extension::<confidential_transfer::ConfidentialTransferMint>()
.unwrap();
let proof_data = confidential_transfer::instruction::TransferData::new(
amount,
(
source_available_balance,
&source_extension.available_balance.try_into().unwrap(),
),
source_elgamal_keypair,
(
&destination_extension.pubkey_elgamal.try_into().unwrap(),
&ct_mint.pubkey_auditor.try_into().unwrap(),
),
)
.map_err(TokenError::Proof)?;
self.process_ixs(
&confidential_transfer::instruction::transfer(
&self.program_id,
source_token_account,
destination_token_account,
&self.pubkey,
new_source_decryptable_available_balance,
&source_token_authority.pubkey(),
&[],
&proof_data,
)?,
&[source_token_authority],
)
.await
}
/// Applies the confidential transfer pending balance to the available balance
pub async fn confidential_transfer_apply_pending_balance<S2: Signer>(
&self,
token_account: &Pubkey,
authority: &S2,
expected_pending_balance_credit_counter: u64,
new_decryptable_available_balance: AeCiphertext,
) -> TokenResult<T::Output> {
self.process_ixs(
&[confidential_transfer::instruction::apply_pending_balance(
&self.program_id,
token_account,
expected_pending_balance_credit_counter,
new_decryptable_available_balance,
&authority.pubkey(),
&[],
)?],
&[authority],
)
.await
}
/// Enable confidential transfer `Deposit` and `Transfer` instructions for a token account
pub async fn confidential_transfer_enable_balance_credits<S2: Signer>(
&self,

View File

@ -6,7 +6,7 @@ use {
program_test::{TestContext, TokenContext},
solana_program_test::tokio,
solana_sdk::{
instruction::InstructionError, signature::Signer, signer::keypair::Keypair,
instruction::InstructionError, pubkey::Pubkey, signature::Signer, signer::keypair::Keypair,
transaction::TransactionError, transport::TransportError,
},
spl_token_2022::{
@ -16,9 +16,15 @@ use {
},
ExtensionType,
},
solana_zk_token_sdk::{encryption::elgamal::*, zk_token_elgamal::pod::Zeroable},
solana_zk_token_sdk::{
encryption::{auth_encryption::*, elgamal::*},
zk_token_elgamal::{self, pod::Zeroable},
},
},
spl_token_client::{
client::SendTransaction,
token::{ExtensionInitializationParams, Token, TokenError as TokenClientError},
},
spl_token_client::token::{ExtensionInitializationParams, TokenError as TokenClientError},
std::convert::TryInto,
};
@ -58,6 +64,79 @@ impl ConfidentialTransferMintWithKeypairs {
}
}
struct ConfidentialTokenAccountMeta {
token_account: Pubkey,
elgamal_keypair: ElGamalKeypair,
ae_key: AeKey,
}
impl ConfidentialTokenAccountMeta {
async fn new<T>(token: &Token<T, Keypair>, owner: &Keypair) -> Self
where
T: SendTransaction,
{
let token_account = token
.create_auxiliary_token_account_with_extension_space(
&Keypair::new(),
&owner.pubkey(),
vec![ExtensionType::ConfidentialTransferAccount],
)
.await
.unwrap();
let (elgamal_keypair, ae_key) = token
.confidential_transfer_configure_token_account_and_keypairs(&token_account, owner)
.await
.unwrap();
Self {
token_account,
elgamal_keypair,
ae_key,
}
}
async fn with_tokens<T>(
token: &Token<T, Keypair>,
owner: &Keypair,
mint_authority: &Keypair,
amount: u64,
decimals: u8,
) -> Self
where
T: SendTransaction,
{
let meta = Self::new(token, owner).await;
token
.mint_to(&meta.token_account, mint_authority, amount)
.await
.unwrap();
token
.confidential_transfer_deposit(
&meta.token_account,
&meta.token_account,
owner,
amount,
decimals,
)
.await
.unwrap();
token
.confidential_transfer_apply_pending_balance(
&meta.token_account,
owner,
1,
meta.ae_key.encrypt(amount),
)
.await
.unwrap();
meta
}
}
#[tokio::test]
async fn ct_initialize_and_update_mint() {
let wrong_keypair = Keypair::new();
@ -144,22 +223,12 @@ async fn ct_configure_token_account() {
.unwrap();
let TokenContext { token, alice, .. } = context.token_context.unwrap();
let alice_meta = ConfidentialTokenAccountMeta::new(&token, &alice).await;
let alice_token_account = token
.create_auxiliary_token_account_with_extension_space(
&alice,
&alice.pubkey(),
vec![ExtensionType::ConfidentialTransferAccount],
)
let state = token
.get_account_info(&alice_meta.token_account)
.await
.unwrap();
let (alice_elgamal_keypair, alice_ae_key) = token
.confidential_transfer_configure_token_account_and_keypairs(&alice_token_account, &alice)
.await
.unwrap();
let state = token.get_account_info(&alice_token_account).await.unwrap();
let extension = state
.get_extension::<ConfidentialTransferAccount>()
.unwrap();
@ -167,21 +236,25 @@ async fn ct_configure_token_account() {
assert!(bool::from(&extension.allow_balance_credits));
assert_eq!(
extension.pubkey_elgamal,
alice_elgamal_keypair.public.into()
alice_meta.elgamal_keypair.public.into()
);
assert_eq!(
alice_ae_key
alice_meta
.ae_key
.decrypt(&(extension.decryptable_available_balance.try_into().unwrap()))
.unwrap(),
0
);
token
.confidential_transfer_approve_account(&alice_token_account, &ct_mint_authority)
.confidential_transfer_approve_account(&alice_meta.token_account, &ct_mint_authority)
.await
.unwrap();
let state = token.get_account_info(&alice_token_account).await.unwrap();
let state = token
.get_account_info(&alice_meta.token_account)
.await
.unwrap();
let extension = state
.get_extension::<ConfidentialTransferAccount>()
.unwrap();
@ -201,36 +274,29 @@ async fn ct_enable_disable_balance_credits() {
.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 _ = token
.confidential_transfer_configure_token_account_and_keypairs(&alice_token_account, &alice)
.await
.unwrap();
let alice_meta = ConfidentialTokenAccountMeta::new(&token, &alice).await;
token
.confidential_transfer_disable_balance_credits(&alice_token_account, &alice)
.confidential_transfer_disable_balance_credits(&alice_meta.token_account, &alice)
.await
.unwrap();
let state = token
.get_account_info(&alice_meta.token_account)
.await
.unwrap();
let state = token.get_account_info(&alice_token_account).await.unwrap();
let extension = state
.get_extension::<ConfidentialTransferAccount>()
.unwrap();
assert!(!bool::from(&extension.allow_balance_credits));
token
.confidential_transfer_enable_balance_credits(&alice_token_account, &alice)
.confidential_transfer_enable_balance_credits(&alice_meta.token_account, &alice)
.await
.unwrap();
let state = token
.get_account_info(&alice_meta.token_account)
.await
.unwrap();
let state = token.get_account_info(&alice_token_account).await.unwrap();
let extension = state
.get_extension::<ConfidentialTransferAccount>()
.unwrap();
@ -251,22 +317,305 @@ async fn ct_new_account_is_empty() {
let TokenContext { token, alice, .. } = context.token_context.unwrap();
let alice_token_account = token
.create_auxiliary_token_account_with_extension_space(
let alice_meta = ConfidentialTokenAccountMeta::new(&token, &alice).await;
token
.confidential_transfer_empty_account(
&alice_meta.token_account,
&alice,
&alice.pubkey(),
vec![ExtensionType::ConfidentialTransferAccount],
&alice_meta.elgamal_keypair,
)
.await
.unwrap();
}
#[tokio::test]
async fn ct_deposit() {
let ConfidentialTransferMintWithKeypairs { ct_mint, .. } =
ConfidentialTransferMintWithKeypairs::new();
let mut context = TestContext::new().await;
context
.init_token_with_mint(vec![
ExtensionInitializationParams::ConfidentialTransferMint { ct_mint },
])
.await
.unwrap();
let TokenContext {
token,
alice,
mint_authority,
decimals,
..
} = context.token_context.unwrap();
let alice_meta = ConfidentialTokenAccountMeta::new(&token, &alice).await;
token
.mint_to(&alice_meta.token_account, &mint_authority, 42)
.await
.unwrap();
let state = token
.get_account_info(&alice_meta.token_account)
.await
.unwrap();
assert_eq!(state.base.amount, 42);
let extension = state
.get_extension::<ConfidentialTransferAccount>()
.unwrap();
assert_eq!(extension.pending_balance_credit_counter, 0.into());
assert_eq!(extension.expected_pending_balance_credit_counter, 0.into());
assert_eq!(extension.actual_pending_balance_credit_counter, 0.into());
assert_eq!(
extension.pending_balance,
zk_token_elgamal::pod::ElGamalCiphertext::zeroed()
);
assert_eq!(
extension.available_balance,
zk_token_elgamal::pod::ElGamalCiphertext::zeroed()
);
token
.confidential_transfer_deposit(
&alice_meta.token_account,
&alice_meta.token_account,
&alice,
42,
decimals,
)
.await
.unwrap();
let (alice_elgamal_keypair, _) = token
.confidential_transfer_configure_token_account_and_keypairs(&alice_token_account, &alice)
let state = token
.get_account_info(&alice_meta.token_account)
.await
.unwrap();
assert_eq!(state.base.amount, 0);
let extension = state
.get_extension::<ConfidentialTransferAccount>()
.unwrap();
assert_eq!(extension.pending_balance_credit_counter, 1.into());
assert_eq!(extension.expected_pending_balance_credit_counter, 0.into());
assert_eq!(extension.actual_pending_balance_credit_counter, 0.into());
assert_ne!(
extension.pending_balance,
zk_token_elgamal::pod::ElGamalCiphertext::zeroed()
);
assert_eq!(
extension.available_balance,
zk_token_elgamal::pod::ElGamalCiphertext::zeroed()
);
let new_decryptable_available_balance = alice_meta.ae_key.encrypt(42_u64);
token
.confidential_transfer_apply_pending_balance(
&alice_meta.token_account,
&alice,
1,
new_decryptable_available_balance.clone(),
)
.await
.unwrap();
let state = token
.get_account_info(&alice_meta.token_account)
.await
.unwrap();
let extension = state
.get_extension::<ConfidentialTransferAccount>()
.unwrap();
assert_eq!(
extension.decryptable_available_balance,
new_decryptable_available_balance.into(),
);
assert_eq!(extension.pending_balance_credit_counter, 1.into());
assert_eq!(extension.expected_pending_balance_credit_counter, 1.into());
assert_eq!(extension.actual_pending_balance_credit_counter, 1.into());
assert_eq!(
extension.pending_balance,
zk_token_elgamal::pod::ElGamalCiphertext::zeroed()
);
assert_ne!(
extension.available_balance,
zk_token_elgamal::pod::ElGamalCiphertext::zeroed()
);
}
#[tokio::test]
async fn ct_withdraw() {
let ConfidentialTransferMintWithKeypairs { ct_mint, .. } =
ConfidentialTransferMintWithKeypairs::new();
let mut context = TestContext::new().await;
context
.init_token_with_mint(vec![
ExtensionInitializationParams::ConfidentialTransferMint { ct_mint },
])
.await
.unwrap();
let TokenContext {
token,
alice,
mint_authority,
decimals,
..
} = context.token_context.unwrap();
let alice_meta =
ConfidentialTokenAccountMeta::with_tokens(&token, &alice, &mint_authority, 42, decimals)
.await;
let state = token
.get_account_info(&alice_meta.token_account)
.await
.unwrap();
assert_eq!(state.base.amount, 0);
token
.confidential_transfer_withdraw(
&alice_meta.token_account,
&alice_meta.token_account,
&alice,
21,
decimals,
42,
&alice_meta.elgamal_keypair,
alice_meta.ae_key.encrypt(42_u64 - 21),
)
.await
.unwrap();
let state = token
.get_account_info(&alice_meta.token_account)
.await
.unwrap();
assert_eq!(state.base.amount, 21);
token
.confidential_transfer_withdraw(
&alice_meta.token_account,
&alice_meta.token_account,
&alice,
21,
decimals,
21,
&alice_meta.elgamal_keypair,
alice_meta.ae_key.encrypt(0u64),
)
.await
.unwrap();
let state = token
.get_account_info(&alice_meta.token_account)
.await
.unwrap();
assert_eq!(state.base.amount, 42);
token
.confidential_transfer_empty_account(
&alice_meta.token_account,
&alice,
&alice_meta.elgamal_keypair,
)
.await
.unwrap();
}
#[tokio::test]
async fn ct_transfer() {
let ConfidentialTransferMintWithKeypairs { ct_mint, .. } =
ConfidentialTransferMintWithKeypairs::new();
let mut context = TestContext::new().await;
context
.init_token_with_mint(vec![
ExtensionInitializationParams::ConfidentialTransferMint { ct_mint },
])
.await
.unwrap();
let TokenContext {
token,
alice,
bob,
mint_authority,
decimals,
..
} = context.token_context.unwrap();
let alice_meta =
ConfidentialTokenAccountMeta::with_tokens(&token, &alice, &mint_authority, 42, decimals)
.await;
let bob_meta = ConfidentialTokenAccountMeta::new(&token, &bob).await;
// Self-transfer of 0 tokens
token
.confidential_transfer_transfer(
&alice_meta.token_account,
&alice_meta.token_account,
&alice,
0, // amount
42, // available balance
&alice_meta.elgamal_keypair,
alice_meta.ae_key.encrypt(42_u64),
)
.await
.unwrap();
// Self-transfer of N tokens
token
.confidential_transfer_transfer(
&alice_meta.token_account,
&alice_meta.token_account,
&alice,
42, // amount
42, // available balance
&alice_meta.elgamal_keypair,
alice_meta.ae_key.encrypt(42_u64),
)
.await
.unwrap();
token
.confidential_transfer_empty_account(&alice_token_account, &alice, &alice_elgamal_keypair)
.confidential_transfer_apply_pending_balance(
&alice_meta.token_account,
&alice,
2,
alice_meta.ae_key.encrypt(42_u64),
)
.await
.unwrap();
token
.confidential_transfer_transfer(
&alice_meta.token_account,
&bob_meta.token_account,
&alice,
42, // amount
42, // available balance
&alice_meta.elgamal_keypair,
alice_meta.ae_key.encrypt(0_u64),
)
.await
.unwrap();
token
.confidential_transfer_empty_account(
&alice_meta.token_account,
&alice,
&alice_meta.elgamal_keypair,
)
.await
.unwrap();
let err = token
.confidential_transfer_empty_account(
&bob_meta.token_account,
&bob,
&bob_meta.elgamal_keypair,
)
.await
.unwrap_err();
assert_eq!(
err,
TokenClientError::Client(Box::new(TransportError::TransactionError(
TransactionError::InstructionError(1, InstructionError::InvalidAccountData)
)))
);
}

View File

@ -624,7 +624,7 @@ pub fn deposit(
decimals: u8,
authority: &Pubkey,
multisig_signers: &[&Pubkey],
) -> Result<Vec<Instruction>, ProgramError> {
) -> Result<Instruction, ProgramError> {
check_program_account(token_program_id)?;
let mut accounts = vec![
AccountMeta::new(*source_token_account, false),
@ -637,7 +637,7 @@ pub fn deposit(
accounts.push(AccountMeta::new_readonly(**multisig_signer, true));
}
Ok(vec![encode_instruction(
Ok(encode_instruction(
token_program_id,
accounts,
ConfidentialTransferInstruction::Deposit,
@ -645,7 +645,7 @@ pub fn deposit(
amount: amount.into(),
decimals,
},
)])
))
}
/// Create a inner `Withdraw` instruction
@ -829,17 +829,15 @@ pub fn apply_pending_balance(
new_decryptable_available_balance: AeCiphertext,
authority: &Pubkey,
multisig_signers: &[&Pubkey],
) -> Result<Vec<Instruction>, ProgramError> {
Ok(vec![
inner_apply_pending_balance(
token_program_id,
token_account,
pending_balance_instructions,
new_decryptable_available_balance.into(),
authority,
multisig_signers,
)?, // calls check_program_account
])
) -> Result<Instruction, ProgramError> {
inner_apply_pending_balance(
token_program_id,
token_account,
pending_balance_instructions,
new_decryptable_available_balance.into(),
authority,
multisig_signers,
) // calls check_program_account
}
fn enable_or_disable_balance_credits(

View File

@ -520,11 +520,11 @@ fn process_transfer(
return Err(TokenError::FeeParametersMismatch.into());
}
let ciphertext_lo= EncryptedBalance::from((
let ciphertext_lo = EncryptedBalance::from((
proof_data.ciphertext_lo.commitment,
proof_data.ciphertext_lo.handle_source,
));
let ciphertext_hi= EncryptedBalance::from((
let ciphertext_hi = EncryptedBalance::from((
proof_data.ciphertext_hi.commitment,
proof_data.ciphertext_hi.handle_source,
));
@ -559,11 +559,11 @@ fn process_transfer(
return Err(TokenError::ConfidentialTransferElGamalPubkeyMismatch.into());
}
let ciphertext_lo= EncryptedBalance::from((
let ciphertext_lo = EncryptedBalance::from((
proof_data.ciphertext_lo.commitment,
proof_data.ciphertext_lo.handle_source,
));
let ciphertext_hi= EncryptedBalance::from((
let ciphertext_hi = EncryptedBalance::from((
proof_data.ciphertext_hi.commitment,
proof_data.ciphertext_hi.handle_source,
));