Skip to content
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
2 changes: 1 addition & 1 deletion program-tests/compressed-token-test/tests/ctoken/create.rs
Original file line number Diff line number Diff line change
Expand Up @@ -206,7 +206,7 @@ async fn test_create_compressible_token_account_failing() {
&mut context,
compressible_data,
"account_already_initialized",
0, // AlreadyInitialized system program cpi fails (for compressible accounts we create the token accounts via cpi)
78, // AlreadyInitialized (our program checks this after Assign+realloc pattern)
)
.await;
}
Expand Down
170 changes: 170 additions & 0 deletions program-tests/compressed-token-test/tests/ctoken/functional_ata.rs
Original file line number Diff line number Diff line change
Expand Up @@ -220,3 +220,173 @@ async fn test_create_ata_idempotent() {
)
.await;
}

/// Test: DoS prevention for ATA creation
/// 1. Derive ATA address
/// 2. Pre-fund the ATA address with lamports (simulating attacker donation)
/// 3. SUCCESS: Create ATA should succeed despite pre-funded lamports
#[tokio::test]
#[serial]
async fn test_create_ata_with_prefunded_lamports() {
let mut context = setup_account_test().await.unwrap();
let payer_pubkey = context.payer.pubkey();
let owner_pubkey = context.owner_keypair.pubkey();

// Derive ATA address
let (ata, bump) = derive_ctoken_ata(&owner_pubkey, &context.mint_pubkey);

// Pre-fund the ATA address with lamports (simulating attacker donation DoS attempt)
let prefund_amount = 1_000; // 1000 lamports
let transfer_ix = solana_sdk::system_instruction::transfer(&payer_pubkey, &ata, prefund_amount);

context
.rpc
.create_and_send_transaction(&[transfer_ix], &payer_pubkey, &[&context.payer])
.await
.unwrap();

// Verify the ATA address now has lamports
let ata_account = context.rpc.get_account(ata).await.unwrap();
assert!(
ata_account.is_some(),
"ATA address should exist with lamports"
);
assert_eq!(
ata_account.unwrap().lamports,
prefund_amount,
"ATA should have pre-funded lamports"
);

// Now create the ATA - this should succeed despite pre-funded lamports
let instruction = CreateAssociatedTokenAccount {
idempotent: false,
bump,
payer: payer_pubkey,
owner: owner_pubkey,
mint: context.mint_pubkey,
associated_token_account: ata,
compressible: None,
}
.instruction()
.unwrap();

context
.rpc
.create_and_send_transaction(&[instruction], &payer_pubkey, &[&context.payer])
.await
.unwrap();

// Verify ATA was created correctly
assert_create_associated_token_account(
&mut context.rpc,
owner_pubkey,
context.mint_pubkey,
None,
)
.await;

// Verify the ATA now has more lamports (rent-exempt + pre-funded)
let final_ata_account = context.rpc.get_account(ata).await.unwrap().unwrap();
assert!(
final_ata_account.lamports > prefund_amount,
"ATA should have rent-exempt balance plus pre-funded amount"
);
}

/// Test: DoS prevention for token account creation with custom rent payer
/// 1. Generate token account keypair
/// 2. Pre-fund the token account address with lamports (simulating attacker donation)
/// 3. SUCCESS: Create token account should succeed despite pre-funded lamports
#[tokio::test]
#[serial]
async fn test_create_token_account_with_prefunded_lamports() {
let mut context = setup_account_test().await.unwrap();
let payer_pubkey = context.payer.pubkey();
let token_account_pubkey = context.token_account_keypair.pubkey();

// Pre-fund the token account address with lamports (simulating attacker donation DoS attempt)
let prefund_amount = 1_000; // 1000 lamports
let transfer_ix = solana_sdk::system_instruction::transfer(
&payer_pubkey,
&token_account_pubkey,
prefund_amount,
);

context
.rpc
.create_and_send_transaction(&[transfer_ix], &payer_pubkey, &[&context.payer])
.await
.unwrap();

// Verify the token account address now has lamports
let token_account = context.rpc.get_account(token_account_pubkey).await.unwrap();
assert!(
token_account.is_some(),
"Token account address should exist with lamports"
);
assert_eq!(
token_account.unwrap().lamports,
prefund_amount,
"Token account should have pre-funded lamports"
);

// Now create the compressible token account - this should succeed despite pre-funded lamports
let compressible_params = CompressibleParams {
compressible_config: context.compressible_config,
rent_sponsor: context.rent_sponsor,
pre_pay_num_epochs: 0,
lamports_per_write: Some(100),
compress_to_account_pubkey: None,
token_account_version: light_ctoken_types::state::TokenDataVersion::ShaFlat,
};

let create_token_account_ix = CreateCTokenAccount::new(
payer_pubkey,
token_account_pubkey,
context.mint_pubkey,
context.owner_keypair.pubkey(),
)
.with_compressible(compressible_params)
.instruction()
.unwrap();

context
.rpc
.create_and_send_transaction(
&[create_token_account_ix],
&payer_pubkey,
&[&context.payer, &context.token_account_keypair],
)
.await
.unwrap();

// Verify token account was created correctly
assert_create_token_account(
&mut context.rpc,
token_account_pubkey,
context.mint_pubkey,
context.owner_keypair.pubkey(),
Some(CompressibleData {
compression_authority: context.compression_authority,
rent_sponsor: context.rent_sponsor,
num_prepaid_epochs: 0,
lamports_per_write: Some(100),
compress_to_pubkey: false,
account_version: light_ctoken_types::state::TokenDataVersion::ShaFlat,
payer: payer_pubkey,
}),
)
.await;

// Verify the token account now has more lamports (rent-exempt + pre-funded)
let final_token_account = context
.rpc
.get_account(token_account_pubkey)
.await
.unwrap()
.unwrap();
assert!(
final_token_account.lamports > prefund_amount,
"Token account should have rent-exempt balance plus pre-funded amount"
);
}
22 changes: 16 additions & 6 deletions program-tests/utils/src/assert_create_token_account.rs
Original file line number Diff line number Diff line change
Expand Up @@ -110,9 +110,14 @@ pub async fn assert_create_token_account_internal(
assert_eq!(actual_token_account, expected_token_account);

// Check if account existed before transaction (for idempotent mode)
let account_existed_before = rpc
.get_pre_transaction_account(&token_account_pubkey)
.is_some();
// Account "existed" only if it had data (was initialized), not just lamports
let pre_tx_account = rpc.get_pre_transaction_account(&token_account_pubkey);
let account_existed_before = pre_tx_account
.as_ref()
.map(|acc| !acc.data.is_empty())
.unwrap_or(false);
// Get pre-existing lamports (e.g., from attacker donation for DoS prevention test)
let pre_existing_lamports = pre_tx_account.map(|acc| acc.lamports).unwrap_or(0);

// Assert payer and rent sponsor balance changes
let payer_balance_before = rpc
Expand Down Expand Up @@ -183,12 +188,17 @@ pub async fn assert_create_token_account_internal(
payer_balance_before - payer_balance_after
);

// Rent sponsor pays: rent_exemption only
// Rent sponsor pays: rent_exemption minus any pre-existing lamports
// (pre-existing lamports from attacker donation are kept in the account)
let expected_rent_sponsor_payment =
rent_exemption.saturating_sub(pre_existing_lamports);
assert_eq!(
rent_sponsor_balance_before - rent_sponsor_balance_after,
expected_rent_sponsor_payment,
"Rent sponsor should have paid {} lamports (rent exemption {} - pre-existing {}), but paid {}",
expected_rent_sponsor_payment,
rent_exemption,
"Rent sponsor should have paid {} lamports (rent exemption only), but paid {}",
rent_exemption,
pre_existing_lamports,
rent_sponsor_balance_before - rent_sponsor_balance_after
);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -115,20 +115,19 @@ pub(crate) fn process_create_associated_token_account_inner<const IDEMPOTENT: bo
} else {
// Create the PDA account (with rent-exempt balance only)
let bump_seed = [bump];
let seeds = [
let ata_seeds = [
Seed::from(owner_bytes.as_ref()),
Seed::from(crate::LIGHT_CPI_SIGNER.program_id.as_ref()),
Seed::from(mint_bytes.as_ref()),
Seed::from(bump_seed.as_ref()),
];

let seeds_inputs = [seeds.as_slice()];

create_pda_account(
fee_payer,
associated_token_account,
token_account_size,
seeds_inputs,
None, // fee_payer is keypair
Some(ata_seeds.as_slice()), // ATA is PDA
None,
)?;
(None, None)
Expand Down Expand Up @@ -185,7 +184,7 @@ fn process_compressible_config<'info>(
compressible_config_ix_data.rent_payment as u64,
);

// Build ATA seeds
// Build ATA seeds (new_account is always a PDA)
let ata_bump_seed = [ata_bump];
let ata_seeds = [
Seed::from(owner_bytes.as_ref()),
Expand All @@ -194,36 +193,30 @@ fn process_compressible_config<'info>(
Seed::from(ata_bump_seed.as_ref()),
];

// Build rent sponsor seeds if needed (must be outside conditional for lifetime)
let rent_sponsor_bump;
let version_bytes;
let rent_sponsor_seeds;
// Build rent sponsor seeds if using rent sponsor PDA as fee_payer
let rent_sponsor_bump = [compressible_config_account.rent_sponsor_bump];
let version_bytes = compressible_config_account.version.to_le_bytes();
let rent_sponsor_seeds = [
Seed::from(b"rent_sponsor".as_ref()),
Seed::from(version_bytes.as_ref()),
Seed::from(rent_sponsor_bump.as_ref()),
];

// Create the PDA account (with rent-exempt balance only)
// rent_payer will be the rent_sponsor PDA for compressible accounts
let seeds_inputs: [&[Seed]; 2] = if custom_rent_payer {
// Only ATA seeds when custom rent payer
[ata_seeds.as_slice(), &[]]
// fee_payer_seeds: Some for rent_sponsor PDA, None for custom keypair
// new_account_seeds: Always Some (ATA is always a PDA)
let fee_payer_seeds = if custom_rent_payer {
None
} else {
// Both rent sponsor PDA seeds and ATA seeds
rent_sponsor_bump = [compressible_config_account.rent_sponsor_bump];
version_bytes = compressible_config_account.version.to_le_bytes();
rent_sponsor_seeds = [
Seed::from(b"rent_sponsor".as_ref()),
Seed::from(version_bytes.as_ref()),
Seed::from(rent_sponsor_bump.as_ref()),
];

[rent_sponsor_seeds.as_slice(), ata_seeds.as_slice()]
Some(rent_sponsor_seeds.as_slice())
};

let additional_lamports = if custom_rent_payer { Some(rent) } else { None };

create_pda_account(
rent_payer,
associated_token_account,
token_account_size,
seeds_inputs,
fee_payer_seeds,
Some(ata_seeds.as_slice()),
additional_lamports,
)?;

Expand Down
Loading