lang, ts: fix init_if_needed missing ATA address check (#1221)

coral-xyz · Jan 2, 2022 · b3720a0 · b3720a0
1 parent 4013ec9
commit b3720a0
Show file tree

Hide file tree

Showing 6 changed files with 79 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -14,7 +14,8 @@ incremented for features.
 ### Fixes
 
 * ts: Change commitment message `recent` to `processed` and `max` to `finalized` ([#1128](https://github.com/project-serum/anchor/pull/1128))
-* ts: fix `translateAddress` which currently leads to failing browser code. Now uses `PublicKey` constructor instead of prototype chain constructor name checking which doesn't work in the presence of code minifying/mangling([1138](https://github.com/project-serum/anchor/pull/1138))
+* ts: fix `translateAddress` which currently leads to failing browser code. Now uses `PublicKey` constructor instead of prototype chain constructor name checking which doesn't work in the presence of code minifying/mangling([#1138](https://github.com/project-serum/anchor/pull/1138))
+* lang: add missing check that verifies that account is ATA when using `init_if_needed` and init is not needed([#1221](https://github.com/project-serum/anchor/pull/1221))
 
 ### Features
 

diff --git a/cli/src/config.rs b/cli/src/config.rs
@@ -422,7 +422,7 @@ impl FromStr for Config {
                 cluster: cfg.provider.cluster.parse()?,
                 wallet: shellexpand::tilde(&cfg.provider.wallet).parse()?,
             },
-            scripts: cfg.scripts.unwrap_or_else(BTreeMap::new),
+            scripts: cfg.scripts.unwrap_or_default(),
             test: cfg.test,
             programs: cfg.programs.map_or(Ok(BTreeMap::new()), deser_programs)?,
             workspace: cfg.workspace.unwrap_or_default(),

diff --git a/lang/src/error.rs b/lang/src/error.rs
@@ -142,6 +142,9 @@ pub enum ErrorCode {
     /// 3013 - The given account is not a program data account
     #[msg("The given account is not a program data account")]
     AccountNotProgramData,
+    /// 3014 - The given account is not the associated token account
+    #[msg("The given account is not the associated token account")]
+    AccountNotAssociatedTokenAccount,
 
     // State.
     /// 4000 - The given state account does not have the correct address

diff --git a/lang/syn/src/codegen/accounts/constraints.rs b/lang/syn/src/codegen/accounts/constraints.rs
@@ -488,6 +488,10 @@ pub fn generate_init(
                         if pa.owner != #owner.key() {
                             return Err(anchor_lang::__private::ErrorCode::ConstraintTokenOwner.into());
                         }
+
+                        if pa.key() != anchor_spl::associated_token::get_associated_token_address(&#owner.key(), &#mint.key()) {
+                            return Err(anchor_lang::__private::ErrorCode::AccountNotAssociatedTokenAccount.into());
+                        }
                     }
                     pa
                 };

diff --git a/tests/misc/tests/misc.js b/tests/misc/tests/misc.js
@@ -1282,6 +1282,70 @@ describe("misc", () => {
     }
   });
 
+  it("init_if_needed throws if token exists with correct owner and mint but is not the ATA", async () => {
+    const mint = anchor.web3.Keypair.generate();
+    await program.rpc.testInitMint({
+      accounts: {
+        mint: mint.publicKey,
+        payer: program.provider.wallet.publicKey,
+        systemProgram: anchor.web3.SystemProgram.programId,
+        tokenProgram: TOKEN_PROGRAM_ID,
+        rent: anchor.web3.SYSVAR_RENT_PUBKEY,
+      },
+      signers: [mint],
+    });
+
+    const associatedToken = await Token.getAssociatedTokenAddress(
+      ASSOCIATED_TOKEN_PROGRAM_ID,
+      TOKEN_PROGRAM_ID,
+      mint.publicKey,
+      program.provider.wallet.publicKey
+    );
+
+    await program.rpc.testInitAssociatedToken({
+      accounts: {
+        token: associatedToken,
+        mint: mint.publicKey,
+        payer: program.provider.wallet.publicKey,
+        rent: anchor.web3.SYSVAR_RENT_PUBKEY,
+        systemProgram: anchor.web3.SystemProgram.programId,
+        tokenProgram: TOKEN_PROGRAM_ID,
+        associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
+      },
+    });
+
+    const token = anchor.web3.Keypair.generate();
+    await program.rpc.testInitToken({
+      accounts: {
+        token: token.publicKey,
+        mint: mint.publicKey,
+        payer: program.provider.wallet.publicKey,
+        systemProgram: anchor.web3.SystemProgram.programId,
+        tokenProgram: TOKEN_PROGRAM_ID,
+        rent: anchor.web3.SYSVAR_RENT_PUBKEY,
+      },
+      signers: [token],
+    });
+
+    try {
+      await program.rpc.testInitAssociatedTokenIfNeeded({
+        accounts: {
+          token: token.publicKey,
+          mint: mint.publicKey,
+          payer: program.provider.wallet.publicKey,
+          rent: anchor.web3.SYSVAR_RENT_PUBKEY,
+          systemProgram: anchor.web3.SystemProgram.programId,
+          tokenProgram: TOKEN_PROGRAM_ID,
+          associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
+          authority: program.provider.wallet.publicKey,
+        },
+      });
+      assert.ok(false);
+    } catch (err) {
+      assert.equal(err.code, 3014);
+    }
+  });
+
   it("Can use multidimensional array", async () => {
     const array2d = new Array(10).fill(new Array(10).fill(99));
     const data = anchor.web3.Keypair.generate();

diff --git a/ts/src/error.ts b/ts/src/error.ts
@@ -98,6 +98,7 @@ const LangErrorCode = {
   AccountNotSystemOwned: 3011,
   AccountNotInitialized: 3012,
   AccountNotProgramData: 3013,
+  AccountNotAssociatedTokenAccount: 3014,
   // State.
   StateInvalidAddress: 4000,
 
@@ -207,6 +208,10 @@ const LangErrorMessage = new Map([
     LangErrorCode.AccountNotProgramData,
     "The given account is not a program data account",
   ],
+  [
+    LangErrorCode.AccountNotAssociatedTokenAccount,
+    "The given account is not the associated token account",
+  ],
 
   // State.
   [