feat: improve withdraw spl conditions when ata doesnt exist (#132)

Author: skosito

programs/gateway/src/errors.rs

@@ -23,4 +23,6 @@ pub enum Errors {
     EmptyReceiver,
     #[msg("InvalidInstructionData")]
     InvalidInstructionData,
+    #[msg("InvalidAtaOwner")]
+    InvalidAtaOwner,
 }

programs/gateway/src/instructions/withdraw.rs

@@ -1,5 +1,6 @@
 use crate::{
     contexts::{Withdraw, WithdrawSPLToken},
+    errors::Errors,
     state::InstructionId,
     utils::{validate_message, verify_ata_match, DEFAULT_GAS_COST},
 };
@@ -89,13 +90,19 @@ pub fn handle_spl(
         &ctx.accounts.recipient_ata.key(),
     )?;
 
+    // Check if account is either empty owned by system program or owner is token program
+    let recipient_ata_account = ctx.accounts.recipient_ata.to_account_info();
+    require!(
+        *recipient_ata_account.owner == anchor_spl::token::ID
+            || (*recipient_ata_account.owner == anchor_lang::system_program::ID
+                && recipient_ata_account.lamports() == 0),
+        Errors::InvalidAtaOwner
+    );
+
     // 3. Create recipient ATA if needed and calculate costs
     let mut cost_ata_create: u64 = 0;
-    let recipient_ata_account = ctx.accounts.recipient_ata.to_account_info();
 
-    if recipient_ata_account.lamports() == 0
-        || *recipient_ata_account.owner == ctx.accounts.system_program.key()
-    {
+    if recipient_ata_account.lamports() == 0 {
         // ATA needs to be created
         msg!(
             "Creating associated token account {:?} for recipient {:?}...",

tests/gateway.ts

@@ -3484,6 +3484,54 @@ describe("Gateway", () => {
     expect(rentPayerPdaBal0 - rentPayerPdaBal1).to.be.eq(to_ata_bal + 5000); // rentPayer pays rent
   });
 
+  it("Withdrawal when a system account exists at the ATA address fails", async () => {
+    // new ata
+    const recipient = anchor.web3.Keypair.generate();
+    const maliciousAta = await spl.getAssociatedTokenAddress(
+      mint.publicKey,
+      recipient.publicKey,
+      false
+    );
+
+    // create system owned account using ata address
+    const lamports = await conn.getMinimumBalanceForRentExemption(8);
+    const maliciousTx = new anchor.web3.Transaction().add(
+      anchor.web3.SystemProgram.transfer({
+        fromPubkey: wallet.publicKey,
+        toPubkey: maliciousAta,
+        lamports,
+      })
+    );
+    await anchor.web3.sendAndConfirmTransaction(conn, maliciousTx, [wallet]);
+
+    // attempt withdrawal — should fail
+    const pdaAta = await spl.getAssociatedTokenAddress(
+      mint.publicKey,
+      pdaAccount,
+      true
+    );
+    const pdaAccountData = await gatewayProgram.account.pda.fetch(pdaAccount);
+    const amount = new anchor.BN(500_000);
+    const nonce = pdaAccountData.nonce;
+
+    try {
+      await withdrawSplToken(
+        mint,
+        usdcDecimals,
+        amount,
+        nonce,
+        pdaAta,
+        maliciousAta,
+        recipient.publicKey,
+        gatewayProgram
+      );
+      throw new Error("Expected error not thrown");
+    } catch (err) {
+      expect(err).to.be.instanceof(anchor.AnchorError);
+      expect(err.message).to.include("InvalidAtaOwner");
+    }
+  });
+
   it("Withdraw SPL token with wrong nonce should fail", async () => {
     let pda_ata = await spl.getAssociatedTokenAddress(
       mint.publicKey,