- Add token_program constraints for init-if-needed mint, token and ata

- Fix ATA `token_program` constraint parsing - Add `token_program` constraint test cases
coral-xyz · Apr 16, 2023 · b8faeca · b8faeca
1 parent 1a89f5a
commit b8faeca
Show file tree

Hide file tree

Showing 7 changed files with 1,319 additions and 23 deletions.
diff --git a/lang/syn/src/codegen/accounts/constraints.rs b/lang/syn/src/codegen/accounts/constraints.rs
@@ -547,7 +547,8 @@ fn generate_constraint_init_group(
                     // Checks that all the required accounts for this operation are present.
                     #optional_checks
 
-                    if !#if_needed || AsRef::<AccountInfo>::as_ref(&#field).owner == &anchor_lang::solana_program::system_program::ID {
+                    let owner_program = AsRef::<AccountInfo>::as_ref(&#field).owner;
+                    if !#if_needed || owner_program == &anchor_lang::solana_program::system_program::ID {
                         #payer_optional_check
 
                         // Create the account with the system program.
@@ -572,6 +573,9 @@ fn generate_constraint_init_group(
                         if pa.owner != #owner.key() {
                             return Err(anchor_lang::error::Error::from(anchor_lang::error::ErrorCode::ConstraintTokenOwner).with_account_name(#name_str).with_pubkeys((pa.owner, #owner.key())));
                         }
+                        if owner_program != &#token_program.key() {
+                            return Err(anchor_lang::error::Error::from(anchor_lang::error::ErrorCode::ConstraintTokenTokenProgram).with_account_name(#name_str).with_pubkeys((*owner_program, #token_program.key())));
+                        }
                     }
                     pa
                 };
@@ -614,7 +618,8 @@ fn generate_constraint_init_group(
                     // Checks that all the required accounts for this operation are present.
                     #optional_checks
 
-                    if !#if_needed || AsRef::<AccountInfo>::as_ref(&#field).owner == &anchor_lang::solana_program::system_program::ID {
+                    let owner_program = AsRef::<AccountInfo>::as_ref(&#field).owner;
+                    if !#if_needed || owner_program == &anchor_lang::solana_program::system_program::ID {
                         #payer_optional_check
 
                         let cpi_program = associated_token_program.to_account_info();
@@ -637,6 +642,9 @@ fn generate_constraint_init_group(
                         if pa.owner != #owner.key() {
                             return Err(anchor_lang::error::Error::from(anchor_lang::error::ErrorCode::ConstraintTokenOwner).with_account_name(#name_str).with_pubkeys((pa.owner, #owner.key())));
                         }
+                        if owner_program != &#token_program.key() {
+                            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()) {
                             return Err(anchor_lang::error::Error::from(anchor_lang::error::ErrorCode::AccountNotAssociatedTokenAccount).with_account_name(#name_str));
@@ -697,7 +705,8 @@ fn generate_constraint_init_group(
                     // Checks that all the required accounts for this operation are present.
                     #optional_checks
 
-                    if !#if_needed || AsRef::<AccountInfo>::as_ref(&#field).owner == &anchor_lang::solana_program::system_program::ID {
+                    let owner_program = AsRef::<AccountInfo>::as_ref(&#field).owner;
+                    if !#if_needed || owner_program == &anchor_lang::solana_program::system_program::ID {
                         // Define payer variable.
                         #payer_optional_check
 
@@ -726,6 +735,9 @@ fn generate_constraint_init_group(
                         if pa.decimals != #decimals {
                             return Err(anchor_lang::error::Error::from(anchor_lang::error::ErrorCode::ConstraintMintDecimals).with_account_name(#name_str).with_values((pa.decimals, #decimals)));
                         }
+                        if owner_program != &#token_program.key() {
+                            return Err(anchor_lang::error::Error::from(anchor_lang::error::ErrorCode::ConstraintMintTokenProgram).with_account_name(#name_str).with_pubkeys((*owner_program, #token_program.key())));
+                        }
                     }
                     pa
                 };
@@ -923,6 +935,7 @@ fn generate_constraint_associated_token(
     quote! {
         {
             #optional_checks
+            #token_program_check
 
             let my_owner = #name.owner;
             let wallet_address = #wallet_address.key();
@@ -934,8 +947,6 @@ fn generate_constraint_associated_token(
             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)));
             }
-
-            #token_program_check
         }
     }
 }

diff --git a/lang/syn/src/parser/accounts/constraints.rs b/lang/syn/src/parser/accounts/constraints.rs
@@ -648,6 +648,12 @@ impl<'ty> ConstraintGroupBuilder<'ty> {
                     auth.span(),
                     "mint must be provided to specify an associated token program derived address",
                 ))
+            },
+            (None, None, Some(token_program)) => {
+                return Err(ParseError::new(
+                    token_program.span(),
+                    "mint and authority must be provided to specify an associated token program derived address",
+                ))
             }
             _ => None,
         };

diff --git a/tests/misc/programs/misc-optional/src/context.rs b/tests/misc/programs/misc-optional/src/context.rs
@@ -47,6 +47,24 @@ pub struct TestInitAssociatedToken<'info> {
     pub associated_token_program: Option<Program<'info, AssociatedToken>>,
 }
 
+#[derive(Accounts)]
+pub struct TestInitAssociatedTokenWithTokenProgram<'info> {
+    #[account(init,
+        payer = payer,
+        associated_token::mint = mint,
+        associated_token::authority = payer,
+        associated_token::token_program = associated_token_token_program,
+    )]
+    pub token: Option<Account<'info, TokenAccount>>,
+    pub mint: Option<Account<'info, Mint>>,
+    #[account(mut)]
+    pub payer: Option<Signer<'info>>,
+    pub system_program: Option<Program<'info, System>>,
+    /// CHECK: ignore
+    pub associated_token_token_program: Option<AccountInfo<'info>>,
+    pub associated_token_program: Option<Program<'info, AssociatedToken>>,
+}
+
 #[derive(Accounts)]
 pub struct TestValidateAssociatedToken<'info> {
     #[account(
@@ -233,6 +251,23 @@ pub struct TestInitMint<'info> {
     pub token_program: Option<Program<'info, Token>>,
 }
 
+#[derive(Accounts)]
+pub struct TestInitMintWithTokenProgram<'info> {
+    #[account(init,
+        payer = payer,
+        mint::decimals = 6,
+        mint::authority = payer,
+        mint::freeze_authority = payer,
+        mint::token_program = mint_token_program,
+    )]
+    pub mint: Option<Account<'info, Mint>>,
+    #[account(mut)]
+    pub payer: Option<Signer<'info>>,
+    pub system_program: Option<Program<'info, System>>,
+    /// CHECK: ignore
+    pub mint_token_program: Option<AccountInfo<'info>>,
+}
+
 #[derive(Accounts)]
 pub struct TestInitToken<'info> {
     #[account(init, token::mint = mint, token::authority = payer, payer = payer, )]
@@ -244,6 +279,23 @@ pub struct TestInitToken<'info> {
     pub token_program: Option<Program<'info, Token>>,
 }
 
+#[derive(Accounts)]
+pub struct TestInitTokenWithTokenProgram<'info> {
+    #[account(init,
+        payer = payer,
+        token::mint = mint,
+        token::authority = payer,
+        token::token_program = token_token_program,
+    )]
+    pub token: Option<Account<'info, TokenAccount>>,
+    pub mint: Option<Account<'info, Mint>>,
+    #[account(mut)]
+    pub payer: Option<Signer<'info>>,
+    pub system_program: Option<Program<'info, System>>,
+    /// CHECK: ignore
+    pub token_token_program: Option<AccountInfo<'info>>,
+}
+
 #[derive(Accounts)]
 pub struct TestFetchAll<'info> {
     #[account(init, payer = authority, space = DataWithFilter::LEN + 8)]
@@ -326,6 +378,27 @@ pub struct TestInitMintIfNeeded<'info> {
     pub freeze_authority: Option<AccountInfo<'info>>,
 }
 
+#[derive(Accounts)]
+pub struct TestInitMintIfNeededWithTokenProgram<'info> {
+    #[account(init_if_needed,
+        payer = payer,
+        mint::decimals = 6,
+        mint::authority = mint_authority,
+        mint::freeze_authority = freeze_authority,
+        mint::token_program = mint_token_program,
+    )]
+    pub mint: Option<Account<'info, Mint>>,
+    #[account(mut)]
+    pub payer: Option<Signer<'info>>,
+    pub system_program: Option<Program<'info, System>>,
+    /// CHECK: ignore
+    pub mint_token_program: Option<AccountInfo<'info>>,
+    /// CHECK: ignore
+    pub mint_authority: Option<AccountInfo<'info>>,
+    /// CHECK: ignore
+    pub freeze_authority: Option<AccountInfo<'info>>,
+}
+
 #[derive(Accounts)]
 pub struct TestInitTokenIfNeeded<'info> {
     #[account(init_if_needed, token::mint = mint, token::authority = authority, payer = payer, )]
@@ -339,10 +412,28 @@ pub struct TestInitTokenIfNeeded<'info> {
     pub authority: Option<AccountInfo<'info>>,
 }
 
+#[derive(Accounts)]
+pub struct TestInitTokenIfNeededWithTokenProgram<'info> {
+    #[account(init_if_needed,
+        payer = payer,
+        token::mint = mint,
+        token::authority = authority,
+        token::token_program = token_token_program
+    )]
+    pub token: Option<Account<'info, TokenAccount>>,
+    pub mint: Option<Account<'info, Mint>>,
+    #[account(mut)]
+    pub payer: Option<Signer<'info>>,
+    pub system_program: Option<Program<'info, System>>,
+    /// CHECK: ignore
+    pub token_token_program: Option<AccountInfo<'info>>,
+    /// CHECK:
+    pub authority: Option<AccountInfo<'info>>,
+}
+
 #[derive(Accounts)]
 pub struct TestInitAssociatedTokenIfNeeded<'info> {
-    #[account(
-        init_if_needed,
+    #[account(init_if_needed,
         payer = payer,
         associated_token::mint = mint,
         associated_token::authority = authority
@@ -358,6 +449,26 @@ pub struct TestInitAssociatedTokenIfNeeded<'info> {
     pub authority: Option<AccountInfo<'info>>,
 }
 
+#[derive(Accounts)]
+pub struct TestInitAssociatedTokenIfNeededWithTokenProgram<'info> {
+    #[account(init_if_needed,
+        payer = payer,
+        associated_token::mint = mint,
+        associated_token::authority = authority,
+        associated_token::token_program = associated_token_token_program,
+    )]
+    pub token: Option<Account<'info, TokenAccount>>,
+    pub mint: Option<Account<'info, Mint>>,
+    #[account(mut)]
+    pub payer: Option<Signer<'info>>,
+    pub system_program: Option<Program<'info, System>>,
+    /// CHECK: ignore
+    pub associated_token_token_program: Option<AccountInfo<'info>>,
+    pub associated_token_program: Option<Program<'info, AssociatedToken>>,
+    /// CHECK: ignore
+    pub authority: Option<AccountInfo<'info>>,
+}
+
 #[derive(Accounts)]
 pub struct TestMultidimensionalArray<'info> {
     #[account(zero)]
@@ -492,6 +603,16 @@ pub struct TestOnlyAuthorityConstraint<'info> {
     pub mint: Option<Account<'info, Mint>>,
     pub payer: Option<Signer<'info>>,
 }
+
+#[derive(Accounts)]
+pub struct TestOnlyTokenProgramConstraint<'info> {
+    #[account(
+        token::token_program = token_token_program
+    )]
+    pub token: Option<Account<'info, TokenAccount>>,
+    pub token_token_program: Option<AccountInfo<'info>>,
+}
+
 #[derive(Accounts)]
 pub struct TestOnlyMintConstraint<'info> {
     #[account(
@@ -554,6 +675,16 @@ pub struct TestMintMissMintAuthConstraint<'info> {
     pub freeze_authority: Option<AccountInfo<'info>>,
 }
 
+#[derive(Accounts)]
+pub struct TestMintOnlyTokenProgramConstraint<'info> {
+    #[account(
+        mint::token_program = mint_token_program,
+    )]
+    pub mint: Option<Account<'info, Mint>>,
+    /// CHECK: ignore
+    pub mint_token_program: Option<AccountInfo<'info>>,
+}
+
 #[derive(Accounts)]
 pub struct TestAssociatedToken<'info> {
     #[account(
@@ -564,3 +695,18 @@ pub struct TestAssociatedToken<'info> {
     pub mint: Option<Account<'info, Mint>>,
     pub authority: Option<AccountInfo<'info>>,
 }
+
+#[derive(Accounts)]
+pub struct TestAssociatedTokenWithTokenProgramConstraint<'info> {
+    #[account(
+        associated_token::mint = mint,
+        associated_token::authority = authority,
+        associated_token::token_program = associated_token_token_program,
+    )]
+    pub token: Option<Account<'info, TokenAccount>>,
+    pub mint: Account<'info, Mint>,
+    /// CHECK: ignore
+    pub authority: AccountInfo<'info>,
+    /// CHECK: ignore
+    pub associated_token_token_program: Option<AccountInfo<'info>>,
+}
diff --git a/tests/misc/programs/misc-optional/src/lib.rs b/tests/misc/programs/misc-optional/src/lib.rs
@@ -171,13 +171,22 @@ pub mod misc_optional {
         Ok(())
     }
 
+    pub fn test_init_mint_with_token_program(_ctx: Context<TestInitMintWithTokenProgram>) -> Result<()> {
+        Ok(())
+    }
+
     pub fn test_init_token(ctx: Context<TestInitToken>) -> Result<()> {
         assert!(
             ctx.accounts.token.as_ref().unwrap().mint == ctx.accounts.mint.as_ref().unwrap().key()
         );
         Ok(())
     }
 
+    pub fn test_init_token_with_token_program(_ctx: Context<TestInitTokenWithTokenProgram>) -> Result<()> {
+        Ok(())
+    }
+
+
     pub fn test_composite_payer(ctx: Context<TestCompositePayer>) -> Result<()> {
         ctx.accounts.composite.data.as_mut().unwrap().data = 1;
         ctx.accounts.data.as_mut().unwrap().udata = 2;
@@ -192,6 +201,10 @@ pub mod misc_optional {
         Ok(())
     }
 
+    pub fn test_init_associated_token_with_token_program(ctx: Context<TestInitAssociatedTokenWithTokenProgram>) -> Result<()> {
+        Ok(())
+    }
+
     pub fn test_validate_associated_token(
         _ctx: Context<TestValidateAssociatedToken>,
     ) -> Result<()> {
@@ -238,16 +251,32 @@ pub mod misc_optional {
         Ok(())
     }
 
+    pub fn test_init_mint_if_needed_with_token_program(
+        _ctx: Context<TestInitMintIfNeededWithTokenProgram>,
+    ) -> Result<()> {
+        Ok(())
+    }
+
     pub fn test_init_token_if_needed(_ctx: Context<TestInitTokenIfNeeded>) -> Result<()> {
         Ok(())
     }
 
+    pub fn test_init_token_if_needed_with_token_program(_ctx: Context<TestInitTokenIfNeededWithTokenProgram>) -> Result<()> {
+        Ok(())
+    }
+
     pub fn test_init_associated_token_if_needed(
         _ctx: Context<TestInitAssociatedTokenIfNeeded>,
     ) -> Result<()> {
         Ok(())
     }
 
+    pub fn test_init_associated_token_if_needed_with_token_program(
+        _ctx: Context<TestInitAssociatedTokenIfNeededWithTokenProgram>,
+    ) -> Result<()> {
+        Ok(())
+    }
+
     pub fn init_with_space(_ctx: Context<InitWithSpace>, data: u16) -> Result<()> {
         Ok(())
     }
@@ -328,6 +357,10 @@ pub mod misc_optional {
         Ok(())
     }
 
+    pub fn test_only_token_program_constraint(_ctx: Context<TestOnlyTokenProgramConstraint>) -> Result<()> {
+        Ok(())
+    }
+
     pub fn test_mint_constraint(_ctx: Context<TestMintConstraint>, _decimals: u8) -> Result<()> {
         Ok(())
     }
@@ -358,7 +391,17 @@ pub mod misc_optional {
         Ok(())
     }
 
+    pub fn test_mint_only_token_program_constraint(
+        _ctx: Context<TestMintOnlyTokenProgramConstraint>,
+    ) -> Result<()> {
+        Ok(())
+    }
+
     pub fn test_associated_constraint(_ctx: Context<TestAssociatedToken>) -> Result<()> {
         Ok(())
     }
+
+    pub fn test_associated_token_with_token_program_constraint(_ctx: Context<TestAssociatedTokenWithTokenProgramConstraint>) -> Result<()> {
+        Ok(())
+    }
 }