fix: robust account creation (#2097)

* first version

* Update programs/compressed-token/program/src/shared/create_pda_account.rs

Co-authored-by: Swen Schäferjohann <swen@lightprotocol.com>

* Update programs/compressed-token/program/src/shared/create_pda_account.rs

Co-authored-by: Swen Schäferjohann <swen@lightprotocol.com>

* Update programs/compressed-token/program/src/shared/create_pda_account.rs

Co-authored-by: Swen Schäferjohann <swen@lightprotocol.com>

---------

Co-authored-by: Swen Schäferjohann <swen@lightprotocol.com>

Author: ananas-block

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;
     }

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"
+    );
+}

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
                 );
             }

programs/compressed-token/program/src/create_associated_token_account.rs

@@ -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)
@@ -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()),
@@ -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,
     )?;