lang: `associated_token` constraints don't work when setting `token_program` (#2603)

Co-authored-by: acheron <acheroncrypto@gmail.com>

CHANGELOG.md

@@ -31,6 +31,7 @@ The minor version will be incremented upon a breaking change and the patch versi
 - client: Compile with Solana `1.14` ([#2572](https://github.com/coral-xyz/anchor/pull/2572)).
 - cli: Fix `anchor build --no-docs` adding docs to the IDL ([#2575](https://github.com/coral-xyz/anchor/pull/2575)).
 - ts: Load workspace programs on-demand rather than loading all of them at once ([#2579](https://github.com/coral-xyz/anchor/pull/2579)).
+- lang: Fix `associated_token::token_program` constraint ([#2603](https://github.com/coral-xyz/anchor/pull/2603)).
 
 ### Breaking
 

lang/syn/src/codegen/accounts/constraints.rs

@@ -646,7 +646,7 @@ fn generate_constraint_init_group(
                             return Err(anchor_lang::error::Error::from(anchor_lang::error::ErrorCode::ConstraintAssociatedTokenTokenProgram).with_account_name(#name_str).with_pubkeys((*owner_program, #token_program.key())));
                         }
 
-                        if pa.key() != ::anchor_spl::associated_token::get_associated_token_address(&#owner.key(), &#mint.key()) {
+                        if pa.key() != ::anchor_spl::associated_token::get_associated_token_address_with_program_id(&#owner.key(), &#mint.key(), &#token_program.key()) {
                             return Err(anchor_lang::error::Error::from(anchor_lang::error::ErrorCode::AccountNotAssociatedTokenAccount).with_account_name(#name_str));
                         }
                     }
@@ -913,6 +913,7 @@ fn generate_constraint_associated_token(
     let name_str = name.to_string();
     let wallet_address = &c.wallet;
     let spl_token_mint_address = &c.mint;
+
     let mut optional_check_scope = OptionalCheckScope::new_with_field(accs, name);
     let wallet_address_optional_check = optional_check_scope.generate_check(wallet_address);
     let spl_token_mint_address_optional_check =
@@ -921,6 +922,7 @@ fn generate_constraint_associated_token(
         #wallet_address_optional_check
         #spl_token_mint_address_optional_check
     };
+
     let token_program_check = match &c.token_program {
         Some(token_program) => {
             let token_program_optional_check = optional_check_scope.generate_check(token_program);
@@ -931,6 +933,14 @@ fn generate_constraint_associated_token(
         }
         None => quote! {},
     };
+    let get_associated_token_address = match &c.token_program {
+        Some(token_program) => quote! {
+            ::anchor_spl::associated_token::get_associated_token_address_with_program_id(&wallet_address, &#spl_token_mint_address.key(), &#token_program.key())
+        },
+        None => quote! {
+            ::anchor_spl::associated_token::get_associated_token_address(&wallet_address, &#spl_token_mint_address.key())
+        },
+    };
 
     quote! {
         {
@@ -942,7 +952,7 @@ fn generate_constraint_associated_token(
             if my_owner != wallet_address {
                 return Err(anchor_lang::error::Error::from(anchor_lang::error::ErrorCode::ConstraintTokenOwner).with_account_name(#name_str).with_pubkeys((my_owner, wallet_address)));
             }
-            let __associated_token_address = ::anchor_spl::associated_token::get_associated_token_address(&wallet_address, &#spl_token_mint_address.key());
+            let __associated_token_address = #get_associated_token_address;
             let my_key = #name.key();
             if my_key != __associated_token_address {
                 return Err(anchor_lang::error::Error::from(anchor_lang::error::ErrorCode::ConstraintAssociated).with_account_name(#name_str).with_pubkeys((my_key, __associated_token_address)));

tests/misc/programs/misc-optional/src/context.rs

@@ -2,6 +2,7 @@ use crate::account::*;
 use anchor_lang::prelude::*;
 use anchor_spl::associated_token::AssociatedToken;
 use anchor_spl::token::{Mint, Token, TokenAccount};
+use anchor_spl::token_interface::{Mint as MintInterface, TokenAccount as TokenAccountInterface};
 
 #[derive(Accounts)]
 pub struct TestTokenSeedsInit<'info> {
@@ -55,8 +56,8 @@ pub struct TestInitAssociatedTokenWithTokenProgram<'info> {
         associated_token::authority = payer,
         associated_token::token_program = associated_token_token_program,
     )]
-    pub token: Option<Account<'info, TokenAccount>>,
-    pub mint: Option<Account<'info, Mint>>,
+    pub token: Option<InterfaceAccount<'info, TokenAccountInterface>>,
+    pub mint: Option<InterfaceAccount<'info, MintInterface>>,
     #[account(mut)]
     pub payer: Option<Signer<'info>>,
     pub system_program: Option<Program<'info, System>>,
@@ -242,7 +243,7 @@ pub struct TestInitMintWithTokenProgram<'info> {
         mint::freeze_authority = payer,
         mint::token_program = mint_token_program,
     )]
-    pub mint: Option<Account<'info, Mint>>,
+    pub mint: Option<InterfaceAccount<'info, MintInterface>>,
     #[account(mut)]
     pub payer: Option<Signer<'info>>,
     pub system_program: Option<Program<'info, System>>,
@@ -439,8 +440,8 @@ pub struct TestInitAssociatedTokenIfNeededWithTokenProgram<'info> {
         associated_token::authority = authority,
         associated_token::token_program = associated_token_token_program,
     )]
-    pub token: Option<Account<'info, TokenAccount>>,
-    pub mint: Option<Account<'info, Mint>>,
+    pub token: Option<InterfaceAccount<'info, TokenAccountInterface>>,
+    pub mint: Option<InterfaceAccount<'info, MintInterface>>,
     #[account(mut)]
     pub payer: Option<Signer<'info>>,
     pub system_program: Option<Program<'info, System>>,
@@ -685,8 +686,8 @@ pub struct TestAssociatedTokenWithTokenProgramConstraint<'info> {
         associated_token::authority = authority,
         associated_token::token_program = associated_token_token_program,
     )]
-    pub token: Option<Account<'info, TokenAccount>>,
-    pub mint: Account<'info, Mint>,
+    pub token: Option<InterfaceAccount<'info, TokenAccountInterface>>,
+    pub mint: InterfaceAccount<'info, MintInterface>,
     /// CHECK: ignore
     pub authority: AccountInfo<'info>,
     /// CHECK: ignore

tests/misc/programs/misc/src/context.rs

@@ -2,6 +2,7 @@ use crate::account::*;
 use anchor_lang::prelude::*;
 use anchor_spl::associated_token::AssociatedToken;
 use anchor_spl::token::{Mint, Token, TokenAccount};
+use anchor_spl::token_interface::{Mint as MintInterface, TokenAccount as TokenAccountInterface};
 
 #[derive(Accounts)]
 pub struct TestTokenSeedsInit<'info> {
@@ -56,8 +57,8 @@ pub struct TestInitAssociatedTokenWithTokenProgram<'info> {
         associated_token::authority = payer,
         associated_token::token_program = associated_token_token_program,
     )]
-    pub token: Account<'info, TokenAccount>,
-    pub mint: Account<'info, Mint>,
+    pub token: InterfaceAccount<'info, TokenAccountInterface>,
+    pub mint: InterfaceAccount<'info, MintInterface>,
     #[account(mut)]
     pub payer: Signer<'info>,
     pub system_program: Program<'info, System>,
@@ -240,7 +241,7 @@ pub struct TestInitMintWithTokenProgram<'info> {
         mint::freeze_authority = payer,
         mint::token_program = mint_token_program,
     )]
-    pub mint: Account<'info, Mint>,
+    pub mint: InterfaceAccount<'info, MintInterface>,
     #[account(mut)]
     pub payer: Signer<'info>,
     pub system_program: Program<'info, System>,
@@ -445,8 +446,8 @@ pub struct TestInitAssociatedTokenIfNeededWithTokenProgram<'info> {
         associated_token::authority = authority,
         associated_token::token_program = associated_token_token_program,
     )]
-    pub token: Account<'info, TokenAccount>,
-    pub mint: Account<'info, Mint>,
+    pub token: InterfaceAccount<'info, TokenAccountInterface>,
+    pub mint: InterfaceAccount<'info, MintInterface>,
     #[account(mut)]
     pub payer: Signer<'info>,
     pub system_program: Program<'info, System>,
@@ -700,8 +701,8 @@ pub struct TestAssociatedTokenWithTokenProgramConstraint<'info> {
         associated_token::authority = authority,
         associated_token::token_program = associated_token_token_program,
     )]
-    pub token: Account<'info, TokenAccount>,
-    pub mint: Account<'info, Mint>,
+    pub token: InterfaceAccount<'info, TokenAccountInterface>,
+    pub mint: InterfaceAccount<'info, MintInterface>,
     /// CHECK: ignore
     pub authority: AccountInfo<'info>,
     /// CHECK: ignore

tests/misc/tests/misc/misc.ts

@@ -13,6 +13,8 @@ import {
   TOKEN_PROGRAM_ID,
   Token,
   ASSOCIATED_TOKEN_PROGRAM_ID,
+  AccountLayout,
+  MintLayout,
 } from "@solana/spl-token";
 import { assert, expect } from "chai";
 
@@ -23,6 +25,10 @@ const utf8 = anchor.utils.bytes.utf8;
 const nativeAssert = require("assert");
 const miscIdl = require("../../target/idl/misc.json");
 
+const TOKEN_2022_PROGRAM_ID = new anchor.web3.PublicKey(
+  "TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb"
+);
+
 const miscTest = (
   program: anchor.Program<Misc> | anchor.Program<MiscOptional>
 ) => {
@@ -830,28 +836,27 @@ const miscTest = (
           mint: newMint.publicKey,
           payer: provider.wallet.publicKey,
           systemProgram: anchor.web3.SystemProgram.programId,
-          mintTokenProgram: TOKEN_PROGRAM_ID,
+          mintTokenProgram: TOKEN_2022_PROGRAM_ID,
         },
         signers: [newMint],
       });
-      const client = new Token(
-        program.provider.connection,
-        newMint.publicKey,
-        TOKEN_PROGRAM_ID,
-        wallet.payer
+      const rawAccount = await provider.connection.getAccountInfo(
+        newMint.publicKey
       );
-      const mintAccount = await client.getMintInfo();
+      const mintAccount = MintLayout.decode(rawAccount.data);
       assert.strictEqual(mintAccount.decimals, 6);
-      assert.isTrue(
-        mintAccount.mintAuthority.equals(provider.wallet.publicKey)
+      assert.strictEqual(
+        new PublicKey(mintAccount.mintAuthority).toString(),
+        provider.wallet.publicKey.toString()
       );
-      assert.isTrue(
-        mintAccount.freezeAuthority.equals(provider.wallet.publicKey)
+      assert.strictEqual(
+        new PublicKey(mintAccount.freezeAuthority).toString(),
+        provider.wallet.publicKey.toString()
       );
-      const accInfo = await program.provider.connection.getAccountInfo(
-        newMint.publicKey
+      assert.strictEqual(
+        rawAccount.owner.toString(),
+        TOKEN_2022_PROGRAM_ID.toString()
       );
-      assert.strictEqual(accInfo.owner.toString(), TOKEN_PROGRAM_ID.toString());
     });
 
     it("Can create a random token account with token program", async () => {
@@ -867,22 +872,20 @@ const miscTest = (
         signers: [token],
       });
 
-      const client = new Token(
-        program.provider.connection,
-        mint.publicKey,
-        TOKEN_PROGRAM_ID,
-        wallet.payer
+      const rawAccount = await provider.connection.getAccountInfo(
+        token.publicKey
       );
-      const account = await client.getAccountInfo(token.publicKey);
-      // @ts-expect-error
-      assert.strictEqual(account.state, 1);
-      assert.strictEqual(account.amount.toNumber(), 0);
-      assert.isTrue(account.isInitialized);
+      const ataAccount = AccountLayout.decode(rawAccount.data);
+      assert.strictEqual(ataAccount.state, 1);
+      assert.strictEqual(new anchor.BN(ataAccount.amount).toNumber(), 0);
       assert.strictEqual(
-        account.owner.toString(),
+        new PublicKey(ataAccount.owner).toString(),
         provider.wallet.publicKey.toString()
       );
-      assert.strictEqual(account.mint.toString(), mint.publicKey.toString());
+      assert.strictEqual(
+        new PublicKey(ataAccount.mint).toString(),
+        mint.publicKey.toString()
+      );
     });
 
     describe("associated_token constraints", () => {
@@ -928,19 +931,19 @@ const miscTest = (
 
       it("Can create an associated token account with token program", async () => {
         const newMint = anchor.web3.Keypair.generate();
-        await program.rpc.testInitMint({
+        await program.rpc.testInitMintWithTokenProgram({
           accounts: {
             mint: newMint.publicKey,
             payer: provider.wallet.publicKey,
             systemProgram: anchor.web3.SystemProgram.programId,
-            tokenProgram: TOKEN_PROGRAM_ID,
+            mintTokenProgram: TOKEN_2022_PROGRAM_ID,
           },
           signers: [newMint],
         });
 
         const associatedToken = await Token.getAssociatedTokenAddress(
           ASSOCIATED_TOKEN_PROGRAM_ID,
-          TOKEN_PROGRAM_ID,
+          TOKEN_2022_PROGRAM_ID,
           newMint.publicKey,
           provider.wallet.publicKey
         );
@@ -951,36 +954,28 @@ const miscTest = (
             mint: newMint.publicKey,
             payer: provider.wallet.publicKey,
             systemProgram: anchor.web3.SystemProgram.programId,
-            associatedTokenTokenProgram: TOKEN_PROGRAM_ID,
+            associatedTokenTokenProgram: TOKEN_2022_PROGRAM_ID,
             associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
           },
         });
 
-        const token = new Token(
-          program.provider.connection,
-          newMint.publicKey,
-          TOKEN_PROGRAM_ID,
-          wallet.payer
+        const rawAta = await provider.connection.getAccountInfo(
+          associatedToken
         );
-        const ataAccount = await token.getAccountInfo(associatedToken);
-        // @ts-expect-error
+        const ataAccount = AccountLayout.decode(rawAta.data);
         assert.strictEqual(ataAccount.state, 1);
-        assert.strictEqual(ataAccount.amount.toNumber(), 0);
-        assert.isTrue(ataAccount.isInitialized);
+        assert.strictEqual(new anchor.BN(ataAccount.amount).toNumber(), 0);
         assert.strictEqual(
-          ataAccount.owner.toString(),
+          new PublicKey(ataAccount.owner).toString(),
           provider.wallet.publicKey.toString()
         );
         assert.strictEqual(
-          ataAccount.mint.toString(),
+          new PublicKey(ataAccount.mint).toString(),
           newMint.publicKey.toString()
         );
-        const rawAta = await provider.connection.getAccountInfo(
-          associatedToken
-        );
         assert.strictEqual(
           rawAta.owner.toBase58(),
-          TOKEN_PROGRAM_ID.toBase58()
+          TOKEN_2022_PROGRAM_ID.toBase58()
         );
       });
 
@@ -1073,31 +1068,31 @@ const miscTest = (
 
       it("associated_token constraints (no init) - Can make with associated_token::token_program", async () => {
         const mint = anchor.web3.Keypair.generate();
-        await program.rpc.testInitMint({
+        await program.rpc.testInitMintWithTokenProgram({
           accounts: {
             mint: mint.publicKey,
             payer: provider.wallet.publicKey,
             systemProgram: anchor.web3.SystemProgram.programId,
-            tokenProgram: TOKEN_PROGRAM_ID,
+            mintTokenProgram: TOKEN_2022_PROGRAM_ID,
           },
           signers: [mint],
         });
 
         const associatedToken = await Token.getAssociatedTokenAddress(
           ASSOCIATED_TOKEN_PROGRAM_ID,
-          TOKEN_PROGRAM_ID,
+          TOKEN_2022_PROGRAM_ID,
           mint.publicKey,
           provider.wallet.publicKey
         );
 
-        await program.rpc.testInitAssociatedToken({
+        await program.rpc.testInitAssociatedTokenWithTokenProgram({
           accounts: {
             token: associatedToken,
             mint: mint.publicKey,
             payer: provider.wallet.publicKey,
             systemProgram: anchor.web3.SystemProgram.programId,
-            tokenProgram: TOKEN_PROGRAM_ID,
             associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
+            associatedTokenTokenProgram: TOKEN_2022_PROGRAM_ID,
           },
           signers: [],
         });
@@ -1106,7 +1101,7 @@ const miscTest = (
             token: associatedToken,
             mint: mint.publicKey,
             authority: provider.wallet.publicKey,
-            associatedTokenTokenProgram: TOKEN_PROGRAM_ID,
+            associatedTokenTokenProgram: TOKEN_2022_PROGRAM_ID,
           },
         });
 
@@ -1115,7 +1110,7 @@ const miscTest = (
         );
         assert.strictEqual(
           account.owner.toString(),
-          TOKEN_PROGRAM_ID.toString()
+          TOKEN_2022_PROGRAM_ID.toString()
         );
       });
 
@@ -1531,25 +1526,21 @@ const miscTest = (
         },
         signers: [newToken],
       });
-      const mintClient = new Token(
-        provider.connection,
-        newMint.publicKey,
-        TOKEN_PROGRAM_ID,
-        wallet.payer
+
+      const rawAccount = await provider.connection.getAccountInfo(
+        newToken.publicKey
       );
-      const tokenAccount = await mintClient.getAccountInfo(newToken.publicKey);
-      assert.strictEqual(tokenAccount.amount.toNumber(), 0);
+      const ataAccount = AccountLayout.decode(rawAccount.data);
+      assert.strictEqual(new anchor.BN(ataAccount.amount).toNumber(), 0);
       assert.strictEqual(
-        tokenAccount.mint.toString(),
+        new PublicKey(ataAccount.mint).toString(),
         newMint.publicKey.toString()
       );
       assert.strictEqual(
-        tokenAccount.owner.toString(),
+        new PublicKey(ataAccount.owner).toString(),
         provider.wallet.publicKey.toString()
       );
-      const rawAccount = await provider.connection.getAccountInfo(
-        newToken.publicKey
-      );
+
       assert.strictEqual(
         rawAccount.owner.toString(),
         TOKEN_PROGRAM_ID.toString()
@@ -1614,19 +1605,19 @@ const miscTest = (
 
     it("init_if_needed creates associated token account if not exists with token program", async () => {
       const newMint = anchor.web3.Keypair.generate();
-      await program.rpc.testInitMint({
+      await program.rpc.testInitMintWithTokenProgram({
         accounts: {
           mint: newMint.publicKey,
           payer: provider.wallet.publicKey,
           systemProgram: anchor.web3.SystemProgram.programId,
-          tokenProgram: TOKEN_PROGRAM_ID,
+          mintTokenProgram: TOKEN_2022_PROGRAM_ID,
         },
         signers: [newMint],
       });
 
       const associatedToken = await Token.getAssociatedTokenAddress(
         ASSOCIATED_TOKEN_PROGRAM_ID,
-        TOKEN_PROGRAM_ID,
+        TOKEN_2022_PROGRAM_ID,
         newMint.publicKey,
         provider.wallet.publicKey
       );
@@ -1638,33 +1629,27 @@ const miscTest = (
           payer: provider.wallet.publicKey,
           systemProgram: anchor.web3.SystemProgram.programId,
           associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
-          associatedTokenTokenProgram: TOKEN_PROGRAM_ID,
+          associatedTokenTokenProgram: TOKEN_2022_PROGRAM_ID,
           authority: provider.wallet.publicKey,
         },
       });
 
-      const mintClient = new Token(
-        provider.connection,
-        newMint.publicKey,
-        TOKEN_PROGRAM_ID,
-        wallet.payer
+      const rawAccount = await provider.connection.getAccountInfo(
+        associatedToken
       );
-      const ataAccount = await mintClient.getAccountInfo(associatedToken);
-      assert.strictEqual(ataAccount.amount.toNumber(), 0);
+      const ataAccount = AccountLayout.decode(rawAccount.data);
+      assert.strictEqual(new anchor.BN(ataAccount.amount).toNumber(), 0);
       assert.strictEqual(
-        ataAccount.mint.toString(),
+        new PublicKey(ataAccount.mint).toString(),
         newMint.publicKey.toString()
       );
       assert.strictEqual(
-        ataAccount.owner.toString(),
+        new PublicKey(ataAccount.owner).toString(),
         provider.wallet.publicKey.toString()
       );
-      const rawAccount = await provider.connection.getAccountInfo(
-        associatedToken
-      );
       assert.strictEqual(
         rawAccount.owner.toString(),
-        TOKEN_PROGRAM_ID.toString()
+        TOKEN_2022_PROGRAM_ID.toString()
       );
     });
 
@@ -2284,7 +2269,7 @@ const miscTest = (
       }
     });
 
-    it("init_if_needed pass if associated token exists with token program", async () => {
+    it("init_if_needed pass if associated token exists", async () => {
       const mint = anchor.web3.Keypair.generate();
       await program.rpc.testInitMint({
         accounts: {
@@ -2314,13 +2299,56 @@ const miscTest = (
         },
       });
 
+      await program.rpc.testInitAssociatedTokenIfNeeded({
+        accounts: {
+          token: associatedToken,
+          mint: mint.publicKey,
+          payer: provider.wallet.publicKey,
+          systemProgram: anchor.web3.SystemProgram.programId,
+          tokenProgram: TOKEN_PROGRAM_ID,
+          associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
+          authority: provider.wallet.publicKey,
+        },
+      });
+    });
+
+    it("init_if_needed pass if associated token exists with token program", async () => {
+      const mint = anchor.web3.Keypair.generate();
+      await program.rpc.testInitMintWithTokenProgram({
+        accounts: {
+          mint: mint.publicKey,
+          payer: provider.wallet.publicKey,
+          systemProgram: anchor.web3.SystemProgram.programId,
+          mintTokenProgram: TOKEN_2022_PROGRAM_ID,
+        },
+        signers: [mint],
+      });
+
+      const associatedToken = await Token.getAssociatedTokenAddress(
+        ASSOCIATED_TOKEN_PROGRAM_ID,
+        TOKEN_2022_PROGRAM_ID,
+        mint.publicKey,
+        provider.wallet.publicKey
+      );
+
+      await program.rpc.testInitAssociatedTokenWithTokenProgram({
+        accounts: {
+          token: associatedToken,
+          mint: mint.publicKey,
+          payer: provider.wallet.publicKey,
+          systemProgram: anchor.web3.SystemProgram.programId,
+          associatedTokenTokenProgram: TOKEN_2022_PROGRAM_ID,
+          associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
+        },
+      });
+
       await program.rpc.testInitAssociatedTokenIfNeededWithTokenProgram({
         accounts: {
           token: associatedToken,
           mint: mint.publicKey,
           payer: provider.wallet.publicKey,
           systemProgram: anchor.web3.SystemProgram.programId,
-          associatedTokenTokenProgram: TOKEN_PROGRAM_ID,
+          associatedTokenTokenProgram: TOKEN_2022_PROGRAM_ID,
           associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
           authority: provider.wallet.publicKey,
         },