transfer-hook-example: Only allow one mint to initialize (#6812)

* transfer-hook-example: Only allow one mint

* Add a crate feature to fix downstream tests easily

token/js/package.json

@@ -33,7 +33,7 @@
         "clean": "shx rm -rf lib **/*.tsbuildinfo || true",
         "build": "tsc --build --verbose tsconfig.all.json",
         "postbuild": "shx echo '{ \"type\": \"commonjs\" }' > lib/cjs/package.json",
-        "build:program": "cargo build-sbf --manifest-path=../program/Cargo.toml && cargo build-sbf --manifest-path=../program-2022/Cargo.toml && cargo build-sbf --manifest-path=../../associated-token-account/program/Cargo.toml && cargo build-sbf --manifest-path=../transfer-hook/example/Cargo.toml",
+        "build:program": "cargo build-sbf --manifest-path=../program/Cargo.toml && cargo build-sbf --manifest-path=../program-2022/Cargo.toml && cargo build-sbf --manifest-path=../../associated-token-account/program/Cargo.toml && cargo build-sbf --no-default-features --manifest-path=../transfer-hook/example/Cargo.toml",
         "watch": "tsc --build --verbose --watch tsconfig.all.json",
         "release": "npm run clean && npm run build",
         "fmt": "prettier --write '{*,**/*}.{ts,tsx,js,jsx,json}'",

token/transfer-hook/cli/Cargo.toml

@@ -30,6 +30,7 @@ serde_yaml = "0.9.34"
 solana-test-validator = ">=1.18.11,<=2"
 spl-token-2022 = { version = "3.0.2", path = "../../program-2022", features = ["no-entrypoint"] }
 spl-token-client = { version = "0.10.0", path = "../../client" }
+spl-transfer-hook-example = { version = "0.6.0", path = "../example" }
 
 [[bin]]
 name = "spl-transfer-hook"

token/transfer-hook/cli/src/main.rs

@@ -263,7 +263,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 Additional accounts with known fixed addresses can be passed at the command line in the format "<PUBKEY>:<ROLE>". The role must be "readonly", "writable". "readonlySigner", or "writableSigner".
 
 Additional accounts requiring seed configurations can be defined in a configuration file using either JSON or YAML. The format is as follows:
-                            
+
 ```json
 {
     "extraMetas": [
@@ -347,7 +347,7 @@ extraMetas:
 Additional accounts with known fixed addresses can be passed at the command line in the format "<PUBKEY>:<ROLE>". The role must be "readonly", "writable". "readonlySigner", or "writableSigner".
 
 Additional accounts requiring seed configurations can be defined in a configuration file using either JSON or YAML. The format is as follows:
-                            
+
 ```json
 {
     "extraMetas": [
@@ -551,19 +551,27 @@ extraMetas:
 mod test {
     use {
         super::*,
-        solana_sdk::{bpf_loader_upgradeable, instruction::AccountMeta, signer::keypair::Keypair},
+        solana_sdk::{
+            account::Account, bpf_loader_upgradeable, instruction::AccountMeta,
+            program_option::COption, signer::keypair::Keypair,
+        },
         solana_test_validator::{TestValidator, TestValidatorGenesis, UpgradeableProgramInfo},
+        spl_token_2022::{
+            extension::{ExtensionType, StateWithExtensionsMut},
+            state::Mint,
+        },
         spl_token_client::{
-            client::{
-                ProgramClient, ProgramRpcClient, ProgramRpcClientSendTransaction, SendTransaction,
-                SimulateTransaction,
-            },
+            client::{ProgramRpcClient, ProgramRpcClientSendTransaction},
             token::Token,
         },
         std::{path::PathBuf, sync::Arc},
     };
 
-    async fn new_validator_for_test(program_id: Pubkey) -> (TestValidator, Keypair) {
+    async fn new_validator_for_test(
+        program_id: Pubkey,
+        mint_authority: &Pubkey,
+        decimals: u8,
+    ) -> (TestValidator, Keypair) {
         solana_logger::setup();
         let mut test_validator_genesis = TestValidatorGenesis::default();
         test_validator_genesis.add_upgradeable_programs_with_path(&[UpgradeableProgramInfo {
@@ -572,54 +580,55 @@ mod test {
             program_path: PathBuf::from("../../../target/deploy/spl_transfer_hook_example.so"),
             upgrade_authority: Pubkey::new_unique(),
         }]);
-        test_validator_genesis.start_async().await
-    }
 
-    async fn setup_mint<T: SendTransaction + SimulateTransaction>(
-        program_id: &Pubkey,
-        mint_authority: &Pubkey,
-        decimals: u8,
-        payer: Arc<dyn Signer>,
-        client: Arc<dyn ProgramClient<T>>,
-    ) -> Token<T> {
-        let mint_account = Keypair::new();
-        let token = Token::new(
-            client,
-            program_id,
-            &mint_account.pubkey(),
-            Some(decimals),
-            payer,
+        let mint_size = ExtensionType::try_calculate_account_len::<Mint>(&[]).unwrap();
+        let mut mint_data = vec![0; mint_size];
+        let mut state =
+            StateWithExtensionsMut::<Mint>::unpack_uninitialized(&mut mint_data).unwrap();
+        let token_amount = 1_000_000_000_000;
+        state.base = Mint {
+            mint_authority: COption::Some(*mint_authority),
+            supply: token_amount,
+            decimals,
+            is_initialized: true,
+            freeze_authority: COption::None,
+        };
+        state.pack_base();
+        test_validator_genesis.add_account(
+            spl_transfer_hook_example::mint::id(),
+            Account {
+                lamports: 1_000_000_000,
+                data: mint_data,
+                owner: spl_token_2022::id(),
+                ..Account::default()
+            }
+            .into(),
         );
-        token
-            .create_mint(mint_authority, None, vec![], &[&mint_account])
-            .await
-            .unwrap();
-        token
+        test_validator_genesis.start_async().await
     }
 
     #[tokio::test]
     async fn test_create() {
         let program_id = Pubkey::new_unique();
 
-        let (test_validator, payer) = new_validator_for_test(program_id).await;
+        let decimals = 2;
+        let mint_authority = Keypair::new();
+        let (test_validator, payer) =
+            new_validator_for_test(program_id, &mint_authority.pubkey(), decimals).await;
         let payer: Arc<dyn Signer> = Arc::new(payer);
         let rpc_client = Arc::new(test_validator.get_async_rpc_client());
         let client = Arc::new(ProgramRpcClient::new(
             rpc_client.clone(),
             ProgramRpcClientSendTransaction,
         ));
 
-        let mint_authority = Keypair::new();
-        let decimals = 2;
-
-        let token = setup_mint(
+        let token = Token::new(
+            client.clone(),
             &spl_token_2022::id(),
-            &mint_authority.pubkey(),
-            decimals,
+            &spl_transfer_hook_example::mint::id(),
+            Some(decimals),
             payer.clone(),
-            client.clone(),
-        )
-        .await;
+        );
 
         let required_address = Pubkey::new_unique();
         let accounts = vec![AccountMeta::new_readonly(required_address, false)];

token/transfer-hook/example/Cargo.toml

@@ -8,8 +8,10 @@ license = "Apache-2.0"
 edition = "2021"
 
 [features]
+default = ["forbid-additional-mints"]
 no-entrypoint = []
 test-sbf = []
+forbid-additional-mints = []
 
 [dependencies]
 arrayref = "0.3.7"

token/transfer-hook/example/src/lib.rs

@@ -16,3 +16,15 @@ mod entrypoint;
 // Export current sdk types for downstream users building with a different sdk
 // version
 pub use solana_program;
+
+/// Place the mint id that you want to target with your transfer hook program.
+/// Any other mint will fail to initialize, protecting the transfer hook program
+/// from rogue mints trying to get access to accounts.
+///
+/// There are many situations where it's reasonable to support multiple mints
+/// with one transfer-hook program, but because it's easy to make something
+/// unsafe, this simple example implementation only allows for one mint.
+#[cfg(feature = "forbid-additional-mints")]
+pub mod mint {
+    solana_program::declare_id!("Mint111111111111111111111111111111111111111");
+}

token/transfer-hook/example/src/processor.rs

@@ -88,6 +88,13 @@ pub fn process_initialize_extra_account_meta_list(
     let authority_info = next_account_info(account_info_iter)?;
     let _system_program_info = next_account_info(account_info_iter)?;
 
+    // check that the one mint we want to target is trying to create extra
+    // account metas
+    #[cfg(feature = "forbid-additional-mints")]
+    if *mint_info.key != crate::mint::id() {
+        return Err(ProgramError::InvalidArgument);
+    }
+
     // check that the mint authority is valid without fully deserializing
     let mint_data = mint_info.try_borrow_data()?;
     let mint = StateWithExtensions::<Mint>::unpack(&mint_data)?;

token/transfer-hook/example/tests/functional.rs

@@ -142,7 +142,7 @@ async fn success_execute() {
 
     let token_program_id = spl_token_2022::id();
     let wallet = Keypair::new();
-    let mint_address = Pubkey::new_unique();
+    let mint_address = spl_transfer_hook_example::mint::id();
     let mint_authority = Keypair::new();
     let mint_authority_pubkey = mint_authority.pubkey();
     let source = Pubkey::new_unique();
@@ -439,7 +439,7 @@ async fn fail_incorrect_derivation() {
 
     let token_program_id = spl_token_2022::id();
     let wallet = Keypair::new();
-    let mint_address = Pubkey::new_unique();
+    let mint_address = spl_transfer_hook_example::mint::id();
     let mint_authority = Keypair::new();
     let mint_authority_pubkey = mint_authority.pubkey();
     let source = Pubkey::new_unique();
@@ -495,6 +495,69 @@ async fn fail_incorrect_derivation() {
     );
 }
 
+#[tokio::test]
+async fn fail_incorrect_mint() {
+    let program_id = Pubkey::new_unique();
+    let mut program_test = setup(&program_id);
+
+    let token_program_id = spl_token_2022::id();
+    let wallet = Keypair::new();
+    // wrong mint, only `spl_transfer_hook_example::mint::id()` allowed
+    let mint_address = Pubkey::new_unique();
+    let mint_authority = Keypair::new();
+    let mint_authority_pubkey = mint_authority.pubkey();
+    let source = Pubkey::new_unique();
+    let destination = Pubkey::new_unique();
+    let decimals = 2;
+    setup_token_accounts(
+        &mut program_test,
+        &token_program_id,
+        &mint_address,
+        &mint_authority_pubkey,
+        &source,
+        &destination,
+        &wallet.pubkey(),
+        decimals,
+        true,
+    );
+
+    let extra_account_metas = get_extra_account_metas_address(&mint_address, &program_id);
+
+    let mut context = program_test.start_with_context().await;
+    let rent = context.banks_client.get_rent().await.unwrap();
+    let rent_lamports = rent.minimum_balance(ExtraAccountMetaList::size_of(0).unwrap());
+
+    let transaction = Transaction::new_signed_with_payer(
+        &[
+            system_instruction::transfer(
+                &context.payer.pubkey(),
+                &extra_account_metas,
+                rent_lamports,
+            ),
+            initialize_extra_account_meta_list(
+                &program_id,
+                &extra_account_metas,
+                &mint_address,
+                &mint_authority_pubkey,
+                &[],
+            ),
+        ],
+        Some(&context.payer.pubkey()),
+        &[&context.payer, &mint_authority],
+        context.last_blockhash,
+    );
+    let error = context
+        .banks_client
+        .process_transaction(transaction)
+        .await
+        .unwrap_err()
+        .unwrap();
+    assert_eq!(
+        error,
+        TransactionError::InstructionError(1, InstructionError::InvalidArgument)
+    );
+}
+
 /// Test program to CPI into default transfer-hook-interface program
 pub fn process_instruction(
     _program_id: &Pubkey,
@@ -530,7 +593,7 @@ async fn success_on_chain_invoke() {
 
     let token_program_id = spl_token_2022::id();
     let wallet = Keypair::new();
-    let mint_address = Pubkey::new_unique();
+    let mint_address = spl_transfer_hook_example::mint::id();
     let mint_authority = Keypair::new();
     let mint_authority_pubkey = mint_authority.pubkey();
     let source = Pubkey::new_unique();
@@ -673,7 +736,7 @@ async fn fail_without_transferring_flag() {
 
     let token_program_id = spl_token_2022::id();
     let wallet = Keypair::new();
-    let mint_address = Pubkey::new_unique();
+    let mint_address = spl_transfer_hook_example::mint::id();
     let mint_authority = Keypair::new();
     let mint_authority_pubkey = mint_authority.pubkey();
     let source = Pubkey::new_unique();
@@ -767,7 +830,7 @@ async fn success_on_chain_invoke_with_updated_extra_account_metas() {
 
     let token_program_id = spl_token_2022::id();
     let wallet = Keypair::new();
-    let mint_address = Pubkey::new_unique();
+    let mint_address = spl_transfer_hook_example::mint::id();
     let mint_authority = Keypair::new();
     let mint_authority_pubkey = mint_authority.pubkey();
     let source = Pubkey::new_unique();
@@ -970,7 +1033,7 @@ async fn success_execute_with_updated_extra_account_metas() {
 
     let token_program_id = spl_token_2022::id();
     let wallet = Keypair::new();
-    let mint_address = Pubkey::new_unique();
+    let mint_address = spl_transfer_hook_example::mint::id();
     let mint_authority = Keypair::new();
     let mint_authority_pubkey = mint_authority.pubkey();
     let source = Pubkey::new_unique();