diff --git a/program-tests/compressed-token-test/tests/ctoken/create.rs b/program-tests/compressed-token-test/tests/ctoken/create.rs index 97cff677bb..a8570c56b6 100644 --- a/program-tests/compressed-token-test/tests/ctoken/create.rs +++ b/program-tests/compressed-token-test/tests/ctoken/create.rs @@ -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; } diff --git a/program-tests/compressed-token-test/tests/ctoken/functional_ata.rs b/program-tests/compressed-token-test/tests/ctoken/functional_ata.rs index 895183a5d1..9c7a1fd517 100644 --- a/program-tests/compressed-token-test/tests/ctoken/functional_ata.rs +++ b/program-tests/compressed-token-test/tests/ctoken/functional_ata.rs @@ -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" + ); +} diff --git a/program-tests/utils/src/assert_create_token_account.rs b/program-tests/utils/src/assert_create_token_account.rs index c4585f4a2c..7f3289d9e3 100644 --- a/program-tests/utils/src/assert_create_token_account.rs +++ b/program-tests/utils/src/assert_create_token_account.rs @@ -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 @@ -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 ); } diff --git a/programs/compressed-token/program/src/create_associated_token_account.rs b/programs/compressed-token/program/src/create_associated_token_account.rs index dc19e6c5cd..00b0024ccf 100644 --- a/programs/compressed-token/program/src/create_associated_token_account.rs +++ b/programs/compressed-token/program/src/create_associated_token_account.rs @@ -115,20 +115,19 @@ pub(crate) fn process_create_associated_token_account_inner( 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()), @@ -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, )?; diff --git a/programs/compressed-token/program/src/create_token_account.rs b/programs/compressed-token/program/src/create_token_account.rs index a86e00f38c..cfd9a7f8cd 100644 --- a/programs/compressed-token/program/src/create_token_account.rs +++ b/programs/compressed-token/program/src/create_token_account.rs @@ -11,12 +11,7 @@ use light_ctoken_types::{ COMPRESSIBLE_TOKEN_ACCOUNT_SIZE, }; use light_program_profiler::profile; -use pinocchio::{ - account_info::AccountInfo, - instruction::Seed, - sysvars::{rent::Rent, Sysvar}, -}; -use pinocchio_system::instructions::CreateAccount; +use pinocchio::{account_info::AccountInfo, instruction::Seed}; use spl_pod::{bytemuck, solana_msg::msg}; use crate::shared::{ @@ -188,42 +183,41 @@ pub fn process_create_token_account( let custom_rent_payer = *compressible.rent_payer.key() != config_account.rent_sponsor.to_bytes(); - if custom_rent_payer { - // custom rent payer for account creation -> pays rent exemption - // rent payer must be signer. - create_account_with_custom_rent_payer( - compressible.rent_payer, - accounts.token_account, - account_size, - rent, - ) - .map_err(convert_program_error)?; - (Some(*config_account), Some(*compressible.rent_payer.key())) + // Build fee_payer seeds (rent_sponsor PDA or None for custom keypair) + let version_bytes = config_account.version.to_le_bytes(); + let bump_seed = [config_account.rent_sponsor_bump]; + let rent_sponsor_seeds = [ + Seed::from(b"rent_sponsor".as_ref()), + Seed::from(version_bytes.as_ref()), + Seed::from(bump_seed.as_ref()), + ]; + + // fee_payer_seeds: Some for rent_sponsor PDA, None for custom keypair + // new_account_seeds: None (token_account is always a keypair signer) + let fee_payer_seeds = if custom_rent_payer { + None } else { - // Rent recipient is fee payer for account creation -> pays rent exemption - let version_bytes = config_account.version.to_le_bytes(); - let bump_seed = [config_account.rent_sponsor_bump]; - let seeds = [ - Seed::from(b"rent_sponsor".as_ref()), - Seed::from(version_bytes.as_ref()), - Seed::from(bump_seed.as_ref()), - ]; - - let seeds_inputs = [seeds.as_slice()]; + Some(rent_sponsor_seeds.as_slice()) + }; - // PDA creates account with only rent-exempt balance - create_pda_account( - compressible.rent_payer, - accounts.token_account, - account_size, - seeds_inputs, - None, - )?; + // Create token account (handles DoS prevention internally) + create_pda_account( + compressible.rent_payer, + accounts.token_account, + account_size, + fee_payer_seeds, + None, // token_account is keypair signer + None, // no additional lamports here + )?; + + // Payer transfers the additional rent (compression incentive) + transfer_lamports_via_cpi(rent, compressible.payer, accounts.token_account) + .map_err(convert_program_error)?; - // Payer transfers the additional rent (compression incentive) - transfer_lamports_via_cpi(rent, compressible.payer, accounts.token_account) - .map_err(convert_program_error)?; + if custom_rent_payer { + (Some(*config_account), Some(*compressible.rent_payer.key())) + } else { (Some(*config_account), None) } } else { @@ -240,24 +234,3 @@ pub fn process_create_token_account( custom_rent_payer, ) } - -#[profile] -#[inline(always)] -fn create_account_with_custom_rent_payer( - rent_payer: &AccountInfo, - token_account: &AccountInfo, - account_size: usize, - rent: u64, -) -> pinocchio::ProgramResult { - let solana_rent = Rent::get()?; - let lamports = solana_rent.minimum_balance(account_size) + rent; - - let create_account = CreateAccount { - from: rent_payer, - to: token_account, - lamports, - space: account_size as u64, - owner: &crate::LIGHT_CPI_SIGNER.program_id, - }; - create_account.invoke() -} diff --git a/programs/compressed-token/program/src/mint_action/actions/create_spl_mint/create_mint_account.rs b/programs/compressed-token/program/src/mint_action/actions/create_spl_mint/create_mint_account.rs index c9217f9a9f..cc231777eb 100644 --- a/programs/compressed-token/program/src/mint_action/actions/create_spl_mint/create_mint_account.rs +++ b/programs/compressed-token/program/src/mint_action/actions/create_spl_mint/create_mint_account.rs @@ -33,19 +33,18 @@ pub fn create_mint_account( // Create account using shared function let bump_seed = [mint_bump]; - let seeds = [ + let mint_seeds = [ Seed::from(COMPRESSED_MINT_SEED), Seed::from(mint_signer.key().as_ref()), Seed::from(bump_seed.as_ref()), ]; - let seeds_inputs = [seeds.as_slice()]; - create_pda_account( executing_accounts.system.fee_payer, mint_account, mint_account_size, - seeds_inputs, + None, // fee_payer is keypair + Some(mint_seeds.as_slice()), // mint is PDA None, ) } diff --git a/programs/compressed-token/program/src/mint_action/actions/create_spl_mint/create_token_pool.rs b/programs/compressed-token/program/src/mint_action/actions/create_spl_mint/create_token_pool.rs index 34045eed54..69dcaad55c 100644 --- a/programs/compressed-token/program/src/mint_action/actions/create_spl_mint/create_token_pool.rs +++ b/programs/compressed-token/program/src/mint_action/actions/create_spl_mint/create_token_pool.rs @@ -44,19 +44,18 @@ pub fn create_token_pool_account_manual( // Create account using shared function let bump_seed = [token_pool_bump]; - let seeds = [ + let pool_seeds = [ Seed::from(POOL_SEED), Seed::from(mint_key.as_ref()), Seed::from(bump_seed.as_ref()), ]; - let seeds_inputs = [seeds.as_slice()]; - create_pda_account( executing_accounts.system.fee_payer, token_pool_pda, token_account_size, - seeds_inputs, + None, // fee_payer is keypair + Some(pool_seeds.as_slice()), // token_pool is PDA None, ) } diff --git a/programs/compressed-token/program/src/shared/create_pda_account.rs b/programs/compressed-token/program/src/shared/create_pda_account.rs index bdecfc9a2d..f7eeb6ab12 100644 --- a/programs/compressed-token/program/src/shared/create_pda_account.rs +++ b/programs/compressed-token/program/src/shared/create_pda_account.rs @@ -6,67 +6,88 @@ use pinocchio::{ pubkey::Pubkey, sysvars::{rent::Rent, Sysvar}, }; -use pinocchio_system::instructions::CreateAccount; +use pinocchio_system::instructions::{Assign, CreateAccount, Transfer}; use crate::{shared::convert_program_error, LIGHT_CPI_SIGNER}; -// /// Configuration for creating a PDA account -// #[derive(Debug)] -// pub struct CreatePdaSeeds<'a> { -// /// The seeds used to derive the PDA (without bump) -// pub seeds: &'a [&'a [u8]], -// /// The bump seed for PDA derivation -// pub bump: u8, -// } - -/// Creates a PDA account with the specified configuration(s). -/// -/// This function abstracts the common PDA account creation pattern used across -/// create_associated_token_account, create_mint_account, and create_token_pool. -/// -/// ## Process -/// 1. Calculates rent based on account size -/// 2. Builds seed arrays with bumps for each config -/// 3. Creates account via system program with specified owner -/// 4. Signs transaction with derived PDA seeds +/// Creates an account with explicit seed parameters for fee_payer and new_account. /// /// ## Parameters -/// - `configs`: ArrayVec of PDA configs. First config is for the new account being created. -/// Additional configs are for fee payer PDAs that need to sign. +/// - `fee_payer`: Account paying for rent (keypair or PDA like rent_sponsor) +/// - `new_account`: Account being created (keypair or PDA like ATA) +/// - `account_size`: Size in bytes for the new account +/// - `fee_payer_seeds`: PDA seeds if fee_payer is a PDA (e.g., rent_sponsor), None if keypair +/// - `new_account_seeds`: PDA seeds if new_account is a PDA (e.g., ATA), None if keypair +/// - `additional_lamports`: Extra lamports beyond rent-exempt minimum (e.g., compression cost) +/// +/// ## Cold Path +/// If new_account already has lamports (e.g., attacker donation), uses +/// Assign + realloc + Transfer pattern instead of CreateAccount which would fail. #[profile] -pub fn create_pda_account( +pub fn create_pda_account( fee_payer: &AccountInfo, new_account: &AccountInfo, account_size: usize, - seeds_inputs: [&[Seed]; N], + fee_payer_seeds: Option<&[Seed]>, + new_account_seeds: Option<&[Seed]>, additional_lamports: Option, ) -> Result<(), ProgramError> { - // Ensure we have at least one config - if seeds_inputs.is_empty() { - return Err(ProgramError::InvalidInstructionData); - } // Calculate rent let rent = Rent::get().map_err(|_| ProgramError::UnsupportedSysvar)?; let lamports = rent.minimum_balance(account_size) + additional_lamports.unwrap_or_default(); - let create_account = CreateAccount { + // Build signers from seeds + let fee_payer_signer: Option = fee_payer_seeds.map(Signer::from); + let new_account_signer: Option = new_account_seeds.map(Signer::from); + + // Cold Path: if account already has lamports (e.g., from attacker donation), + // use Assign + realloc + Transfer instead of CreateAccount which would fail. + if new_account.lamports() > 0 { + let current_lamports = new_account.lamports(); + + Assign { + account: new_account, + owner: &LIGHT_CPI_SIGNER.program_id, + } + .invoke_signed(new_account_signer.as_slice()) + .map_err(convert_program_error)?; + + new_account + .resize(account_size) + .map_err(convert_program_error)?; + + // Transfer remaining lamports for rent-exemption if needed + if lamports > current_lamports { + Transfer { + from: fee_payer, + to: new_account, + lamports: lamports - current_lamports, + } + .invoke_signed(fee_payer_signer.as_slice()) + .map_err(convert_program_error)?; + } + + return Ok(()); + } + + // Normal path: CreateAccount (requires both to sign) + let mut signers = arrayvec::ArrayVec::::new(); + if let Some(s) = fee_payer_signer { + signers.push(s); + } + if let Some(s) = new_account_signer { + signers.push(s); + } + + CreateAccount { from: fee_payer, to: new_account, lamports, space: account_size as u64, owner: &LIGHT_CPI_SIGNER.program_id, - }; - - let mut signers = arrayvec::ArrayVec::::new(); - for seeds in seeds_inputs.iter() { - if !seeds.is_empty() { - signers.push(Signer::from(*seeds)); - } } - - create_account - .invoke_signed(signers.as_slice()) - .map_err(convert_program_error) + .invoke_signed(signers.as_slice()) + .map_err(convert_program_error) } /// Verifies that the provided account matches the expected PDA