diff --git a/client/tests/rpc_client.rs b/client/tests/rpc_client.rs index 1ea59e50fb..a8d158f78a 100644 --- a/client/tests/rpc_client.rs +++ b/client/tests/rpc_client.rs @@ -141,6 +141,7 @@ async fn test_all_endpoints() { vec![payer_pubkey], None, false, + 0, ); let tx = Transaction::new_signed_with_payer( diff --git a/examples/token-escrow/programs/token-escrow/src/escrow_with_pda/sdk.rs b/examples/token-escrow/programs/token-escrow/src/escrow_with_pda/sdk.rs index d91aab031c..9ce2e7bad2 100644 --- a/examples/token-escrow/programs/token-escrow/src/escrow_with_pda/sdk.rs +++ b/examples/token-escrow/programs/token-escrow/src/escrow_with_pda/sdk.rs @@ -132,6 +132,7 @@ pub fn create_withdrawal_escrow_instruction( None, None, None, + &[], ); let merkle_tree_indices = add_and_get_remaining_account_indices( diff --git a/js/compressed-token/src/idl/light_compressed_token.ts b/js/compressed-token/src/idl/light_compressed_token.ts index 4f8396b8ea..72a1408996 100644 --- a/js/compressed-token/src/idl/light_compressed_token.ts +++ b/js/compressed-token/src/idl/light_compressed_token.ts @@ -45,6 +45,57 @@ export type LightCompressedToken = { ]; args: []; }, + { + name: 'addTokenPool'; + docs: [ + 'This instruction creates an additional token pool for a given mint.', + 'The maximum number of token pools per mint is 5.', + ]; + accounts: [ + { + name: 'feePayer'; + isMut: true; + isSigner: true; + docs: ['UNCHECKED: only pays fees.']; + }, + { + name: 'tokenPoolPda'; + isMut: true; + isSigner: false; + }, + { + name: 'existingTokenPoolPda'; + isMut: false; + isSigner: false; + }, + { + name: 'systemProgram'; + isMut: false; + isSigner: false; + }, + { + name: 'mint'; + isMut: true; + isSigner: false; + }, + { + name: 'tokenProgram'; + isMut: false; + isSigner: false; + }, + { + name: 'cpiAuthorityPda'; + isMut: false; + isSigner: false; + }, + ]; + args: [ + { + name: 'tokenPoolBump'; + type: 'u8'; + }, + ]; + }, { name: 'mintTo'; docs: [ @@ -1597,6 +1648,57 @@ export const IDL: LightCompressedToken = { ], args: [], }, + { + name: 'addTokenPool', + docs: [ + 'This instruction creates an additional token pool for a given mint.', + 'The maximum number of token pools per mint is 5.', + ], + accounts: [ + { + name: 'feePayer', + isMut: true, + isSigner: true, + docs: ['UNCHECKED: only pays fees.'], + }, + { + name: 'tokenPoolPda', + isMut: true, + isSigner: false, + }, + { + name: 'existingTokenPoolPda', + isMut: false, + isSigner: false, + }, + { + name: 'systemProgram', + isMut: false, + isSigner: false, + }, + { + name: 'mint', + isMut: true, + isSigner: false, + }, + { + name: 'tokenProgram', + isMut: false, + isSigner: false, + }, + { + name: 'cpiAuthorityPda', + isMut: false, + isSigner: false, + }, + ], + args: [ + { + name: 'tokenPoolBump', + type: 'u8', + }, + ], + }, { name: 'mintTo', docs: [ diff --git a/js/stateless.js/src/idls/account_compression.ts b/js/stateless.js/src/idls/account_compression.ts index 26d02cbec8..f90f8a4ad2 100644 --- a/js/stateless.js/src/idls/account_compression.ts +++ b/js/stateless.js/src/idls/account_compression.ts @@ -940,6 +940,45 @@ export type AccountCompression = { }, ]; }, + { + name: 'migrateState'; + accounts: [ + { + name: 'authority'; + isMut: false; + isSigner: true; + }, + { + name: 'registeredProgramPda'; + isMut: false; + isSigner: false; + isOptional: true; + }, + { + name: 'logWrapper'; + isMut: false; + isSigner: false; + }, + { + name: 'merkleTree'; + isMut: true; + isSigner: false; + }, + { + name: 'outputQueue'; + isMut: true; + isSigner: false; + }, + ]; + args: [ + { + name: 'input'; + type: { + defined: 'MigrateLeafParams'; + }; + }, + ]; + }, ]; accounts: [ { @@ -1151,6 +1190,39 @@ export type AccountCompression = { ]; }; }, + { + name: 'MigrateLeafParams'; + type: { + kind: 'struct'; + fields: [ + { + name: 'changeLogIndex'; + type: 'u64'; + }, + { + name: 'leaf'; + type: { + array: ['u8', 32]; + }; + }, + { + name: 'leafIndex'; + type: 'u64'; + }, + { + name: 'proof'; + type: { + array: [ + { + array: ['u8', 32]; + }, + 16, + ]; + }; + }, + ]; + }; + }, { name: 'AddressQueueConfig'; type: { @@ -1314,6 +1386,14 @@ export type AccountCompression = { code: 6031; name: 'StateMerkleTreeAccountDiscriminatorMismatch'; }, + { + code: 6032; + name: 'RegistryProgramIsNone'; + }, + { + code: 6033; + name: 'EmptyLeaf'; + }, ]; }; @@ -2259,6 +2339,45 @@ export const IDL: AccountCompression = { }, ], }, + { + name: 'migrateState', + accounts: [ + { + name: 'authority', + isMut: false, + isSigner: true, + }, + { + name: 'registeredProgramPda', + isMut: false, + isSigner: false, + isOptional: true, + }, + { + name: 'logWrapper', + isMut: false, + isSigner: false, + }, + { + name: 'merkleTree', + isMut: true, + isSigner: false, + }, + { + name: 'outputQueue', + isMut: true, + isSigner: false, + }, + ], + args: [ + { + name: 'input', + type: { + defined: 'MigrateLeafParams', + }, + }, + ], + }, ], accounts: [ { @@ -2470,6 +2589,39 @@ export const IDL: AccountCompression = { ], }, }, + { + name: 'MigrateLeafParams', + type: { + kind: 'struct', + fields: [ + { + name: 'changeLogIndex', + type: 'u64', + }, + { + name: 'leaf', + type: { + array: ['u8', 32], + }, + }, + { + name: 'leafIndex', + type: 'u64', + }, + { + name: 'proof', + type: { + array: [ + { + array: ['u8', 32], + }, + 16, + ], + }, + }, + ], + }, + }, { name: 'AddressQueueConfig', type: { @@ -2633,5 +2785,13 @@ export const IDL: AccountCompression = { code: 6031, name: 'StateMerkleTreeAccountDiscriminatorMismatch', }, + { + code: 6032, + name: 'RegistryProgramIsNone', + }, + { + code: 6033, + name: 'EmptyLeaf', + }, ], }; diff --git a/js/stateless.js/src/idls/light_compressed_token.ts b/js/stateless.js/src/idls/light_compressed_token.ts index 4f8396b8ea..72a1408996 100644 --- a/js/stateless.js/src/idls/light_compressed_token.ts +++ b/js/stateless.js/src/idls/light_compressed_token.ts @@ -45,6 +45,57 @@ export type LightCompressedToken = { ]; args: []; }, + { + name: 'addTokenPool'; + docs: [ + 'This instruction creates an additional token pool for a given mint.', + 'The maximum number of token pools per mint is 5.', + ]; + accounts: [ + { + name: 'feePayer'; + isMut: true; + isSigner: true; + docs: ['UNCHECKED: only pays fees.']; + }, + { + name: 'tokenPoolPda'; + isMut: true; + isSigner: false; + }, + { + name: 'existingTokenPoolPda'; + isMut: false; + isSigner: false; + }, + { + name: 'systemProgram'; + isMut: false; + isSigner: false; + }, + { + name: 'mint'; + isMut: true; + isSigner: false; + }, + { + name: 'tokenProgram'; + isMut: false; + isSigner: false; + }, + { + name: 'cpiAuthorityPda'; + isMut: false; + isSigner: false; + }, + ]; + args: [ + { + name: 'tokenPoolBump'; + type: 'u8'; + }, + ]; + }, { name: 'mintTo'; docs: [ @@ -1597,6 +1648,57 @@ export const IDL: LightCompressedToken = { ], args: [], }, + { + name: 'addTokenPool', + docs: [ + 'This instruction creates an additional token pool for a given mint.', + 'The maximum number of token pools per mint is 5.', + ], + accounts: [ + { + name: 'feePayer', + isMut: true, + isSigner: true, + docs: ['UNCHECKED: only pays fees.'], + }, + { + name: 'tokenPoolPda', + isMut: true, + isSigner: false, + }, + { + name: 'existingTokenPoolPda', + isMut: false, + isSigner: false, + }, + { + name: 'systemProgram', + isMut: false, + isSigner: false, + }, + { + name: 'mint', + isMut: true, + isSigner: false, + }, + { + name: 'tokenProgram', + isMut: false, + isSigner: false, + }, + { + name: 'cpiAuthorityPda', + isMut: false, + isSigner: false, + }, + ], + args: [ + { + name: 'tokenPoolBump', + type: 'u8', + }, + ], + }, { name: 'mintTo', docs: [ diff --git a/js/stateless.js/src/idls/light_registry.ts b/js/stateless.js/src/idls/light_registry.ts index 59dbe116d6..05d9e23e78 100644 --- a/js/stateless.js/src/idls/light_registry.ts +++ b/js/stateless.js/src/idls/light_registry.ts @@ -1274,6 +1274,63 @@ export type LightRegistry = { }, ]; }, + { + name: 'migrateState'; + accounts: [ + { + name: 'registeredForesterPda'; + isMut: true; + isSigner: false; + }, + { + name: 'authority'; + isMut: false; + isSigner: true; + }, + { + name: 'cpiAuthority'; + isMut: false; + isSigner: false; + }, + { + name: 'registeredProgramPda'; + isMut: false; + isSigner: false; + }, + { + name: 'accountCompressionProgram'; + isMut: false; + isSigner: false; + }, + { + name: 'logWrapper'; + isMut: false; + isSigner: false; + }, + { + name: 'merkleTree'; + isMut: true; + isSigner: false; + }, + { + name: 'outputQueue'; + isMut: true; + isSigner: false; + }, + ]; + args: [ + { + name: 'bump'; + type: 'u8'; + }, + { + name: 'inputs'; + type: { + defined: 'MigrateLeafParams'; + }; + }, + ]; + }, ]; accounts: [ { @@ -2959,6 +3016,63 @@ export const IDL: LightRegistry = { }, ], }, + { + name: 'migrateState', + accounts: [ + { + name: 'registeredForesterPda', + isMut: true, + isSigner: false, + }, + { + name: 'authority', + isMut: false, + isSigner: true, + }, + { + name: 'cpiAuthority', + isMut: false, + isSigner: false, + }, + { + name: 'registeredProgramPda', + isMut: false, + isSigner: false, + }, + { + name: 'accountCompressionProgram', + isMut: false, + isSigner: false, + }, + { + name: 'logWrapper', + isMut: false, + isSigner: false, + }, + { + name: 'merkleTree', + isMut: true, + isSigner: false, + }, + { + name: 'outputQueue', + isMut: true, + isSigner: false, + }, + ], + args: [ + { + name: 'bump', + type: 'u8', + }, + { + name: 'inputs', + type: { + defined: 'MigrateLeafParams', + }, + }, + ], + }, ], accounts: [ { diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 29342a0b4e..0a42d50d1c 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -398,6 +398,10 @@ importers: specifier: ^1.6.0 version: 1.6.0(@types/node@22.5.5)(@vitest/browser@1.6.0)(terser@5.31.0) + hasher.rs/src/main/wasm: {} + + hasher.rs/src/main/wasm-simd: {} + js/compressed-token: dependencies: '@coral-xyz/anchor': diff --git a/programs/compressed-token/src/burn.rs b/programs/compressed-token/src/burn.rs index 05eb7d392d..2f61c4059a 100644 --- a/programs/compressed-token/src/burn.rs +++ b/programs/compressed-token/src/burn.rs @@ -8,13 +8,14 @@ use light_system_program::{ use light_utils::hash_to_bn254_field_size_be; use crate::{ - constants::NOT_FROZEN, + constants::{NOT_FROZEN, NUM_MAX_POOL_ACCOUNTS}, process_transfer::{ add_token_data_to_input_compressed_accounts, cpi_execute_compressed_transaction_transfer, create_output_compressed_accounts, get_cpi_signer_seeds, get_input_compressed_accounts_with_merkle_context_and_check_signer, DelegatedTransfer, InputTokenDataWithContext, }, + spl_compression::check_spl_token_pool_derivation, BurnInstruction, ErrorCode, }; @@ -68,10 +69,61 @@ pub fn burn_spl_from_pool_pda<'info>( ctx: &Context<'_, '_, '_, 'info, BurnInstruction<'info>>, inputs: &CompressedTokenInstructionDataBurn, ) -> Result<()> { - let pre_token_balance = ctx.accounts.token_pool_pda.amount; + let mut token_pool_bumps = (0..NUM_MAX_POOL_ACCOUNTS).collect::>(); + let mut amount = inputs.burn_amount; + let mut token_pool_pda = &ctx.accounts.token_pool_pda; + let mint_bytes = ctx.accounts.mint.key().to_bytes(); + for i in 0..NUM_MAX_POOL_ACCOUNTS { + if i != 0 { + token_pool_pda = &ctx.remaining_accounts[i as usize - 1]; + } + let token_pool_amount = + TokenAccount::try_deserialize(&mut &token_pool_pda.data.borrow()[..]) + .map_err(|_| crate::ErrorCode::InvalidTokenPoolPda)? + .amount; + let withdrawal_amount = std::cmp::min(amount, token_pool_amount); + if withdrawal_amount == 0 { + continue; + } + let mut remove_index = 0; + for (index, i) in token_pool_bumps.iter().enumerate() { + match check_spl_token_pool_derivation( + mint_bytes.as_slice(), + &token_pool_pda.key(), + &[*i], + ) { + true => { + burn( + ctx, + token_pool_pda.to_account_info(), + withdrawal_amount, + token_pool_amount, + )?; + + remove_index = index; + } + false => {} + } + } + token_pool_bumps.remove(remove_index); + + amount = amount.saturating_sub(withdrawal_amount); + if amount == 0 { + return Ok(()); + } + } + err!(crate::ErrorCode::FailedToDecompress) +} + +fn burn<'info>( + ctx: &Context<'_, '_, '_, 'info, BurnInstruction<'info>>, + token_pool_pda: AccountInfo<'info>, + burn_amount: u64, + pre_token_balance: u64, +) -> Result<()> { let cpi_accounts = anchor_spl::token_interface::Burn { mint: ctx.accounts.mint.to_account_info(), - from: ctx.accounts.token_pool_pda.to_account_info(), + from: token_pool_pda.to_account_info(), authority: ctx.accounts.cpi_authority_pda.to_account_info(), }; let signer_seeds = get_cpi_signer_seeds(); @@ -81,19 +133,15 @@ pub fn burn_spl_from_pool_pda<'info>( cpi_accounts, signer_seeds_ref, ); - anchor_spl::token_interface::burn(cpi_ctx, inputs.burn_amount)?; - - let post_token_balance = TokenAccount::try_deserialize( - &mut &ctx.accounts.token_pool_pda.to_account_info().data.borrow()[..], - )? - .amount; - // Guard against unexpected behavior of the SPL token program. - if post_token_balance != pre_token_balance - inputs.burn_amount { + anchor_spl::token_interface::burn(cpi_ctx, burn_amount)?; + let post_token_balance = + TokenAccount::try_deserialize(&mut &token_pool_pda.data.borrow()[..])?.amount; + if post_token_balance != pre_token_balance - burn_amount { msg!( "post_token_balance {} != pre_token_balance {} - burn_amount {}", post_token_balance, pre_token_balance, - inputs.burn_amount + burn_amount ); return err!(crate::ErrorCode::SplTokenSupplyMismatch); } @@ -181,7 +229,7 @@ pub mod sdk { use solana_sdk::{instruction::Instruction, pubkey::Pubkey}; use crate::{ - get_token_pool_pda, + get_token_pool_pda_with_bump, process_transfer::{ get_cpi_authority_pda, transfer_sdk::{ @@ -207,6 +255,8 @@ pub mod sdk { pub burn_amount: u64, pub signer_is_delegate: bool, pub is_token_22: bool, + pub token_pool_bump: u8, + pub additonal_pool_accounts: Vec, } pub fn create_burn_instruction( @@ -214,7 +264,11 @@ pub mod sdk { ) -> Result { let (remaining_accounts, input_token_data_with_context, _) = create_input_output_and_remaining_accounts( - &[inputs.change_account_merkle_tree], + &[ + inputs.additonal_pool_accounts, + vec![inputs.change_account_merkle_tree], + ] + .concat(), &inputs.input_token_data, &inputs.input_compressed_accounts, &inputs.input_merkle_contexts, @@ -254,7 +308,7 @@ pub mod sdk { } .data(); - let token_pool_pda = get_token_pool_pda(&inputs.mint); + let token_pool_pda = get_token_pool_pda_with_bump(&inputs.mint, inputs.token_pool_bump); let token_program = if inputs.is_token_22 { anchor_spl::token_2022::ID } else { diff --git a/programs/compressed-token/src/constants.rs b/programs/compressed-token/src/constants.rs index 10943b488c..67b9ab70f8 100644 --- a/programs/compressed-token/src/constants.rs +++ b/programs/compressed-token/src/constants.rs @@ -2,3 +2,7 @@ pub const TOKEN_COMPRESSED_ACCOUNT_DISCRIMINATOR: [u8; 8] = [2, 0, 0, 0, 0, 0, 0, 0]; pub const BUMP_CPI_AUTHORITY: u8 = 254; pub const NOT_FROZEN: bool = false; +pub const POOL_SEED: &[u8] = b"pool"; + +/// Maximum number of pool accounts that can be created for each mint. +pub const NUM_MAX_POOL_ACCOUNTS: u8 = 5; diff --git a/programs/compressed-token/src/instructions/burn.rs b/programs/compressed-token/src/instructions/burn.rs index c670567b99..86925cd5cf 100644 --- a/programs/compressed-token/src/instructions/burn.rs +++ b/programs/compressed-token/src/instructions/burn.rs @@ -1,12 +1,12 @@ use account_compression::{program::AccountCompression, utils::constants::CPI_AUTHORITY_PDA_SEED}; use anchor_lang::prelude::*; -use anchor_spl::token_interface::{Mint, TokenAccount, TokenInterface}; +use anchor_spl::token_interface::{Mint, TokenInterface}; use light_system_program::{ program::LightSystemProgram, sdk::accounts::{InvokeAccounts, SignerAccounts}, }; -use crate::{program::LightCompressedToken, POOL_SEED}; +use crate::program::LightCompressedToken; #[derive(Accounts)] pub struct BurnInstruction<'info> { @@ -24,9 +24,9 @@ pub struct BurnInstruction<'info> { /// CHECK: is used to burn tokens. #[account(mut)] pub mint: InterfaceAccount<'info, Mint>, - /// CHECK: (seed constraint) is derived from mint account. - #[account(mut, seeds = [POOL_SEED, mint.key().as_ref()], bump)] - pub token_pool_pda: InterfaceAccount<'info, TokenAccount>, + /// CHECK: with is_valid_token_pool_pda(). + #[account(mut)] + pub token_pool_pda: AccountInfo<'info>, pub token_program: Interface<'info, TokenInterface>, pub light_system_program: Program<'info, LightSystemProgram>, /// CHECK: (account compression program). diff --git a/programs/compressed-token/src/instructions/create_token_pool.rs b/programs/compressed-token/src/instructions/create_token_pool.rs index 6f93c5a73c..caa671ac16 100644 --- a/programs/compressed-token/src/instructions/create_token_pool.rs +++ b/programs/compressed-token/src/instructions/create_token_pool.rs @@ -6,7 +6,10 @@ use spl_token_2022::{ pod::PodMint, }; -pub const POOL_SEED: &[u8] = b"pool"; +use crate::{ + constants::{NUM_MAX_POOL_ACCOUNTS, POOL_SEED}, + spl_compression::check_spl_token_pool_derivation, +}; /// Creates an SPL or token-2022 token pool account, which is owned by the token authority PDA. #[derive(Accounts)] @@ -36,7 +39,16 @@ pub struct CreateTokenPoolInstruction<'info> { } pub fn get_token_pool_pda(mint: &Pubkey) -> Pubkey { - let seeds = &[POOL_SEED, mint.as_ref()]; + get_token_pool_pda_with_bump(mint, 0) +} + +pub fn get_token_pool_pda_with_bump(mint: &Pubkey, token_pool_bump: u8) -> Pubkey { + let seeds = &[POOL_SEED, mint.as_ref(), &[token_pool_bump]]; + let seeds = if token_pool_bump == 0 { + &seeds[..2] + } else { + &seeds[..] + }; let (address, _) = Pubkey::find_program_address(seeds, &crate::ID); address } @@ -62,3 +74,81 @@ pub fn assert_mint_extensions(account_data: &[u8]) -> Result<()> { } Ok(()) } + +/// Creates an SPL or token-2022 token pool account, which is owned by the token authority PDA. +#[derive(Accounts)] +#[instruction(token_pool_bump: u8)] +pub struct AddTokenPoolInstruction<'info> { + /// UNCHECKED: only pays fees. + #[account(mut)] + pub fee_payer: Signer<'info>, + #[account( + init, + seeds = [ + POOL_SEED, &mint.key().to_bytes(), &[token_pool_bump], + ], + bump, + payer = fee_payer, + token::mint = mint, + token::authority = cpi_authority_pda, + )] + pub token_pool_pda: InterfaceAccount<'info, TokenAccount>, + pub existing_token_pool_pda: InterfaceAccount<'info, TokenAccount>, + pub system_program: Program<'info, System>, + /// CHECK: is mint account. + #[account(mut)] + pub mint: InterfaceAccount<'info, Mint>, + pub token_program: Interface<'info, TokenInterface>, + /// CHECK: (seeds anchor constraint). + #[account(seeds = [CPI_AUTHORITY_PDA_SEED], bump)] + pub cpi_authority_pda: AccountInfo<'info>, +} + +/// Checks if the token pool PDA is valid. +/// Iterates over all possible bump seeds to check if the token pool PDA is valid. +#[inline(always)] +pub fn is_valid_token_pool_pda(token_pool_pda: &Pubkey, mint: &Pubkey) -> Result<()> { + let mint_bytes = mint.to_bytes(); + let is_valid_token_pool_pda = (0..NUM_MAX_POOL_ACCOUNTS) + .any(|i| check_spl_token_pool_derivation(mint_bytes.as_slice(), token_pool_pda, &[i])); + if !is_valid_token_pool_pda { + err!(crate::ErrorCode::InvalidTokenPoolPda) + } else { + Ok(()) + } +} + +#[cfg(test)] +mod test { + use super::*; + + /// Test: + /// 1. Functional: test_is_valid_token_pool_pda_valid + /// 2. Failing: test_is_valid_token_pool_pda_invalid_derivation + /// 3. Failing: test_is_valid_token_pool_pda_bump_seed_equal_to_num_max_accounts + /// 4. Failing: test_is_valid_token_pool_pda_bump_seed_larger_than_num_max_accounts + #[test] + fn test_is_valid_token_pool_pda() { + // 1. Functional: test_is_valid_token_pool_pda_valid + let mint = Pubkey::new_unique(); + for i in 0..NUM_MAX_POOL_ACCOUNTS { + let valid_pda = get_token_pool_pda_with_bump(&mint, i); + assert!(is_valid_token_pool_pda(&valid_pda, &mint).is_ok()); + } + + // 2. Failing: test_is_valid_token_pool_pda_invalid_derivation + let mint = Pubkey::new_unique(); + let invalid_pda = Pubkey::new_unique(); + assert!(is_valid_token_pool_pda(&invalid_pda, &mint).is_err()); + + // 3. Failing: test_is_valid_token_pool_pda_bump_seed_equal_to_num_max_accounts + let mint = Pubkey::new_unique(); + let invalid_pda = get_token_pool_pda_with_bump(&mint, NUM_MAX_POOL_ACCOUNTS); + assert!(is_valid_token_pool_pda(&invalid_pda, &mint).is_err()); + + // 4. Failing: test_is_valid_token_pool_pda_bump_seed_larger_than_num_max_accounts + let mint = Pubkey::new_unique(); + let invalid_pda = get_token_pool_pda_with_bump(&mint, NUM_MAX_POOL_ACCOUNTS + 1); + assert!(is_valid_token_pool_pda(&invalid_pda, &mint).is_err()); + } +} diff --git a/programs/compressed-token/src/lib.rs b/programs/compressed-token/src/lib.rs index 417e3d8551..d54250a302 100644 --- a/programs/compressed-token/src/lib.rs +++ b/programs/compressed-token/src/lib.rs @@ -32,7 +32,8 @@ solana_security_txt::security_txt! { pub mod light_compressed_token { use super::*; - use constants::NOT_FROZEN; + use constants::{NOT_FROZEN, NUM_MAX_POOL_ACCOUNTS}; + use spl_compression::spl_token_pool_derivation; /// This instruction creates a token pool for a given mint. Every spl mint /// can have one token pool. When a token is compressed the tokens are @@ -46,6 +47,23 @@ pub mod light_compressed_token { ) } + /// This instruction creates an additional token pool for a given mint. + /// The maximum number of token pools per mint is 5. + pub fn add_token_pool<'info>( + ctx: Context<'_, '_, '_, 'info, AddTokenPoolInstruction<'info>>, + token_pool_bump: u8, + ) -> Result<()> { + if token_pool_bump >= NUM_MAX_POOL_ACCOUNTS { + return err!(ErrorCode::InvalidTokenPoolBump); + } + // Check that token pool account with previous bump already exists. + spl_token_pool_derivation( + &ctx.accounts.mint.key().to_bytes(), + &ctx.accounts.existing_token_pool_pda.key(), + &[token_pool_bump.saturating_sub(1)], + ) + } + /// Mints tokens from an spl token mint to a list of compressed accounts. /// Minted tokens are transferred to a pool account owned by the compressed /// token program. The instruction creates one compressed output account for @@ -209,4 +227,7 @@ pub enum ErrorCode { MintWithInvalidExtension, #[msg("The token account balance is less than the remaining amount.")] InsufficientTokenAccountBalance, + #[msg("Max number of token pools reached.")] + InvalidTokenPoolBump, + FailedToDecompress, } diff --git a/programs/compressed-token/src/process_compress_spl_token_account.rs b/programs/compressed-token/src/process_compress_spl_token_account.rs index 1e5035d4f4..2b74362ba0 100644 --- a/programs/compressed-token/src/process_compress_spl_token_account.rs +++ b/programs/compressed-token/src/process_compress_spl_token_account.rs @@ -48,7 +48,7 @@ pub fn process_compress_spl_token_account<'info>( #[cfg(not(target_os = "solana"))] pub mod sdk { - use crate::get_token_pool_pda; + use crate::get_token_pool_pda_with_bump; use anchor_lang::prelude::AccountMeta; use anchor_lang::InstructionData; use anchor_lang::ToAccountMetas; @@ -67,6 +67,7 @@ pub mod sdk { output_merkle_tree: &Pubkey, token_account: &Pubkey, is_token_22: bool, + token_pool_bump: u8, ) -> Instruction { let instruction_data = crate::instruction::CompressSplTokenAccount { owner: *owner, @@ -74,7 +75,7 @@ pub mod sdk { cpi_context, }; let (cpi_authority_pda, _) = crate::process_transfer::get_cpi_authority_pda(); - let token_pool_pda = get_token_pool_pda(mint); + let token_pool_pda = get_token_pool_pda_with_bump(mint, token_pool_bump); let token_program = if is_token_22 { Some(anchor_spl::token_2022::ID) } else { diff --git a/programs/compressed-token/src/process_mint.rs b/programs/compressed-token/src/process_mint.rs index a3b32c6098..de8a28ad85 100644 --- a/programs/compressed-token/src/process_mint.rs +++ b/programs/compressed-token/src/process_mint.rs @@ -3,7 +3,7 @@ use anchor_lang::prelude::*; use anchor_spl::token_interface::{Mint, TokenAccount, TokenInterface}; use light_system_program::{program::LightSystemProgram, OutputCompressedAccountWithPackedContext}; -use crate::{program::LightCompressedToken, POOL_SEED}; +use crate::{is_valid_token_pool_pda, program::LightCompressedToken}; #[cfg(target_os = "solana")] use { @@ -275,6 +275,7 @@ pub fn serialize_mint_to_cpi_instruction_data( #[inline(never)] pub fn mint_spl_to_pool_pda(ctx: &Context, amounts: &[u64]) -> Result<()> { + is_valid_token_pool_pda(&ctx.accounts.token_pool_pda.key(), &ctx.accounts.mint.key())?; let mut mint_amount: u64 = 0; for amount in amounts.iter() { mint_amount = mint_amount @@ -325,7 +326,8 @@ pub struct MintToInstruction<'info> { @ crate::ErrorCode::InvalidAuthorityMint )] pub mint: InterfaceAccount<'info, Mint>, - #[account(mut, seeds = [POOL_SEED, mint.key().as_ref()], bump)] + /// CHECK: with is_valid_token_pool_pda(). + #[account(mut)] pub token_pool_pda: InterfaceAccount<'info, TokenAccount>, pub token_program: Interface<'info, TokenInterface>, pub light_system_program: Program<'info, LightSystemProgram>, @@ -352,7 +354,9 @@ pub struct MintToInstruction<'info> { #[cfg(not(target_os = "solana"))] pub mod mint_sdk { - use crate::{get_token_pool_pda, process_transfer::get_cpi_authority_pda}; + use crate::{ + get_token_pool_pda, get_token_pool_pda_with_bump, process_transfer::get_cpi_authority_pda, + }; use anchor_lang::{system_program, InstructionData, ToAccountMetas}; use light_system_program::sdk::invoke::get_sol_pool_pda; use solana_sdk::{instruction::Instruction, pubkey::Pubkey}; @@ -386,6 +390,39 @@ pub mod mint_sdk { } } + pub fn create_add_token_pool_instruction( + fee_payer: &Pubkey, + mint: &Pubkey, + token_pool_bump: u8, + is_token_22: bool, + ) -> Instruction { + let token_pool_pda = get_token_pool_pda_with_bump(mint, token_pool_bump); + let existing_token_pool_pda = + get_token_pool_pda_with_bump(mint, token_pool_bump.saturating_sub(1)); + let instruction_data = crate::instruction::AddTokenPool { token_pool_bump }; + + let token_program: Pubkey = if is_token_22 { + anchor_spl::token_2022::ID + } else { + anchor_spl::token::ID + }; + let accounts = crate::accounts::AddTokenPoolInstruction { + fee_payer: *fee_payer, + token_pool_pda, + system_program: system_program::ID, + mint: *mint, + token_program, + cpi_authority_pda: get_cpi_authority_pda().0, + existing_token_pool_pda, + }; + + Instruction { + program_id: crate::ID, + accounts: accounts.to_account_metas(Some(true)), + data: instruction_data.data(), + } + } + #[allow(clippy::too_many_arguments)] pub fn create_mint_to_instruction( fee_payer: &Pubkey, @@ -396,8 +433,9 @@ pub mod mint_sdk { public_keys: Vec, lamports: Option, token_2022: bool, + token_pool_bump: u8, ) -> Instruction { - let token_pool_pda = get_token_pool_pda(mint); + let token_pool_pda = get_token_pool_pda_with_bump(mint, token_pool_bump); let instruction_data = crate::instruction::MintTo { amounts, @@ -448,6 +486,7 @@ pub mod mint_sdk { } #[cfg(test)] + mod test { use super::*; use crate::{ @@ -460,6 +499,7 @@ mod test { sdk::compressed_account::{CompressedAccount, CompressedAccountData}, OutputCompressedAccountWithPackedContext, }; + #[test] fn test_manual_ix_data_serialization_borsh_compat() { let pubkeys = vec![Pubkey::new_unique(), Pubkey::new_unique()]; diff --git a/programs/compressed-token/src/process_transfer.rs b/programs/compressed-token/src/process_transfer.rs index 77c63c12c3..38d6e41703 100644 --- a/programs/compressed-token/src/process_transfer.rs +++ b/programs/compressed-token/src/process_transfer.rs @@ -644,6 +644,7 @@ pub mod transfer_sdk { delegate_change_account_index: Option, lamports_change_account_merkle_tree: Option, is_token_22: bool, + additional_token_pools: &[Pubkey], ) -> Result { let (remaining_accounts, mut inputs_struct) = create_inputs_and_remaining_accounts( input_token_data, @@ -658,6 +659,7 @@ pub mod transfer_sdk { compress_or_decompress_amount, delegate_change_account_index, lamports_change_account_merkle_tree, + additional_token_pools, ); if sort { inputs_struct @@ -760,6 +762,7 @@ pub mod transfer_sdk { compress_or_decompress_amount, delegate_change_account_index, lamports_change_account_merkle_tree, + &[], ); Ok((remaining_accounts, compressed_accounts_ix_data)) } @@ -778,11 +781,13 @@ pub mod transfer_sdk { compress_or_decompress_amount: Option, delegate_change_account_index: Option, lamports_change_account_merkle_tree: Option, + accounts: &[Pubkey], ) -> ( HashMap, CompressedTokenInstructionDataTransfer, ) { let mut additonal_accounts = Vec::new(); + additonal_accounts.extend_from_slice(accounts); if let Some(delegate) = delegate { additonal_accounts.push(delegate); for account in input_token_data.iter() { diff --git a/programs/compressed-token/src/spl_compression.rs b/programs/compressed-token/src/spl_compression.rs index 85d4525c83..377c85db17 100644 --- a/programs/compressed-token/src/spl_compression.rs +++ b/programs/compressed-token/src/spl_compression.rs @@ -1,15 +1,16 @@ #![allow(deprecated)] use anchor_lang::{prelude::*, solana_program::account_info::AccountInfo}; -use anchor_spl::token_interface; +use anchor_spl::{token::TokenAccount, token_interface}; use crate::{ - process_transfer::get_cpi_signer_seeds, CompressedTokenInstructionDataTransfer, - TransferInstruction, POOL_SEED, + constants::{NUM_MAX_POOL_ACCOUNTS, POOL_SEED}, + process_transfer::get_cpi_signer_seeds, + CompressedTokenInstructionDataTransfer, TransferInstruction, }; -pub fn process_compression_or_decompression( +pub fn process_compression_or_decompression<'info>( inputs: &CompressedTokenInstructionDataTransfer, - ctx: &Context, + ctx: &Context<'_, '_, '_, 'info, TransferInstruction<'info>>, ) -> Result<()> { if inputs.is_compress { compress_spl_tokens(inputs, ctx) @@ -19,80 +20,144 @@ pub fn process_compression_or_decompression( } pub fn spl_token_pool_derivation( - mint: &Pubkey, - program_id: &Pubkey, + mint_bytes: &[u8], token_pool_pubkey: &Pubkey, + bump: &[u8], ) -> Result<()> { - let seeds = &[POOL_SEED, &mint.to_bytes()[..]]; - let (pda, _bump_seed) = Pubkey::find_program_address(seeds, program_id); - if pda == *token_pool_pubkey { + if check_spl_token_pool_derivation(mint_bytes, token_pool_pubkey, bump) { Ok(()) } else { err!(crate::ErrorCode::InvalidTokenPoolPda) } } -pub fn decompress_spl_tokens( +pub fn check_spl_token_pool_derivation( + mint_bytes: &[u8], + token_pool_pubkey: &Pubkey, + bump: &[u8], +) -> bool { + let seeds = [POOL_SEED, mint_bytes, bump]; + let seeds = if bump[0] == 0 { + &seeds[..2] + } else { + &seeds[..] + }; + let (pda, _) = Pubkey::find_program_address(seeds, &crate::ID); + pda == *token_pool_pubkey +} + +pub fn decompress_spl_tokens<'info>( inputs: &CompressedTokenInstructionDataTransfer, - ctx: &Context, + ctx: &Context<'_, '_, '_, 'info, TransferInstruction<'info>>, ) -> Result<()> { let recipient = match ctx.accounts.compress_or_decompress_token_account.as_ref() { Some(compression_recipient) => compression_recipient.to_account_info(), None => return err!(crate::ErrorCode::DecompressRecipientUndefinedForDecompress), }; - let token_pool_pda = match ctx.accounts.token_pool_pda.as_ref() { + let mut token_pool_pda = match ctx.accounts.token_pool_pda.as_ref() { Some(token_pool_pda) => token_pool_pda.to_account_info(), None => return err!(crate::ErrorCode::CompressedPdaUndefinedForDecompress), }; - spl_token_pool_derivation(&inputs.mint, &crate::ID, &token_pool_pda.key())?; - - let amount = match inputs.compress_or_decompress_amount { + let mut amount = match inputs.compress_or_decompress_amount { Some(amount) => amount, None => return err!(crate::ErrorCode::DeCompressAmountUndefinedForDecompress), }; + let mint_bytes = inputs.mint.to_bytes(); - transfer( - token_pool_pda, - recipient, - ctx.accounts.cpi_authority_pda.to_account_info(), - ctx.accounts - .token_program - .as_ref() - .unwrap() - .to_account_info(), - amount, - ) + let mut token_pool_bumps = (0..NUM_MAX_POOL_ACCOUNTS).collect::>(); + + // + for i in 0..NUM_MAX_POOL_ACCOUNTS { + if i != 0 { + token_pool_pda = ctx.remaining_accounts[i as usize - 1].to_account_info(); + } + let token_pool_amount = + TokenAccount::try_deserialize(&mut &token_pool_pda.data.borrow()[..]) + .map_err(|_| crate::ErrorCode::InvalidTokenPoolPda)? + .amount; + let withdrawal_amount = std::cmp::min(amount, token_pool_amount); + if withdrawal_amount == 0 { + continue; + } + let mut remove_index = 0; + for (index, i) in token_pool_bumps.iter().enumerate() { + match check_spl_token_pool_derivation( + mint_bytes.as_slice(), + &token_pool_pda.key(), + &[*i], + ) { + true => { + transfer( + token_pool_pda.to_account_info(), + recipient.to_account_info(), + ctx.accounts.cpi_authority_pda.to_account_info(), + ctx.accounts + .token_program + .as_ref() + .unwrap() + .to_account_info(), + withdrawal_amount, + )?; + remove_index = index; + } + false => {} + } + } + token_pool_bumps.remove(remove_index); + + amount = amount.saturating_sub(withdrawal_amount); + if amount == 0 { + return Ok(()); + } + } + msg!("Remaining amount: {}.", amount); + msg!("Token pool account balance insufficient for decompression. \nTry to pass more token pool accounts."); + err!(crate::ErrorCode::FailedToDecompress) } -pub fn compress_spl_tokens( +pub fn compress_spl_tokens<'info>( inputs: &CompressedTokenInstructionDataTransfer, - ctx: &Context, + ctx: &Context<'_, '_, '_, 'info, TransferInstruction<'info>>, ) -> Result<()> { let recipient_token_pool = match ctx.accounts.token_pool_pda.as_ref() { - Some(token_pool_pda) => token_pool_pda, + Some(token_pool_pda) => token_pool_pda.to_account_info(), None => return err!(crate::ErrorCode::CompressedPdaUndefinedForCompress), }; - spl_token_pool_derivation(&inputs.mint, &crate::ID, &recipient_token_pool.key())?; let amount = match inputs.compress_or_decompress_amount { Some(amount) => amount, None => return err!(crate::ErrorCode::DeCompressAmountUndefinedForCompress), }; - transfer_compress( - ctx.accounts - .compress_or_decompress_token_account - .as_ref() - .unwrap() - .to_account_info(), - recipient_token_pool.to_account_info(), - ctx.accounts.authority.to_account_info(), - ctx.accounts - .token_program - .as_ref() - .unwrap() - .to_account_info(), - amount, - ) + let mint_bytes = inputs.mint.to_bytes(); + + for i in 0..NUM_MAX_POOL_ACCOUNTS { + match check_spl_token_pool_derivation( + mint_bytes.as_slice(), + &recipient_token_pool.key(), + &[i], + ) { + true => { + transfer_compress( + ctx.accounts + .compress_or_decompress_token_account + .as_ref() + .unwrap() + .to_account_info(), + recipient_token_pool.to_account_info(), + ctx.accounts.authority.to_account_info(), + ctx.accounts + .token_program + .as_ref() + .unwrap() + .to_account_info(), + amount, + )?; + return Ok(()); + } + false => {} + } + } + err!(crate::ErrorCode::InvalidTokenPoolPda) } pub fn transfer<'info>( diff --git a/test-programs/compressed-token-test/tests/test.rs b/test-programs/compressed-token-test/tests/test.rs index ae6dad97ea..620228a7dd 100644 --- a/test-programs/compressed-token-test/tests/test.rs +++ b/test-programs/compressed-token-test/tests/test.rs @@ -1,6 +1,7 @@ #![cfg(feature = "test-sbf")] use account_compression::errors::AccountCompressionErrorCode; +use anchor_lang::prelude::AccountMeta; use anchor_lang::{ system_program, AnchorDeserialize, AnchorSerialize, InstructionData, ToAccountMetas, }; @@ -12,13 +13,14 @@ use light_compressed_token::delegation::sdk::{ CreateRevokeInstructionInputs, }; use light_compressed_token::freeze::sdk::{create_instruction, CreateInstructionInputs}; -use light_compressed_token::get_token_pool_pda; +use light_compressed_token::get_token_pool_pda_with_bump; use light_compressed_token::mint_sdk::create_create_token_pool_instruction; use light_compressed_token::mint_sdk::create_mint_to_instruction; use light_compressed_token::process_transfer::transfer_sdk::create_transfer_instruction; use light_compressed_token::process_transfer::{get_cpi_authority_pda, TokenTransferOutputData}; use light_compressed_token::spl_compression::spl_token_pool_derivation; use light_compressed_token::token_data::AccountState; +use light_compressed_token::{constants::NUM_MAX_POOL_ACCOUNTS, get_token_pool_pda}; use light_compressed_token::{token_data::TokenData, ErrorCode}; use light_program_test::test_env::setup_test_programs_with_accounts; use light_program_test::test_rpc::ProgramTestRpcConnection; @@ -27,8 +29,6 @@ use light_system_program::{ invoke::processor::CompressedProof, sdk::compressed_account::{CompressedAccountWithMerkleContext, MerkleContext}, }; -use light_test_utils::spl::mint_tokens_helper_with_lamports; -use light_test_utils::spl::revoke_test; use light_test_utils::spl::thaw_test; use light_test_utils::spl::BurnInstructionMode; use light_test_utils::spl::{approve_test, create_mint_22_helper}; @@ -40,16 +40,20 @@ use light_test_utils::spl::{ use light_test_utils::spl::{ compressed_transfer_22_test, create_burn_test_instruction, perform_compress_spl_token_account, }; +use light_test_utils::spl::{create_additional_token_pools, mint_tokens_helper_with_lamports}; use light_test_utils::spl::{create_token_2022_account, freeze_test}; use light_test_utils::spl::{mint_spl_tokens, mint_wrapped_sol}; +use light_test_utils::spl::{mint_tokens_22_helper_with_lamports_and_bump, revoke_test}; use light_test_utils::{ airdrop_lamports, assert_rpc_error, create_account_instruction, Indexer, RpcConnection, RpcError, TokenDataWithContext, }; use light_test_utils::{assert_custom_error_or_program_error, indexer::TestIndexer}; use light_verifier::VerifierError; -use rand::Rng; +use rand::seq::SliceRandom; +use rand::{thread_rng, Rng}; use serial_test::serial; +use solana_sdk::signature::Signature; use solana_sdk::system_instruction; use solana_sdk::{ instruction::{Instruction, InstructionError}, @@ -65,7 +69,14 @@ use spl_token::{error::TokenError, instruction::initialize_mint}; async fn test_create_mint() { let (mut rpc, _) = setup_test_programs_with_accounts(None).await; let payer = rpc.get_payer().insecure_clone(); - create_mint_helper(&mut rpc, &payer).await; + let mint = create_mint_helper(&mut rpc, &payer).await; + create_additional_token_pools(&mut rpc, &payer, &mint, false, NUM_MAX_POOL_ACCOUNTS) + .await + .unwrap(); + let mint_22 = create_mint_22_helper(&mut rpc, &payer).await; + create_additional_token_pools(&mut rpc, &payer, &mint_22, true, NUM_MAX_POOL_ACCOUNTS) + .await + .unwrap(); } #[serial] @@ -313,16 +324,257 @@ async fn test_failing_create_token_pool() { let token_pool_pubkey = get_token_pool_pda(&mint.pubkey()); let token_pool_account = rpc.get_account(token_pool_pubkey).await.unwrap().unwrap(); - spl_token_pool_derivation( - &mint.pubkey(), - &light_compressed_token::ID, - &token_pool_pubkey, - ) - .unwrap(); + spl_token_pool_derivation(&mint.pubkey().to_bytes(), &token_pool_pubkey, &[0]).unwrap(); assert_eq!(token_pool_account.data.len(), TokenAccount::LEN); } } +#[serial] +#[tokio::test] +async fn failing_tests_add_token_pool() { + for is_token_22 in vec![false, true] { + let (mut rpc, _) = setup_test_programs_with_accounts(None).await; + let payer = rpc.get_payer().insecure_clone(); + + let mint = if !is_token_22 { + create_mint_helper(&mut rpc, &payer).await + } else { + create_mint_22_helper(&mut rpc, &payer).await + }; + let invalid_mint = if !is_token_22 { + create_mint_helper(&mut rpc, &payer).await + } else { + create_mint_22_helper(&mut rpc, &payer).await + }; + let mut current_token_pool_bump = 1; + create_additional_token_pools(&mut rpc, &payer, &mint, is_token_22, 2) + .await + .unwrap(); + create_additional_token_pools(&mut rpc, &payer, &invalid_mint, is_token_22, 2) + .await + .unwrap(); + current_token_pool_bump += 2; + // 1. failing invalid existing token pool pda + { + let result = add_token_pool( + &mut rpc, + &payer, + &mint, + None, + current_token_pool_bump, + is_token_22, + FailingTestsAddTokenPool::InvalidExistingTokenPoolPda, + ) + .await; + assert_rpc_error(result, 0, ErrorCode::InvalidTokenPoolPda.into()).unwrap(); + } + // 2. failing InvalidTokenPoolPda + { + let result = add_token_pool( + &mut rpc, + &payer, + &mint, + None, + current_token_pool_bump, + is_token_22, + FailingTestsAddTokenPool::InvalidTokenPoolPda, + ) + .await; + assert_rpc_error( + result, + 0, + anchor_lang::error::ErrorCode::ConstraintSeeds.into(), + ) + .unwrap(); + } + // 3. failing invalid system program id + { + let result = add_token_pool( + &mut rpc, + &payer, + &mint, + None, + current_token_pool_bump, + is_token_22, + FailingTestsAddTokenPool::InvalidSystemProgramId, + ) + .await; + assert_rpc_error( + result, + 0, + anchor_lang::error::ErrorCode::InvalidProgramId.into(), + ) + .unwrap(); + } + // 4. failing invalid mint + { + let result = add_token_pool( + &mut rpc, + &payer, + &mint, + None, + current_token_pool_bump, + is_token_22, + FailingTestsAddTokenPool::InvalidMint, + ) + .await; + assert_rpc_error( + result, + 0, + anchor_lang::error::ErrorCode::AccountNotInitialized.into(), + ) + .unwrap(); + } + // 5. failing inconsistent mints + { + let result = add_token_pool( + &mut rpc, + &payer, + &mint, + Some(invalid_mint), + current_token_pool_bump, + is_token_22, + FailingTestsAddTokenPool::InconsistentMints, + ) + .await; + assert_rpc_error(result, 0, ErrorCode::InvalidTokenPoolPda.into()).unwrap(); + } + // 6. failing invalid program id + { + let result = add_token_pool( + &mut rpc, + &payer, + &mint, + None, + current_token_pool_bump, + is_token_22, + FailingTestsAddTokenPool::InvalidTokenProgramId, + ) + .await; + assert_rpc_error( + result, + 0, + anchor_lang::error::ErrorCode::InvalidProgramId.into(), + ) + .unwrap(); + } + // 7. failing invalid cpi authority pda + { + let result = add_token_pool( + &mut rpc, + &payer, + &mint, + None, + current_token_pool_bump, + is_token_22, + FailingTestsAddTokenPool::InvalidCpiAuthorityPda, + ) + .await; + assert_rpc_error( + result, + 0, + anchor_lang::error::ErrorCode::ConstraintSeeds.into(), + ) + .unwrap(); + } + // create all remaining token pools + create_additional_token_pools(&mut rpc, &payer, &mint, is_token_22, 5) + .await + .unwrap(); + // 8. failing invalid token pool bump (too large) + { + let result = add_token_pool( + &mut rpc, + &payer, + &mint, + None, + NUM_MAX_POOL_ACCOUNTS, + is_token_22, + FailingTestsAddTokenPool::Functional, + ) + .await; + assert_rpc_error(result, 0, ErrorCode::InvalidTokenPoolBump.into()).unwrap(); + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum FailingTestsAddTokenPool { + Functional, + InvalidMint, + InconsistentMints, + InvalidTokenPoolPda, + InvalidSystemProgramId, + InvalidExistingTokenPoolPda, + InvalidCpiAuthorityPda, + InvalidTokenProgramId, +} + +pub async fn add_token_pool( + rpc: &mut R, + fee_payer: &Keypair, + mint: &Pubkey, + invalid_mint: Option, + token_pool_bump: u8, + is_token_22: bool, + mode: FailingTestsAddTokenPool, +) -> Result { + let token_pool_pda = if mode == FailingTestsAddTokenPool::InvalidTokenPoolPda { + Pubkey::new_unique() + } else { + get_token_pool_pda_with_bump(mint, token_pool_bump) + }; + let existing_token_pool_pda = if mode == FailingTestsAddTokenPool::InvalidExistingTokenPoolPda { + get_token_pool_pda_with_bump(mint, token_pool_bump.saturating_sub(2)) + } else if let Some(invalid_mint) = invalid_mint { + get_token_pool_pda_with_bump(&invalid_mint, token_pool_bump.saturating_sub(1)) + } else { + get_token_pool_pda_with_bump(mint, token_pool_bump.saturating_sub(1)) + }; + let instruction_data = light_compressed_token::instruction::AddTokenPool { token_pool_bump }; + + let token_program: Pubkey = if mode == FailingTestsAddTokenPool::InvalidTokenProgramId { + Pubkey::new_unique() + } else if is_token_22 { + anchor_spl::token_2022::ID + } else { + anchor_spl::token::ID + }; + let cpi_authority_pda = if mode == FailingTestsAddTokenPool::InvalidCpiAuthorityPda { + Pubkey::new_unique() + } else { + get_cpi_authority_pda().0 + }; + let system_program = if mode == FailingTestsAddTokenPool::InvalidSystemProgramId { + Pubkey::new_unique() + } else { + system_program::ID + }; + let mint = if mode == FailingTestsAddTokenPool::InvalidMint { + Pubkey::new_unique() + } else { + *mint + }; + + let accounts = light_compressed_token::accounts::AddTokenPoolInstruction { + fee_payer: fee_payer.pubkey(), + token_pool_pda, + system_program, + mint, + token_program, + cpi_authority_pda, + existing_token_pool_pda, + }; + + let instruction = Instruction { + program_id: light_compressed_token::ID, + accounts: accounts.to_account_metas(Some(true)), + data: instruction_data.data(), + }; + rpc.create_and_send_transaction(&[instruction], &fee_payer.pubkey(), &[fee_payer]) + .await +} + #[serial] #[tokio::test] async fn test_wrapped_sol() { @@ -393,6 +645,8 @@ async fn test_wrapped_sol() { &token_account_keypair.pubkey(), None, is_token_22, + 0, + None, ) .await; let input_compressed_accounts = @@ -407,6 +661,8 @@ async fn test_wrapped_sol() { &token_account_keypair.pubkey(), None, is_token_22, + 0, + None, ) .await; } @@ -502,6 +758,7 @@ async fn compress_spl_account() { &merkle_tree_pubkey, None, is_token_22, + 0, ) .await .unwrap(); @@ -528,6 +785,7 @@ async fn compress_spl_account() { &merkle_tree_pubkey, Some(first_token_account_balance + 1), // invalid remaining amount is_token_22, + 0, ) .await; assert_rpc_error(result, 0, ErrorCode::InsufficientTokenAccountBalance.into()).unwrap(); @@ -542,6 +800,7 @@ async fn compress_spl_account() { &merkle_tree_pubkey, Some(1), is_token_22, + 0, ) .await .unwrap(); @@ -705,6 +964,7 @@ async fn test_mint_to_failing() { recipients.clone(), None, is_token_22, + 0, ); let result = rpc .create_and_send_transaction(&[instruction], &payer_2.pubkey(), &[&payer_2]) @@ -724,6 +984,7 @@ async fn test_mint_to_failing() { recipients.clone(), None, is_token_22, + 0, ); let result = rpc .create_and_send_transaction(&[instruction], &payer_1.pubkey(), &[&payer_1]) @@ -991,6 +1252,7 @@ async fn test_mint_to_failing() { recipients.clone(), None, is_token_22, + 0, ); let result = rpc .create_and_send_transaction(&[instruction], &payer_1.pubkey(), &[&payer_1]) @@ -1018,6 +1280,7 @@ async fn test_mint_to_failing() { recipients.clone(), None, is_token_22, + 0, ); let result = rpc .create_and_send_transaction(&[instruction], &payer_1.pubkey(), &[&payer_1]) @@ -1036,6 +1299,7 @@ async fn test_mint_to_failing() { recipients.clone(), None, is_token_22, + 0, ); // The first mint is still below `u64::MAX`. rpc.create_and_send_transaction(&[instruction.clone()], &payer_1.pubkey(), &[&payer_1]) @@ -1271,6 +1535,8 @@ async fn test_decompression() { &token_account_keypair.pubkey(), None, is_token_22, + 0, + None, ) .await; println!("5"); @@ -1284,12 +1550,314 @@ async fn test_decompression() { &token_account_keypair.pubkey(), None, is_token_22, + 0, + None, ) .await; } kill_prover(); } +pub async fn mint_tokens_to_all_token_pools( + rpc: &mut R, + test_indexer: &mut TestIndexer, + merkle_tree_pubkey: &Pubkey, + mint_authority: &Keypair, + mint: &Pubkey, + amounts: Vec, + recipients: Vec, + is_token22: bool, + invert_order: bool, +) -> Result<(), RpcError> { + let iterator = (0..NUM_MAX_POOL_ACCOUNTS).collect::>(); + let iterator = if invert_order { + iterator.iter().rev().cloned().collect::>() + } else { + iterator + }; + for bump in iterator { + let token_pool_pda = get_token_pool_pda_with_bump(mint, bump); + let token_pool_account = rpc.get_account(token_pool_pda).await?; + if token_pool_account.is_some() { + mint_tokens_22_helper_with_lamports_and_bump( + rpc, + test_indexer, + merkle_tree_pubkey, + mint_authority, + mint, + amounts.clone(), + recipients.clone(), + None, + is_token22, + bump, + ) + .await; + } + } + Ok(()) +} +use anchor_lang::AccountDeserialize; +/// Assert that every token pool account contains `amount` tokens. +pub async fn assert_minted_to_all_token_pools( + rpc: &mut R, + amount: u64, + mint: &Pubkey, +) -> Result<(), RpcError> { + for bump in 0..NUM_MAX_POOL_ACCOUNTS { + let token_pool_pda = get_token_pool_pda_with_bump(&mint, bump); + let mut token_pool_account = rpc.get_account(token_pool_pda).await?.unwrap(); + let token_pool_data = + TokenAccount::try_deserialize_unchecked(&mut &*token_pool_account.data.as_mut_slice()) + .unwrap(); + assert_eq!(token_pool_data.amount, amount); + } + + Ok(()) +} + +#[serial] +#[tokio::test] +async fn test_mint_to_and_burn_from_all_token_pools() { + spawn_prover( + false, + ProverConfig { + run_mode: None, + circuits: vec![ProofType::Inclusion], + }, + ) + .await; + for is_token_22 in vec![false, true] { + let (mut rpc, env) = setup_test_programs_with_accounts(None).await; + let payer = rpc.get_payer().insecure_clone(); + let merkle_tree_pubkey = env.merkle_tree_pubkey; + let mut test_indexer = + TestIndexer::::init_from_env(&payer, &env, None).await; + let mint = if is_token_22 { + create_mint_22_helper(&mut rpc, &payer).await + } else { + create_mint_helper(&mut rpc, &payer).await + }; + create_additional_token_pools(&mut rpc, &payer, &mint, is_token_22, NUM_MAX_POOL_ACCOUNTS) + .await + .unwrap(); + let amount = 123; + mint_tokens_to_all_token_pools( + &mut rpc, + &mut test_indexer, + &merkle_tree_pubkey, + &payer, + &mint, + vec![amount], + vec![payer.pubkey()], + is_token_22, + is_token_22, // invert order + ) + .await + .unwrap(); + assert_minted_to_all_token_pools(&mut rpc, amount, &mint) + .await + .unwrap(); + let iterator = (0..NUM_MAX_POOL_ACCOUNTS).collect::>(); + let iterator = if !is_token_22 { + iterator.iter().rev().cloned().collect::>() + } else { + iterator + }; + for i in iterator { + let input_compressed_account = + test_indexer.get_compressed_token_accounts_by_owner(&payer.pubkey())[0].clone(); + let change_account_merkle_tree = input_compressed_account + .compressed_account + .merkle_context + .merkle_tree_pubkey; + burn_test( + &payer, + &mut rpc, + &mut test_indexer, + vec![input_compressed_account], + &change_account_merkle_tree, + amount, + false, + None, + is_token_22, + i, + ) + .await; + } + assert_minted_to_all_token_pools(&mut rpc, 0, &mint) + .await + .unwrap(); + } +} + +#[serial] +#[tokio::test] +async fn test_multiple_decompression() { + spawn_prover( + true, + ProverConfig { + run_mode: None, + circuits: vec![ProofType::Inclusion], + }, + ) + .await; + let rng = &mut rand::thread_rng(); + for is_token_22 in vec![false, true] { + println!("is_token_22: {}", is_token_22); + let (mut context, env) = setup_test_programs_with_accounts(None).await; + let payer = context.get_payer().insecure_clone(); + let merkle_tree_pubkey = env.merkle_tree_pubkey; + let mut test_indexer = + TestIndexer::::init_from_env(&payer, &env, None).await; + let sender = Keypair::new(); + airdrop_lamports(&mut context, &sender.pubkey(), 1_000_000_000) + .await + .unwrap(); + let mint = if is_token_22 { + create_mint_22_helper(&mut context, &payer).await + } else { + create_mint_helper(&mut context, &payer).await + }; + let amount = 10000u64; + create_additional_token_pools( + &mut context, + &payer, + &mint, + is_token_22, + NUM_MAX_POOL_ACCOUNTS, + ) + .await + .unwrap(); + + mint_tokens_to_all_token_pools( + &mut context, + &mut test_indexer, + &merkle_tree_pubkey, + &payer, + &mint, + vec![amount], + vec![sender.pubkey()], + is_token_22, + is_token_22, + ) + .await + .unwrap(); + println!("3"); + let token_account_keypair = Keypair::new(); + create_token_2022_account( + &mut context, + &mint, + &token_account_keypair, + &sender, + is_token_22, + ) + .await + .unwrap(); + println!("4"); + + // 1. functional - decompress from any token pool + let mut iterator = vec![0, 1, 2, 3, 4]; + iterator.shuffle(rng); + for i in iterator { + let input_compressed_account = test_indexer + .get_compressed_token_accounts_by_owner(&sender.pubkey()) + .iter() + .filter(|x| x.token_data.amount != 0) + .collect::>()[0] + .clone(); + println!("i = {}", i); + println!("input_compressed_account = {:?}", input_compressed_account); + decompress_test( + &sender, + &mut context, + &mut test_indexer, + vec![input_compressed_account], + amount, + &merkle_tree_pubkey, + &token_account_keypair.pubkey(), + None, + is_token_22, + i, + None, + ) + .await; + } + + println!("5"); + + // 2. functional - compress to any token pool + let mut iterator = vec![0, 1, 2, 3, 4]; + iterator.shuffle(rng); + for i in iterator { + compress_test( + &sender, + &mut context, + &mut test_indexer, + amount, + &mint, + &merkle_tree_pubkey, + &token_account_keypair.pubkey(), + None, + is_token_22, + i, + None, + ) + .await; + } + + // Decompress from all token pools + { + let input_compressed_accounts = test_indexer + .get_compressed_token_accounts_by_owner(&sender.pubkey())[0..4] + .to_vec(); + let amount = input_compressed_accounts + .iter() + .map(|x| x.token_data.amount) + .sum(); + let mut add_token_pool_accounts = (0..4) + .map(|x| get_token_pool_pda_with_bump(&mint, x.clone())) + .collect::>(); + add_token_pool_accounts.shuffle(rng); + decompress_test( + &sender, + &mut context, + &mut test_indexer, + input_compressed_accounts, + amount, + &merkle_tree_pubkey, + &token_account_keypair.pubkey(), + None, + is_token_22, + 4, + Some(add_token_pool_accounts.clone()), + ) + .await; + let input_compressed_accounts = test_indexer + .get_compressed_token_accounts_by_owner(&sender.pubkey()) + .iter() + .filter(|x| x.token_data.amount != 0) + .collect::>()[0] + .clone(); + let amount = input_compressed_accounts.token_data.amount; + decompress_test( + &sender, + &mut context, + &mut test_indexer, + vec![input_compressed_accounts], + amount, + &merkle_tree_pubkey, + &token_account_keypair.pubkey(), + None, + is_token_22, + 4, + Some(add_token_pool_accounts), + ) + .await; + } + } + kill_prover(); +} + /// Test delegation: /// 1. Delegate tokens with approve /// 2. Delegate transfers a part of the delegated tokens @@ -2267,7 +2835,7 @@ async fn test_revoke_failing() { #[tokio::test] async fn test_burn() { spawn_prover( - true, + false, ProverConfig { run_mode: None, circuits: vec![ProofType::Inclusion], @@ -2295,6 +2863,9 @@ async fn test_burn() { create_mint_helper(&mut rpc, &payer).await }; let amount = 10000u64; + create_additional_token_pools(&mut rpc, &payer, &mint, is_token_22, NUM_MAX_POOL_ACCOUNTS) + .await + .unwrap(); mint_tokens_22_helper_with_lamports( &mut rpc, &mut test_indexer, @@ -2326,6 +2897,7 @@ async fn test_burn() { false, None, is_token_22, + 0, ) .await; } @@ -2365,49 +2937,149 @@ async fn test_burn() { let change_account_merkle_tree = input_compressed_accounts[0] .compressed_account .merkle_context - .merkle_tree_pubkey; - burn_test( - &delegate, + .merkle_tree_pubkey; + burn_test( + &delegate, + &mut rpc, + &mut test_indexer, + input_compressed_accounts, + &change_account_merkle_tree, + burn_amount, + true, + None, + is_token_22, + 0, + ) + .await; + } + // 3. Burn all delegated tokens + { + let input_compressed_accounts = + test_indexer.get_compressed_token_accounts_by_owner(&sender.pubkey()); + let input_compressed_accounts = input_compressed_accounts + .iter() + .filter(|x| x.token_data.delegate.is_some()) + .cloned() + .collect::>(); + let burn_amount = input_compressed_accounts + .iter() + .map(|x| x.token_data.amount) + .sum::(); + let change_account_merkle_tree = input_compressed_accounts[0] + .compressed_account + .merkle_context + .merkle_tree_pubkey; + burn_test( + &delegate, + &mut rpc, + &mut test_indexer, + input_compressed_accounts, + &change_account_merkle_tree, + burn_amount, + true, + None, + is_token_22, + 0, + ) + .await; + } + // 5. Burn tokens from multiple token pools + { + let amount = 123; + mint_tokens_to_all_token_pools( + &mut rpc, + &mut test_indexer, + &env.merkle_tree_pubkey, + &payer, + &mint, + vec![amount], + vec![sender.pubkey()], + is_token_22, + false, + ) + .await + .unwrap(); + let input_compressed_accounts = test_indexer + .get_compressed_token_accounts_by_owner(&sender.pubkey()) + .iter() + .filter(|x| x.token_data.amount != 0) + .cloned() + .collect::>()[0..4] + .to_vec(); + let burn_amount = input_compressed_accounts + .iter() + .map(|x| x.token_data.amount) + .sum(); + let invalid_change_account_merkle_tree = input_compressed_accounts[0] + .compressed_account + .merkle_context + .nullifier_queue_pubkey; + let mut additional_token_pool_accounts = (0..4) + .map(|x| get_token_pool_pda_with_bump(&mint, x)) + .collect::>(); + let rng = &mut thread_rng(); + additional_token_pool_accounts.shuffle(rng); + let (_, _, _, _, instruction) = create_burn_test_instruction( + &sender, &mut rpc, &mut test_indexer, - input_compressed_accounts, - &change_account_merkle_tree, + input_compressed_accounts.as_slice(), + &invalid_change_account_merkle_tree, burn_amount, - true, - None, + false, + BurnInstructionMode::Normal, is_token_22, + 4, + Some(additional_token_pool_accounts.clone()), ) .await; - } - // 3. Burn all delegated tokens - { - let input_compressed_accounts = - test_indexer.get_compressed_token_accounts_by_owner(&sender.pubkey()); - let input_compressed_accounts = input_compressed_accounts + + let (event, _, _) = rpc + .create_and_send_transaction_with_event( + &[instruction], + &payer.pubkey(), + &[&payer, &sender], + None, + ) + .await + .unwrap() + .unwrap(); + let slot = rpc.get_slot().await.unwrap(); + test_indexer.add_event_and_compressed_accounts(slot, &event); + let input_compressed_accounts = test_indexer + .get_compressed_token_accounts_by_owner(&sender.pubkey()) .iter() - .filter(|x| x.token_data.delegate.is_some()) + .filter(|x| x.token_data.amount != 0) .cloned() - .collect::>(); + .collect::>(); let burn_amount = input_compressed_accounts .iter() .map(|x| x.token_data.amount) - .sum::(); - let change_account_merkle_tree = input_compressed_accounts[0] - .compressed_account - .merkle_context - .merkle_tree_pubkey; - burn_test( - &delegate, + .sum(); + + additional_token_pool_accounts.shuffle(rng); + + let (_, _, _, _, instruction) = create_burn_test_instruction( + &sender, &mut rpc, &mut test_indexer, - input_compressed_accounts, - &change_account_merkle_tree, + &input_compressed_accounts, + &merkle_tree_pubkey, burn_amount, - true, - None, + false, + BurnInstructionMode::Normal, is_token_22, + 4, + Some(additional_token_pool_accounts), ) .await; + + rpc.create_and_send_transaction(&[instruction], &payer.pubkey(), &[&payer, &sender]) + .await + .unwrap(); + assert_minted_to_all_token_pools(&mut rpc, 0, &mint) + .await + .unwrap(); } } } @@ -2497,6 +3169,8 @@ async fn failing_tests_burn() { false, BurnInstructionMode::InvalidProof, is_token_22, + 0, + None, ) .await; let res = rpc @@ -2523,6 +3197,8 @@ async fn failing_tests_burn() { true, BurnInstructionMode::Normal, is_token_22, + 0, + None, ) .await; let res = rpc @@ -2558,6 +3234,8 @@ async fn failing_tests_burn() { true, BurnInstructionMode::Normal, is_token_22, + 0, + None, ) .await; let res = rpc @@ -2584,6 +3262,8 @@ async fn failing_tests_burn() { false, BurnInstructionMode::Normal, is_token_22, + 0, + None, ) .await; let res = rpc @@ -2614,6 +3294,8 @@ async fn failing_tests_burn() { false, BurnInstructionMode::InvalidMint, is_token_22, + 0, + None, ) .await; let res = rpc @@ -2645,6 +3327,8 @@ async fn failing_tests_burn() { false, BurnInstructionMode::Normal, is_token_22, + 0, + None, ) .await; let res = rpc @@ -2657,6 +3341,65 @@ async fn failing_tests_burn() { ) .unwrap(); } + // 6. invalid token pool (not initialized) + { + let input_compressed_accounts = + test_indexer.get_compressed_token_accounts_by_owner(&sender.pubkey()); + let burn_amount = 1; + let invalid_change_account_merkle_tree = input_compressed_accounts[0] + .compressed_account + .merkle_context + .nullifier_queue_pubkey; + let (_, _, _, _, instruction) = create_burn_test_instruction( + &sender, + &mut rpc, + &mut test_indexer, + &input_compressed_accounts, + &invalid_change_account_merkle_tree, + burn_amount, + false, + BurnInstructionMode::Normal, + is_token_22, + 1, + None, + ) + .await; + let res = rpc + .create_and_send_transaction(&[instruction], &payer.pubkey(), &[&payer, &sender]) + .await; + assert_rpc_error(res, 0, ErrorCode::InvalidTokenPoolPda.into()).unwrap(); + } + // 7. invalid token pool (invalid mint) + { + let input_compressed_accounts = + test_indexer.get_compressed_token_accounts_by_owner(&sender.pubkey()); + let burn_amount = 1; + let invalid_change_account_merkle_tree = input_compressed_accounts[0] + .compressed_account + .merkle_context + .nullifier_queue_pubkey; + let (_, _, _, _, mut instruction) = create_burn_test_instruction( + &sender, + &mut rpc, + &mut test_indexer, + &input_compressed_accounts, + &invalid_change_account_merkle_tree, + burn_amount, + false, + BurnInstructionMode::Normal, + is_token_22, + 0, + None, + ) + .await; + let mint = create_mint_helper(&mut rpc, &payer).await; + let token_pool = get_token_pool_pda(&mint); + instruction.accounts[4] = AccountMeta::new(token_pool, false); + let res = rpc + .create_and_send_transaction(&[instruction], &payer.pubkey(), &[&payer, &sender]) + .await; + assert_rpc_error(res, 0, ErrorCode::InvalidTokenPoolPda.into()).unwrap(); + } } } @@ -3380,14 +4123,20 @@ async fn test_failing_thaw() { /// 4. Invalid decompression amount +1 /// 5. Invalid decompression amount 0 /// 6: invalid token recipient -/// 7. Invalid compression amount -1 -/// 8. Invalid compression amount +1 -/// 9. Invalid compression amount 0 +/// 7. invalid token pool pda (in struct) +/// 8. invalid token pool pda (in remaining accounts) +/// 9. FailedToDecompress pass multiple correct token accounts with insufficient balance +/// 10. invalid token pool pda from invalid mint (in struct) +/// 11. invalid token pool pda from invalid mint (in remaining accounts) +/// 12. Invalid compression amount -1 +/// 13. Invalid compression amount +1 +/// 14. Invalid compression amount 0 +/// 15. Invalid token pool pda compress (in struct) #[serial] #[tokio::test] async fn test_failing_decompression() { spawn_prover( - true, + false, ProverConfig { run_mode: None, circuits: vec![ProofType::Inclusion], @@ -3409,6 +4158,15 @@ async fn test_failing_decompression() { } else { create_mint_helper(&mut context, &payer).await }; + create_additional_token_pools( + &mut context, + &payer, + &mint, + is_token_22, + NUM_MAX_POOL_ACCOUNTS, + ) + .await + .unwrap(); let amount = 10000u64; mint_tokens_22_helper_with_lamports( &mut context, @@ -3452,6 +4210,7 @@ async fn test_failing_decompression() { &mint, 0, //ProgramError::InvalidAccountData.into(), error code 17179869184 does not fit u32 is_token_22, + None, ) .await .unwrap_err(); @@ -3482,6 +4241,7 @@ async fn test_failing_decompression() { &mint, ErrorCode::InvalidTokenPoolPda.into(), is_token_22, + None, ) .await .unwrap(); @@ -3510,6 +4270,7 @@ async fn test_failing_decompression() { &mint, ErrorCode::InvalidTokenPoolPda.into(), is_token_22, + None, ) .await .unwrap(); @@ -3530,6 +4291,7 @@ async fn test_failing_decompression() { &mint, ErrorCode::SumCheckFailed.into(), is_token_22, + None, ) .await .unwrap(); @@ -3550,6 +4312,7 @@ async fn test_failing_decompression() { &mint, ErrorCode::ComputeOutputSumFailed.into(), is_token_22, + None, ) .await .unwrap(); @@ -3570,6 +4333,7 @@ async fn test_failing_decompression() { &mint, ErrorCode::SumCheckFailed.into(), is_token_22, + None, ) .await .unwrap(); @@ -3590,6 +4354,133 @@ async fn test_failing_decompression() { &mint, ErrorCode::IsTokenPoolPda.into(), is_token_22, + None, + ) + .await + .unwrap(); + } + // Test 7: invalid token pool pda (in struct) + { + failing_compress_decompress( + &sender, + &mut context, + &mut test_indexer, + input_compressed_account.clone(), + decompress_amount, // needs to be consistent with compression amount + &merkle_tree_pubkey, + decompress_amount, + false, + &token_account_keypair.pubkey(), + Some(get_token_pool_pda_with_bump(&mint, NUM_MAX_POOL_ACCOUNTS)), + &mint, + anchor_lang::error::ErrorCode::AccountNotInitialized.into(), + is_token_22, + None, + ) + .await + .unwrap(); + } + // Test 8: invalid token pool pda (in remaining accounts) + { + failing_compress_decompress( + &sender, + &mut context, + &mut test_indexer, + input_compressed_account.clone(), + decompress_amount, // needs to be consistent with compression amount + &merkle_tree_pubkey, + decompress_amount, + false, + &token_account_keypair.pubkey(), + Some(get_token_pool_pda_with_bump(&mint, 3)), + &mint, + ErrorCode::InvalidTokenPoolPda.into(), + is_token_22, + Some(vec![get_token_pool_pda_with_bump( + &mint, + NUM_MAX_POOL_ACCOUNTS, + )]), + ) + .await + .unwrap(); + } + // Test 9: FailedToDecompress pass multiple correct token accounts with insufficient balance + { + let token_pool = get_token_pool_pda_with_bump(&mint, 3); + let mut account = context.get_account(token_pool).await.unwrap().unwrap(); + println!("token pool account {:?}", token_pool); + let amount = + TokenAccount::try_deserialize_unchecked(&mut &*account.data.as_mut_slice()) + .unwrap() + .amount; + println!("{:?}", amount); + failing_compress_decompress( + &sender, + &mut context, + &mut test_indexer, + input_compressed_account.clone(), + decompress_amount, // needs to be consistent with compression amount + &merkle_tree_pubkey, + decompress_amount, + false, + &token_account_keypair.pubkey(), + Some(token_pool), + &mint, + ErrorCode::FailedToDecompress.into(), + is_token_22, + Some(vec![ + token_pool, + get_token_pool_pda_with_bump(&mint, 1), + get_token_pool_pda_with_bump(&mint, 2), + get_token_pool_pda_with_bump(&mint, 4), + ]), + ) + .await + .unwrap(); + } + + let invalid_mint = create_mint_22_helper(&mut context, &payer).await; + // Test 10: invalid token pool pda from invalid mint (in struct) + { + failing_compress_decompress( + &sender, + &mut context, + &mut test_indexer, + input_compressed_account.clone(), + decompress_amount, // needs to be consistent with compression amount + &merkle_tree_pubkey, + decompress_amount, + false, + &token_account_keypair.pubkey(), + Some(get_token_pool_pda_with_bump(&invalid_mint, 0)), + &mint, + ErrorCode::InvalidTokenPoolPda.into(), + is_token_22, + Some(vec![get_token_pool_pda_with_bump( + &mint, + NUM_MAX_POOL_ACCOUNTS, + )]), + ) + .await + .unwrap(); + } + // Test 11: invalid token pool pda from invalid mint (in remaining accounts) + { + failing_compress_decompress( + &sender, + &mut context, + &mut test_indexer, + input_compressed_account.clone(), + decompress_amount, // needs to be consistent with compression amount + &merkle_tree_pubkey, + decompress_amount, + false, + &token_account_keypair.pubkey(), + Some(get_token_pool_pda_with_bump(&mint, 4)), + &mint, + ErrorCode::InvalidTokenPoolPda.into(), + is_token_22, + Some(vec![get_token_pool_pda_with_bump(&invalid_mint, 0)]), ) .await .unwrap(); @@ -3606,10 +4497,12 @@ async fn test_failing_decompression() { &token_account_keypair.pubkey(), None, is_token_22, + 0, + None, ) .await; let compress_amount = decompress_amount - 100; - // Test 7: invalid compression amount -1 + // Test 12: invalid compression amount -1 { failing_compress_decompress( &sender, @@ -3625,11 +4518,12 @@ async fn test_failing_decompression() { &mint, ErrorCode::ComputeOutputSumFailed.into(), is_token_22, + None, ) .await .unwrap(); } - // Test 8: invalid compression amount +1 + // Test 13: invalid compression amount +1 { failing_compress_decompress( &sender, @@ -3645,11 +4539,12 @@ async fn test_failing_decompression() { &mint, ErrorCode::SumCheckFailed.into(), is_token_22, + None, ) .await .unwrap(); } - // Test 9: invalid compression amount 0 + // Test 14: invalid compression amount 0 { failing_compress_decompress( &sender, @@ -3665,6 +4560,28 @@ async fn test_failing_decompression() { &mint, ErrorCode::ComputeOutputSumFailed.into(), is_token_22, + None, + ) + .await + .unwrap(); + } + // Test 15: invalid token pool pda (in struct) + { + failing_compress_decompress( + &sender, + &mut context, + &mut test_indexer, + Vec::new(), + compress_amount, // needs to be consistent with compression amount + &merkle_tree_pubkey, + compress_amount, + true, + &token_account_keypair.pubkey(), + Some(get_token_pool_pda_with_bump(&invalid_mint, 0)), + &mint, + ErrorCode::InvalidTokenPoolPda.into(), + is_token_22, + None, ) .await .unwrap(); @@ -3680,6 +4597,8 @@ async fn test_failing_decompression() { &token_account_keypair.pubkey(), None, is_token_22, + 0, + None, ) .await; } @@ -3701,6 +4620,7 @@ pub async fn failing_compress_decompress( mint: &Pubkey, error_code: u32, is_token_22: bool, + additional_token_pools: Option>, ) -> Result<(), RpcError> { let max_amount: u64 = input_compressed_accounts .iter() @@ -3774,6 +4694,7 @@ pub async fn failing_compress_decompress( None, None, is_token_22, + &additional_token_pools.unwrap_or_default(), ) .unwrap(); let instructions = if !is_compress { @@ -4239,6 +5160,7 @@ async fn perform_transfer_failing_test( None, None, false, + &[], ) .unwrap(); diff --git a/test-programs/system-cpi-test/tests/test_program_owned_trees.rs b/test-programs/system-cpi-test/tests/test_program_owned_trees.rs index 33d44c646d..8f73fbb67d 100644 --- a/test-programs/system-cpi-test/tests/test_program_owned_trees.rs +++ b/test-programs/system-cpi-test/tests/test_program_owned_trees.rs @@ -87,6 +87,7 @@ async fn test_program_owned_merkle_tree() { vec![recipient_keypair.pubkey(); 1], None, false, + 0, ); let pre_merkle_tree = get_concurrent_merkle_tree::< StateMerkleTreeAccount, @@ -152,6 +153,7 @@ async fn test_program_owned_merkle_tree() { vec![recipient_keypair.pubkey(); 1], None, false, + 0, ); let latest_blockhash = rpc.get_latest_blockhash().await.unwrap(); diff --git a/test-utils/src/assert_token_tx.rs b/test-utils/src/assert_token_tx.rs index 6dba2dc066..b7697ad300 100644 --- a/test-utils/src/assert_token_tx.rs +++ b/test-utils/src/assert_token_tx.rs @@ -5,10 +5,7 @@ use crate::assert_compressed_tx::{ use anchor_lang::AnchorSerialize; use forester_utils::indexer::{Indexer, TokenDataWithContext}; use light_client::rpc::RpcConnection; -use light_compressed_token::{ - get_token_pool_pda, - process_transfer::{get_cpi_authority_pda, TokenTransferOutputData}, -}; +use light_compressed_token::process_transfer::{get_cpi_authority_pda, TokenTransferOutputData}; use light_system_program::sdk::{ compressed_account::CompressedAccountWithMerkleContext, event::PublicTransactionEvent, }; @@ -200,6 +197,7 @@ pub async fn assert_mint_to<'a, R: RpcConnection, I: Indexer>( created_token_accounts: &[TokenDataWithContext], previous_mint_supply: u64, previous_sol_pool_amount: u64, + token_pool_pda: Pubkey, ) { let mut created_token_accounts = created_token_accounts.to_vec(); for (recipient, amount) in recipients.iter().zip(amounts) { @@ -221,10 +219,10 @@ pub async fn assert_mint_to<'a, R: RpcConnection, I: Indexer>( let sum_amounts = amounts.iter().sum::(); assert_eq!(mint_account.supply, previous_mint_supply + sum_amounts); - let pool = get_token_pool_pda(&mint); - let pool_account = - spl_token::state::Account::unpack(&rpc.get_account(pool).await.unwrap().unwrap().data) - .unwrap(); + let pool_account = spl_token::state::Account::unpack( + &rpc.get_account(token_pool_pda).await.unwrap().unwrap().data, + ) + .unwrap(); assert_eq!(pool_account.amount, previous_sol_pool_amount + sum_amounts); } diff --git a/test-utils/src/e2e_test_env.rs b/test-utils/src/e2e_test_env.rs index 6080a3db0b..bec98cbf2d 100644 --- a/test-utils/src/e2e_test_env.rs +++ b/test-utils/src/e2e_test_env.rs @@ -1829,6 +1829,7 @@ where false, transaction_paramets, false, + 0, ) .await; self.stats.spl_burned += 1; @@ -1974,6 +1975,8 @@ where &token_account, transaction_paramets, false, + 0, // TODO: make random + None, ) .await; self.stats.spl_compress += 1; @@ -2051,6 +2054,8 @@ where &token_account, transaction_paramets, false, + 0, // TODO: make random + None, ) .await; self.stats.spl_decompress += 1; diff --git a/test-utils/src/spl.rs b/test-utils/src/spl.rs index f40c16f61f..438bf674a4 100644 --- a/test-utils/src/spl.rs +++ b/test-utils/src/spl.rs @@ -1,6 +1,7 @@ use anchor_spl::token::{Mint, TokenAccount}; use forester_utils::create_account_instruction; use forester_utils::indexer::{Indexer, TokenDataWithContext}; +use light_compressed_token::mint_sdk::create_add_token_pool_instruction; use light_compressed_token::process_compress_spl_token_account::sdk::create_compress_spl_token_account_instruction; use light_compressed_token::{ burn::sdk::{create_burn_instruction, CreateBurnInstructionInputs}, @@ -15,6 +16,7 @@ use light_compressed_token::{ token_data::AccountState, TokenData, }; +use light_compressed_token::{constants::NUM_MAX_POOL_ACCOUNTS, get_token_pool_pda_with_bump}; use light_hasher::Poseidon; use light_system_program::{ invoke::processor::CompressedProof, @@ -132,6 +134,34 @@ pub async fn mint_tokens_22_helper_with_lamports recipients: Vec, lamports: Option, token_22: bool, +) { + mint_tokens_22_helper_with_lamports_and_bump( + rpc, + test_indexer, + merkle_tree_pubkey, + mint_authority, + mint, + amounts, + recipients, + lamports, + token_22, + 0, + ) + .await +} + +#[allow(clippy::too_many_arguments)] +pub async fn mint_tokens_22_helper_with_lamports_and_bump>( + rpc: &mut R, + test_indexer: &mut I, + merkle_tree_pubkey: &Pubkey, + mint_authority: &Keypair, + mint: &Pubkey, + amounts: Vec, + recipients: Vec, + lamports: Option, + token_22: bool, + token_pool_bump: u8, ) { let payer_pubkey = mint_authority.pubkey(); let instruction = create_mint_to_instruction( @@ -143,6 +173,7 @@ pub async fn mint_tokens_22_helper_with_lamports recipients.clone(), lamports, token_22, + token_pool_bump, ); let output_merkle_tree_accounts = @@ -154,7 +185,7 @@ pub async fn mint_tokens_22_helper_with_lamports .unwrap() .supply; - let pool: Pubkey = get_token_pool_pda(mint); + let pool: Pubkey = get_token_pool_pda_with_bump(mint, token_pool_bump); let previous_pool_amount = spl_token::state::Account::unpack(&rpc.get_account(pool).await.unwrap().unwrap().data) .unwrap() @@ -182,6 +213,7 @@ pub async fn mint_tokens_22_helper_with_lamports &created_token_accounts, previous_mint_supply, previous_pool_amount, + pool, ) .await; } @@ -338,6 +370,40 @@ pub fn create_initialize_mint_22_instructions( ) } +pub async fn create_additional_token_pools( + rpc: &mut R, + payer: &Keypair, + mint: &Pubkey, + is_token_22: bool, + num: u8, +) -> Result, RpcError> { + let mut instructions = Vec::new(); + let mut created_token_pools = Vec::new(); + + for token_pool_bump in 0..NUM_MAX_POOL_ACCOUNTS { + if instructions.len() == num as usize { + break; + } + let token_pool_pda = get_token_pool_pda_with_bump(mint, token_pool_bump); + let account = rpc.get_account(token_pool_pda).await.unwrap(); + println!("bump {}", token_pool_bump); + println!("account exists {:?}", account.is_some()); + if account.is_none() { + created_token_pools.push(token_pool_pda); + let instruction = create_add_token_pool_instruction( + &payer.pubkey(), + mint, + token_pool_bump, + is_token_22, + ); + instructions.push(instruction); + } + } + rpc.create_and_send_transaction(&instructions, &payer.pubkey(), &[payer]) + .await?; + Ok(created_token_pools) +} + /// Creates a spl token account and initializes it with the given mint and owner. /// This function is useful to create token accounts for spl compression and decompression tests. pub async fn create_token_account( @@ -568,6 +634,7 @@ pub async fn compressed_transfer_22_test>( delegate_change_account_index, None, token_22, + &[], ) .unwrap(); let sum_input_lamports = input_compressed_accounts @@ -661,13 +728,18 @@ pub async fn decompress_test>( recipient_token_account: &Pubkey, transaction_params: Option, is_token_22: bool, + token_pool_bump: u8, + additonal_pool_accounts: Option>, ) { let max_amount: u64 = input_compressed_accounts .iter() .map(|x| x.token_data.amount) .sum(); + println!("max_amount: {}", max_amount); + println!("amount: {}", amount); + let output_amount = max_amount - amount; let change_out_compressed_account = TokenTransferOutputData { - amount: max_amount - amount, + amount: output_amount, owner: payer.pubkey(), lamports: None, merkle_tree: *output_merkle_tree_pubkey, @@ -690,6 +762,8 @@ pub async fn decompress_test>( ) .await; let mint = input_compressed_accounts[0].token_data.mint; + let token_pool_pda = get_token_pool_pda_with_bump(&mint, token_pool_bump); + let instruction = create_transfer_instruction( &rpc.get_payer().pubkey(), &payer.pubkey(), // authority @@ -710,16 +784,20 @@ pub async fn decompress_test>( .map(|x| &x.compressed_account.compressed_account) .cloned() .collect::>(), - mint, // mint - None, // owner_if_delegate_change_account_index - false, // is_compress - Some(amount), // compression_amount - Some(get_token_pool_pda(&mint)), // token_pool_pda - Some(*recipient_token_account), // compress_or_decompress_token_account + mint, // mint + None, // owner_if_delegate_change_account_index + false, // is_compress + Some(amount), // compression_amount + Some(token_pool_pda), // token_pool_pda + Some(*recipient_token_account), // compress_or_decompress_token_account true, None, None, is_token_22, + additonal_pool_accounts + .clone() + .unwrap_or_default() + .as_slice(), ) .unwrap(); let output_merkle_tree_pubkeys = vec![*output_merkle_tree_pubkey]; @@ -739,6 +817,30 @@ pub async fn decompress_test>( .data, ) .unwrap(); + let mut token_pool_pre_balances = vec![ + spl_token::state::Account::unpack( + &rpc.get_account(token_pool_pda).await.unwrap().unwrap().data, + ) + .unwrap() + .amount, + ]; + for additional_pool_account in additonal_pool_accounts + .clone() + .unwrap_or_default() + .as_slice() + { + token_pool_pre_balances.push( + spl_token::state::Account::unpack( + &rpc.get_account(*additional_pool_account) + .await + .unwrap() + .unwrap() + .data, + ) + .unwrap() + .amount, + ); + } let context_payer = rpc.get_payer().insecure_clone(); let (event, _signature, _) = rpc .create_and_send_transaction_with_event::( @@ -769,7 +871,6 @@ pub async fn decompress_test>( None, ) .await; - let recipient_token_account_data = spl_token::state::Account::unpack( &rpc.get_account(*recipient_token_account) .await @@ -778,10 +879,44 @@ pub async fn decompress_test>( .data, ) .unwrap(); + println!("amount: {}", amount); assert_eq!( recipient_token_account_data.amount, recipient_token_account_data_pre.amount + amount ); + let token_pool_post_balance = spl_token::state::Account::unpack( + &rpc.get_account(token_pool_pda).await.unwrap().unwrap().data, + ) + .unwrap() + .amount; + assert_eq!( + token_pool_post_balance, + token_pool_pre_balances[0].saturating_sub(amount) + ); + let mut amount = amount - token_pool_pre_balances[0]; + for (i, additional_account) in additonal_pool_accounts + .unwrap_or_default() + .iter() + .enumerate() + { + let post_balance = spl_token::state::Account::unpack( + &rpc.get_account(*additional_account) + .await + .unwrap() + .unwrap() + .data, + ) + .unwrap() + .amount; + amount -= token_pool_pre_balances[i + 1]; + if amount == 0 { + break; + } + assert_eq!( + post_balance, + token_pool_pre_balances[i + 1].saturating_sub(amount) + ); + } } #[allow(clippy::too_many_arguments)] @@ -795,6 +930,7 @@ pub async fn perform_compress_spl_token_account> merkle_tree_pubkey: &Pubkey, remaining_amount: Option, is_token_22: bool, + token_pool_bump: u8, ) -> Result<(), RpcError> { let pre_token_account_amount = spl_token::state::Account::unpack( &rpc.get_account(*token_account).await.unwrap().unwrap().data, @@ -811,6 +947,7 @@ pub async fn perform_compress_spl_token_account> merkle_tree_pubkey, token_account, is_token_22, + token_pool_bump, ); let (event, _, _) = rpc .create_and_send_transaction_with_event::( @@ -867,6 +1004,8 @@ pub async fn compress_test>( sender_token_account: &Pubkey, transaction_params: Option, is_token_22: bool, + token_pool_bump: u8, + additonal_pool_accounts: Option>, ) { let output_compressed_account = TokenTransferOutputData { amount, @@ -882,18 +1021,19 @@ pub async fn compress_test>( &[output_compressed_account], // output_compressed_accounts &Vec::new(), // root_indices &None, - &Vec::new(), // input_token_data - &Vec::new(), // input_compressed_accounts - *mint, // mint - None, // owner_if_delegate_is_signer - true, // is_compress - Some(amount), // compression_amount - Some(get_token_pool_pda(mint)), // token_pool_pda - Some(*sender_token_account), // compress_or_decompress_token_account + &Vec::new(), // input_token_data + &Vec::new(), // input_compressed_accounts + *mint, // mint + None, // owner_if_delegate_is_signer + true, // is_compress + Some(amount), // compression_amount + Some(get_token_pool_pda_with_bump(mint, token_pool_bump)), // token_pool_pda + Some(*sender_token_account), // compress_or_decompress_token_account true, None, None, is_token_22, + additonal_pool_accounts.unwrap_or_default().as_slice(), ) .unwrap(); let output_merkle_tree_pubkeys = vec![*output_merkle_tree_pubkey]; @@ -1424,6 +1564,7 @@ pub async fn burn_test>( signer_is_delegate: bool, transaction_params: Option, is_token_22: bool, + token_pool_bump: u8, ) { let ( input_compressed_account_hashes, @@ -1441,6 +1582,8 @@ pub async fn burn_test>( signer_is_delegate, BurnInstructionMode::Normal, is_token_22, + token_pool_bump, + None, ) .await; let output_merkle_tree_pubkeys = vec![*change_account_merkle_tree; 1]; @@ -1453,7 +1596,7 @@ pub async fn burn_test>( Vec::new() }; - let token_pool_pda_address = get_token_pool_pda(&mint); + let token_pool_pda_address = get_token_pool_pda_with_bump(&mint, token_pool_bump); let pre_token_pool_account = rpc .get_account(token_pool_pda_address) .await @@ -1564,6 +1707,8 @@ pub async fn create_burn_test_instruction>( signer_is_delegate: bool, mode: BurnInstructionMode, is_token_22: bool, + token_pool_bump: u8, + additonal_pool_accounts: Option>, ) -> (Vec<[u8; 32]>, Vec, Pubkey, u64, Instruction) { let input_compressed_account_hashes = input_compressed_accounts .iter() @@ -1619,6 +1764,8 @@ pub async fn create_burn_test_instruction>( signer_is_delegate, burn_amount, is_token_22, + token_pool_bump, + additonal_pool_accounts: additonal_pool_accounts.unwrap_or_default(), }; let input_amount_sum = input_compressed_accounts .iter()