associated-token-account: Fail idempotent for non-ata (#2903)

Author: joncinque

associated-token-account/program/src/processor.rs

@@ -72,6 +72,17 @@ fn process_create_associated_token_account(
     let spl_token_program_info = next_account_info(account_info_iter)?;
     let spl_token_program_id = spl_token_program_info.key;
 
+    let (associated_token_address, bump_seed) = get_associated_token_address_and_bump_seed_internal(
+        wallet_account_info.key,
+        spl_token_mint_info.key,
+        program_id,
+        spl_token_program_id,
+    );
+    if associated_token_address != *associated_token_account_info.key {
+        msg!("Error: Associated address does not match seed derivation");
+        return Err(ProgramError::InvalidSeeds);
+    }
+
     if create_mode == CreateMode::Idempotent
         && associated_token_account_info.owner == spl_token_program_id
     {
@@ -94,17 +105,6 @@ fn process_create_associated_token_account(
 
     let rent = Rent::get()?;
 
-    let (associated_token_address, bump_seed) = get_associated_token_address_and_bump_seed_internal(
-        wallet_account_info.key,
-        spl_token_mint_info.key,
-        program_id,
-        spl_token_program_id,
-    );
-    if associated_token_address != *associated_token_account_info.key {
-        msg!("Error: Associated address does not match seed derivation");
-        return Err(ProgramError::InvalidSeeds);
-    }
-
     let associated_token_account_signer_seeds: &[&[_]] = &[
         &wallet_account_info.key.to_bytes(),
         &spl_token_program_id.to_bytes(),

associated-token-account/program/tests/create_idempotent.rs

@@ -11,6 +11,8 @@ use {
         program_option::COption,
         program_pack::Pack,
         signature::Signer,
+        signer::keypair::Keypair,
+        system_instruction::create_account,
         transaction::{Transaction, TransactionError},
     },
     spl_associated_token_account::{
@@ -22,6 +24,7 @@ use {
     },
     spl_token::{
         extension::ExtensionType,
+        instruction::initialize_account,
         state::{Account, AccountState},
     },
 };
@@ -173,3 +176,64 @@ async fn fail_account_exists_with_wrong_owner() {
         )
     );
 }
+
+#[tokio::test]
+async fn fail_non_ata() {
+    let token_mint_address = Pubkey::new_unique();
+    let (mut banks_client, payer, recent_blockhash) =
+        program_test(token_mint_address, true).start().await;
+
+    let rent = banks_client.get_rent().await.unwrap();
+    let token_account_len = ExtensionType::get_account_len::<spl_token::state::Account>(&[
+        ExtensionType::ImmutableOwner,
+    ]);
+    let token_account_balance = rent.minimum_balance(token_account_len);
+
+    let wallet_address = Pubkey::new_unique();
+    let account = Keypair::new();
+    let transaction = Transaction::new_signed_with_payer(
+        &[
+            create_account(
+                &payer.pubkey(),
+                &account.pubkey(),
+                token_account_balance,
+                token_account_len as u64,
+                &spl_token::id(),
+            ),
+            initialize_account(
+                &spl_token::id(),
+                &account.pubkey(),
+                &token_mint_address,
+                &wallet_address,
+            )
+            .unwrap(),
+        ],
+        Some(&payer.pubkey()),
+        &[&payer, &account],
+        recent_blockhash,
+    );
+    banks_client.process_transaction(transaction).await.unwrap();
+
+    let mut instruction = create_associated_token_account_idempotent(
+        &payer.pubkey(),
+        &wallet_address,
+        &token_mint_address,
+        &spl_token::id(),
+    );
+    instruction.accounts[1] = AccountMeta::new(account.pubkey(), false); // <-- Invalid associated_account_address
+
+    let transaction = Transaction::new_signed_with_payer(
+        &[instruction],
+        Some(&payer.pubkey()),
+        &[&payer],
+        recent_blockhash,
+    );
+    assert_eq!(
+        banks_client
+            .process_transaction(transaction)
+            .await
+            .unwrap_err()
+            .unwrap(),
+        TransactionError::InstructionError(0, InstructionError::InvalidSeeds)
+    );
+}

associated-token-account/program/tests/process_create_associated_token_account.rs

@@ -213,7 +213,7 @@ async fn test_create_account_mismatch() {
             .await
             .unwrap_err()
             .unwrap(),
-        TransactionError::InstructionError(0, InstructionError::IllegalOwner)
+        TransactionError::InstructionError(0, InstructionError::InvalidSeeds)
     );
 
     let mut instruction = create_associated_token_account(