From 28ef6b3ce71c3ccc05a7f0870a03536dffb9b9be Mon Sep 17 00:00:00 2001 From: ananas-block <58553958+ananas-block@users.noreply.github.com> Date: Wed, 25 Dec 2024 15:22:20 +0000 Subject: [PATCH] fix: invoke_token_program_with_multiple_token_pool_accounts (#1427) --- programs/compressed-token/src/lib.rs | 1 + .../compressed-token/src/spl_compression.rs | 60 +++++++++++------ .../compressed-token-test/tests/test.rs | 67 +++++++++++++++++-- 3 files changed, 102 insertions(+), 26 deletions(-) diff --git a/programs/compressed-token/src/lib.rs b/programs/compressed-token/src/lib.rs index eee49faf6..df73cd46f 100644 --- a/programs/compressed-token/src/lib.rs +++ b/programs/compressed-token/src/lib.rs @@ -233,4 +233,5 @@ pub enum ErrorCode { InvalidTokenPoolBump, FailedToDecompress, FailedToBurnSplTokensFromTokenPool, + NoMatchingBumpFound, } diff --git a/programs/compressed-token/src/spl_compression.rs b/programs/compressed-token/src/spl_compression.rs index 317734201..8bce08b16 100644 --- a/programs/compressed-token/src/spl_compression.rs +++ b/programs/compressed-token/src/spl_compression.rs @@ -5,7 +5,7 @@ use anchor_spl::{token::TokenAccount, token_interface}; use crate::{ constants::{NUM_MAX_POOL_ACCOUNTS, POOL_SEED}, process_transfer::get_cpi_signer_seeds, - CompressedTokenInstructionDataTransfer, TransferInstruction, + CompressedTokenInstructionDataTransfer, ErrorCode, TransferInstruction, }; pub fn process_compression_or_decompression<'info>( @@ -27,7 +27,7 @@ pub fn spl_token_pool_derivation( if check_spl_token_pool_derivation(mint_bytes, token_pool_pubkey, bump) { Ok(()) } else { - err!(crate::ErrorCode::InvalidTokenPoolPda) + err!(ErrorCode::InvalidTokenPoolPda) } } @@ -52,15 +52,15 @@ pub fn decompress_spl_tokens<'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), + None => return err!(ErrorCode::DecompressRecipientUndefinedForDecompress), }; let 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), + None => return err!(ErrorCode::CompressedPdaUndefinedForDecompress), }; let amount = match inputs.compress_or_decompress_amount { Some(amount) => amount, - None => return err!(crate::ErrorCode::DeCompressAmountUndefinedForDecompress), + None => return err!(ErrorCode::DeCompressAmountUndefinedForDecompress), }; invoke_token_program_with_multiple_token_pool_accounts::( ctx.remaining_accounts, @@ -80,6 +80,20 @@ pub fn decompress_spl_tokens<'info>( /// Executes a token program instruction with multiple token pool accounts. /// Supported instructions are burn and transfer to decompress spl tokens. +/// Logic: +/// 1. iterate over at most NUM_MAX_POOL_ACCOUNTS token pool accounts. +/// 2. start with passed in token pool account. +/// 3. determine whether complete amount can be transferred or burned. +/// 4. Skip if action amount is zero. +/// 5. check if the token pool account is derived from the mint. +/// 6. return error if the token pool account is not derived +/// from any combination of mint and bump. +/// 7. burn or transfer the amount from the token pool account. +/// 8. remove bump from the list of bumps. +/// 9. reduce the amount by the transferred or burned amount. +/// 10. continue until the amount is zero. +/// 11. Return if complete amount has been transferred or burned. +/// 12. return error if the amount is not zero and the number of accounts has been exhausted. #[allow(clippy::too_many_arguments)] pub fn invoke_token_program_with_multiple_token_pool_accounts<'info, const IS_BURN: bool>( remaining_accounts: &[AccountInfo<'info>], @@ -90,22 +104,25 @@ pub fn invoke_token_program_with_multiple_token_pool_accounts<'info, const IS_BU token_program: AccountInfo<'info>, mut token_pool_pda: AccountInfo<'info>, mut amount: u64, -) -> std::result::Result<(), Error> { - let mut token_pool_bumps = (0..NUM_MAX_POOL_ACCOUNTS).collect::>(); - +) -> Result<()> { + let mut token_pool_bumps: Vec = (0..NUM_MAX_POOL_ACCOUNTS).collect(); + // 1. iterate over at most NUM_MAX_POOL_ACCOUNTS token pool accounts. for i in 0..NUM_MAX_POOL_ACCOUNTS { + // 2. start with passed in token pool account.token_pool_bumps if i != 0 { token_pool_pda = 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)? + .map_err(|_| ErrorCode::InvalidTokenPoolPda)? .amount; + // 3. determine whether complete amount can be transferred or burned. let action_amount = std::cmp::min(amount, token_pool_amount); + // 4. Skip if action amount is zero. if action_amount == 0 { continue; } - let mut remove_index = 0; + // 5. check if the token pool account is derived from the mint for any bump. for (index, i) in token_pool_bumps.iter().enumerate() { if check_spl_token_pool_derivation(mint_bytes.as_slice(), &token_pool_pda.key(), &[*i]) { @@ -127,25 +144,30 @@ pub fn invoke_token_program_with_multiple_token_pool_accounts<'info, const IS_BU action_amount, )?; } - - remove_index = index; + token_pool_bumps.remove(index); + amount = amount.saturating_sub(action_amount); + break; + } else if index == token_pool_bumps.len() - 1 { + // 6. return error if the token pool account is not derived + // from any combination of mint and bump. + return err!(crate::ErrorCode::NoMatchingBumpFound); } } - token_pool_bumps.remove(remove_index); - amount = amount.saturating_sub(action_amount); + // 11. Return if complete amount has been transferred or burned. if amount == 0 { return Ok(()); } } + // 12. return error if the amount is not zero and the number of accounts has been exhausted. msg!("Remaining amount: {}.", amount); if IS_BURN { msg!("Token pool account balances insufficient for burn. \nTry to pass more token pool accounts."); - err!(crate::ErrorCode::FailedToBurnSplTokensFromTokenPool) + err!(ErrorCode::FailedToBurnSplTokensFromTokenPool) } else { msg!("Token pool account balances insufficient for decompression. \nTry to pass more token pool accounts."); - err!(crate::ErrorCode::FailedToDecompress) + err!(ErrorCode::FailedToDecompress) } } @@ -155,11 +177,11 @@ pub fn compress_spl_tokens<'info>( ) -> Result<()> { let recipient_token_pool = match ctx.accounts.token_pool_pda.as_ref() { Some(token_pool_pda) => token_pool_pda.to_account_info(), - None => return err!(crate::ErrorCode::CompressedPdaUndefinedForCompress), + None => return err!(ErrorCode::CompressedPdaUndefinedForCompress), }; let amount = match inputs.compress_or_decompress_amount { Some(amount) => amount, - None => return err!(crate::ErrorCode::DeCompressAmountUndefinedForCompress), + None => return err!(ErrorCode::DeCompressAmountUndefinedForCompress), }; let mint_bytes = inputs.mint.to_bytes(); @@ -185,7 +207,7 @@ pub fn compress_spl_tokens<'info>( return Ok(()); } } - err!(crate::ErrorCode::InvalidTokenPoolPda) + err!(ErrorCode::InvalidTokenPoolPda) } /// Invoke the spl token burn instruction with cpi authority pda as signer. diff --git a/test-programs/compressed-token-test/tests/test.rs b/test-programs/compressed-token-test/tests/test.rs index 2d948dad3..a4c41a75f 100644 --- a/test-programs/compressed-token-test/tests/test.rs +++ b/test-programs/compressed-token-test/tests/test.rs @@ -4112,6 +4112,7 @@ async fn test_failing_thaw() { /// 6: invalid token recipient /// 7. invalid token pool pda (in struct) /// 8. invalid token pool pda (in remaining accounts) +/// 8.1. invalid derived token pool pda (in struct and 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) @@ -4390,17 +4391,69 @@ async fn test_failing_decompression() { ) .await .unwrap(); + } // Test 8.1: invalid derived token pool pda (in struct and remaining accounts) + { + let token_account_keypair_2 = Keypair::new(); + create_token_2022_account( + &mut context, + &mint, + &token_account_keypair_2, + &sender, + is_token_22, + ) + .await + .unwrap(); + mint_spl_tokens( + &mut context, + &mint, + &token_account_keypair_2.pubkey(), + &payer.pubkey(), + &payer, + amount, + is_token_22, + ) + .await + .unwrap(); + 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_account_keypair_2.pubkey()), + &mint, + ErrorCode::NoMatchingBumpFound.into(), + is_token_22, + Some(vec![token_account_keypair_2.pubkey()]), + ) + .await + .unwrap(); + 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::NoMatchingBumpFound.into(), + is_token_22, + Some(vec![token_account_keypair_2.pubkey()]), + ) + .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,