Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: (de)compression transfer ix token pool check #1032

Merged
merged 2 commits into from
Jul 26, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions programs/compressed-token/src/instructions/transfer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@ use light_system_program::{
sdk::accounts::{InvokeAccounts, SignerAccounts},
};

use crate::POOL_SEED;
#[derive(Accounts)]
pub struct TransferInstruction<'info> {
/// UNCHECKED: only pays fees.
Expand Down Expand Up @@ -34,7 +33,7 @@ pub struct TransferInstruction<'info> {
/// (different program) checked in light system program to derive
/// cpi_authority_pda and check that this program is the signer of the cpi.
pub self_program: Program<'info, crate::program::LightCompressedToken>,
#[account(mut, seeds = [POOL_SEED, &token_pool_pda.mint.key().to_bytes()],bump)]
#[account(mut)]
pub token_pool_pda: Option<Account<'info, TokenAccount>>,
#[account(mut, constraint= if token_pool_pda.is_some() {Ok(token_pool_pda.as_ref().unwrap().key() != compress_or_decompress_token_account.key())}else {err!(crate::ErrorCode::TokenPoolPdaUndefined)}? @crate::ErrorCode::IsTokenPoolPda)]
pub compress_or_decompress_token_account: Option<Account<'info, TokenAccount>>,
Expand Down
1 change: 1 addition & 0 deletions programs/compressed-token/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -183,4 +183,5 @@ pub enum ErrorCode {
TokenPoolPdaUndefined,
#[msg("Compress or decompress recipient is the same account as the token pool pda.")]
IsTokenPoolPda,
InvalidTokenPoolPda,
}
23 changes: 20 additions & 3 deletions programs/compressed-token/src/spl_compression.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use anchor_spl::token::Transfer;

use crate::{
process_transfer::get_cpi_signer_seeds, CompressedTokenInstructionDataTransfer,
TransferInstruction,
TransferInstruction, POOL_SEED,
};

pub fn process_compression_or_decompression<'info>(
Expand All @@ -17,6 +17,20 @@ pub fn process_compression_or_decompression<'info>(
}
}

pub fn spl_token_pool_derivation(
mint: &Pubkey,
program_id: &Pubkey,
token_pool_pubkey: &Pubkey,
) -> Result<()> {
let seeds = &[POOL_SEED, &mint.to_bytes()[..]];
let (pda, _bump_seed) = Pubkey::find_program_address(seeds, program_id);
if pda == *token_pool_pubkey {
Ok(())
} else {
err!(crate::ErrorCode::InvalidTokenPoolPda)
}
}

pub fn decompress_spl_tokens<'info>(
inputs: &CompressedTokenInstructionDataTransfer,
ctx: &Context<'_, '_, '_, 'info, TransferInstruction<'info>>,
Expand All @@ -29,6 +43,8 @@ pub fn decompress_spl_tokens<'info>(
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 {
Some(amount) => amount,
None => return err!(crate::ErrorCode::DeCompressAmountUndefinedForDecompress),
Expand All @@ -50,10 +66,11 @@ pub fn compress_spl_tokens<'info>(
inputs: &CompressedTokenInstructionDataTransfer,
ctx: &Context<'_, '_, '_, 'info, TransferInstruction<'info>>,
) -> Result<()> {
let recipient = match ctx.accounts.token_pool_pda.as_ref() {
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),
};
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),
Expand All @@ -65,7 +82,7 @@ pub fn compress_spl_tokens<'info>(
.as_ref()
.unwrap()
.to_account_info(),
&recipient,
&recipient_token_pool,
&ctx.accounts.authority.to_account_info(),
&ctx.accounts
.token_program
Expand Down
45 changes: 33 additions & 12 deletions test-programs/compressed-token-test/tests/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2655,7 +2655,7 @@ async fn test_failing_decompression() {
&mut context,
&mut test_indexer,
input_compressed_account.clone(),
decompress_amount, // need to be consistent with compression amount
decompress_amount, // needs to be consistent with compression amount
&merkle_tree_pubkey,
decompress_amount,
false,
Expand All @@ -2667,7 +2667,7 @@ async fn test_failing_decompression() {
.await
.unwrap_err();
}
// Test 2: invalid token pool pda
// Test 2: invalid token pool pda (compress and decompress)
{
let invalid_token_account_keypair = Keypair::new();
create_token_account(&mut context, &mint, &invalid_token_account_keypair, &payer)
Expand All @@ -2678,14 +2678,35 @@ async fn test_failing_decompression() {
&mut context,
&mut test_indexer,
input_compressed_account.clone(),
decompress_amount, // need to be consistent with compression amount
decompress_amount, // needs to be consistent with compression amount
&merkle_tree_pubkey,
decompress_amount - 1,
decompress_amount,
false,
&token_account_keypair.pubkey(),
Some(invalid_token_account_keypair.pubkey()),
&mint,
anchor_lang::error::ErrorCode::ConstraintSeeds.into(),
ErrorCode::InvalidTokenPoolPda.into(),
)
.await
.unwrap();

let invalid_token_account_keypair = Keypair::new();
create_token_account(&mut context, &mint, &invalid_token_account_keypair, &payer)
.await
.unwrap();
failing_compress_decompress(
&sender,
&mut context,
&mut test_indexer,
input_compressed_account.clone(),
0, // needs to be consistent with compression amount
&merkle_tree_pubkey,
0,
true,
&token_account_keypair.pubkey(),
Some(invalid_token_account_keypair.pubkey()),
&mint,
ErrorCode::InvalidTokenPoolPda.into(),
)
.await
.unwrap();
Expand All @@ -2697,7 +2718,7 @@ async fn test_failing_decompression() {
&mut context,
&mut test_indexer,
input_compressed_account.clone(),
decompress_amount, // need to be consistent with compression amount
decompress_amount, // needs to be consistent with compression amount
&merkle_tree_pubkey,
decompress_amount - 1,
false,
Expand All @@ -2716,7 +2737,7 @@ async fn test_failing_decompression() {
&mut context,
&mut test_indexer,
input_compressed_account.clone(),
decompress_amount, // need to be consistent with compression amount
decompress_amount, // needs to be consistent with compression amount
&merkle_tree_pubkey,
decompress_amount + 1,
false,
Expand All @@ -2735,7 +2756,7 @@ async fn test_failing_decompression() {
&mut context,
&mut test_indexer,
input_compressed_account.clone(),
decompress_amount, // need to be consistent with compression amount
decompress_amount, // needs to be consistent with compression amount
&merkle_tree_pubkey,
0,
false,
Expand All @@ -2754,7 +2775,7 @@ async fn test_failing_decompression() {
&mut context,
&mut test_indexer,
input_compressed_account.clone(),
decompress_amount, // need to be consistent with compression amount
decompress_amount, // needs to be consistent with compression amount
&merkle_tree_pubkey,
decompress_amount,
false,
Expand Down Expand Up @@ -2787,7 +2808,7 @@ async fn test_failing_decompression() {
&mut context,
&mut test_indexer,
Vec::new(),
compress_amount, // need to be consistent with compression amount
compress_amount, // needs to be consistent with compression amount
&merkle_tree_pubkey,
compress_amount - 1,
true,
Expand All @@ -2806,7 +2827,7 @@ async fn test_failing_decompression() {
&mut context,
&mut test_indexer,
Vec::new(),
compress_amount, // need to be consistent with compression amount
compress_amount, // needs to be consistent with compression amount
&merkle_tree_pubkey,
compress_amount + 1,
true,
Expand All @@ -2825,7 +2846,7 @@ async fn test_failing_decompression() {
&mut context,
&mut test_indexer,
Vec::new(),
compress_amount, // need to be consistent with compression amount
compress_amount, // needs to be consistent with compression amount
&merkle_tree_pubkey,
0,
true,
Expand Down