From 4658815e9f8204d61d0e089cf450c0381c340a5a Mon Sep 17 00:00:00 2001
From: Jon Cinque <jon.cinque@gmail.com>
Date: Tue, 19 Oct 2021 01:39:57 +0200
Subject: [PATCH] stake-pool: Force pools to only use the SPL token program
 (#2521)

---
 stake-pool/program/src/processor.rs    |   9 +++
 stake-pool/program/tests/initialize.rs | 101 ++++++++++++++++++++++++-
 2 files changed, 107 insertions(+), 3 deletions(-)

diff --git a/stake-pool/program/src/processor.rs b/stake-pool/program/src/processor.rs
index 9c065f1050e..82ff479a610 100644
--- a/stake-pool/program/src/processor.rs
+++ b/stake-pool/program/src/processor.rs
@@ -565,6 +565,15 @@ impl Processor {
             return Err(StakePoolError::FeeTooHigh.into());
         }
 
+        if *token_program_info.key != spl_token::id() {
+            msg!(
+                "Only the SPL token program is currently supported, expected {}, received {}",
+                spl_token::id(),
+                *token_program_info.key
+            );
+            return Err(ProgramError::IncorrectProgramId);
+        }
+
         if manager_fee_info.owner != token_program_info.key {
             return Err(ProgramError::IncorrectProgramId);
         }
diff --git a/stake-pool/program/tests/initialize.rs b/stake-pool/program/tests/initialize.rs
index 1c4044a68e7..e9aa1160e1a 100644
--- a/stake-pool/program/tests/initialize.rs
+++ b/stake-pool/program/tests/initialize.rs
@@ -454,10 +454,107 @@ async fn fail_with_wrong_token_program_id() {
     .await
     .unwrap();
 
+    create_token_account(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &stake_pool_accounts.pool_fee_account,
+        &stake_pool_accounts.pool_mint.pubkey(),
+        &stake_pool_accounts.manager.pubkey(),
+    )
+    .await
+    .unwrap();
+
     let rent = banks_client.get_rent().await.unwrap();
+    let rent_stake_pool = rent.minimum_balance(get_packed_len::<state::StakePool>());
+    let validator_list_size = get_instance_packed_len(&state::ValidatorList::new(
+        stake_pool_accounts.max_validators,
+    ))
+    .unwrap();
+    let rent_validator_list = rent.minimum_balance(validator_list_size);
 
-    let account_rent = rent.minimum_balance(spl_token::state::Account::LEN);
     let mut transaction = Transaction::new_with_payer(
+        &[
+            system_instruction::create_account(
+                &payer.pubkey(),
+                &stake_pool_accounts.stake_pool.pubkey(),
+                rent_stake_pool,
+                get_packed_len::<state::StakePool>() as u64,
+                &id(),
+            ),
+            system_instruction::create_account(
+                &payer.pubkey(),
+                &stake_pool_accounts.validator_list.pubkey(),
+                rent_validator_list,
+                validator_list_size as u64,
+                &id(),
+            ),
+            instruction::initialize(
+                &id(),
+                &stake_pool_accounts.stake_pool.pubkey(),
+                &stake_pool_accounts.manager.pubkey(),
+                &stake_pool_accounts.staker.pubkey(),
+                &stake_pool_accounts.validator_list.pubkey(),
+                &stake_pool_accounts.reserve_stake.pubkey(),
+                &stake_pool_accounts.pool_mint.pubkey(),
+                &stake_pool_accounts.pool_fee_account.pubkey(),
+                &wrong_token_program.pubkey(),
+                None,
+                stake_pool_accounts.epoch_fee,
+                stake_pool_accounts.withdrawal_fee,
+                stake_pool_accounts.deposit_fee,
+                stake_pool_accounts.referral_fee,
+                stake_pool_accounts.max_validators,
+            ),
+        ],
+        Some(&payer.pubkey()),
+    );
+    transaction.sign(
+        &[
+            &payer,
+            &stake_pool_accounts.stake_pool,
+            &stake_pool_accounts.validator_list,
+            &stake_pool_accounts.manager,
+        ],
+        recent_blockhash,
+    );
+    let transaction_error = banks_client
+        .process_transaction(transaction)
+        .await
+        .err()
+        .unwrap();
+
+    match transaction_error {
+        TransportError::TransactionError(TransactionError::InstructionError(_, error)) => {
+            assert_eq!(error, InstructionError::IncorrectProgramId);
+        }
+        _ => panic!(
+            "Wrong error occurs while try to initialize stake pool with wrong token program ID"
+        ),
+    }
+}
+
+#[tokio::test]
+async fn fail_with_fee_owned_by_wrong_token_program_id() {
+    let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
+    let stake_pool_accounts = StakePoolAccounts::new();
+
+    let wrong_token_program = Keypair::new();
+
+    create_mint(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &stake_pool_accounts.pool_mint,
+        &stake_pool_accounts.withdraw_authority,
+    )
+    .await
+    .unwrap();
+
+    let rent = banks_client.get_rent().await.unwrap();
+
+    let account_rent = rent.minimum_balance(spl_token::state::Account::LEN);
+    let transaction = Transaction::new_signed_with_payer(
         &[system_instruction::create_account(
             &payer.pubkey(),
             &stake_pool_accounts.pool_fee_account.pubkey(),
@@ -466,8 +563,6 @@ async fn fail_with_wrong_token_program_id() {
             &wrong_token_program.pubkey(),
         )],
         Some(&payer.pubkey()),
-    );
-    transaction.sign(
         &[&payer, &stake_pool_accounts.pool_fee_account],
         recent_blockhash,
     );