From 3b1039b46908647df42dbd66c7967b7ef3dbe27e Mon Sep 17 00:00:00 2001
From: Hardhat Chad <hardhatchad@gmail.com>
Date: Tue, 2 Jul 2024 13:08:35 +0000
Subject: [PATCH] use strict treasury_tokens loader

---
 api/src/consts.rs    | 13 +++++++++++++
 api/src/loaders.rs   | 14 ++++++++++++++
 program/src/claim.rs | 10 ++--------
 program/src/reset.rs |  7 +------
 program/src/stake.rs |  7 +------
 5 files changed, 31 insertions(+), 20 deletions(-)

diff --git a/api/src/consts.rs b/api/src/consts.rs
index 92f437a..d0d7fae 100644
--- a/api/src/consts.rs
+++ b/api/src/consts.rs
@@ -131,3 +131,16 @@ pub const TREASURY_ADDRESS: Pubkey =
 
 /// The bump of the treasury account, for cpis.
 pub const TREASURY_BUMP: u8 = ed25519::derive_program_address(&[TREASURY], &PROGRAM_ID).1;
+
+/// The address of the treasury token account.
+pub const TREASURY_TOKENS_ADDRESS: Pubkey = Pubkey::new_from_array(
+    ed25519::derive_program_address(
+        &[
+            unsafe { &*(&TREASURY_ADDRESS as *const Pubkey as *const [u8; 32]) },
+            unsafe { &*(&spl_token::id() as *const Pubkey as *const [u8; 32]) },
+            unsafe { &*(&MINT_ADDRESS as *const Pubkey as *const [u8; 32]) },
+        ],
+        unsafe { &*(&spl_token::id() as *const Pubkey as *const [u8; 32]) },
+    )
+    .0,
+);
diff --git a/api/src/loaders.rs b/api/src/loaders.rs
index 0b95d4b..998dca6 100644
--- a/api/src/loaders.rs
+++ b/api/src/loaders.rs
@@ -252,6 +252,20 @@ pub fn load_treasury<'a, 'info>(
     Ok(())
 }
 
+/// Errors if:
+/// - Address does not match the expected treasury tokens address.
+/// - Cannot load as a token account
+pub fn load_treasury_tokens<'a, 'info>(
+    info: &'a AccountInfo<'info>,
+    is_writable: bool,
+) -> Result<(), ProgramError> {
+    if info.key.ne(&TREASURY_TOKENS_ADDRESS) {
+        return Err(ProgramError::InvalidSeeds);
+    }
+
+    load_token_account(info, Some(&TREASURY_ADDRESS), &MINT_ADDRESS, is_writable)
+}
+
 /// Errors if:
 /// - Owner is not SPL token program.
 /// - Address does not match the expected mint address.
diff --git a/program/src/claim.rs b/program/src/claim.rs
index 9a5ae4e..c6a68b9 100644
--- a/program/src/claim.rs
+++ b/program/src/claim.rs
@@ -24,22 +24,16 @@ pub fn process_claim<'a, 'info>(
     let amount = u64::from_le_bytes(args.amount);
 
     // Load accounts
-    let [signer, beneficiary_info, mint_info, proof_info, treasury_info, treasury_tokens_info, token_program] =
+    let [signer, beneficiary_info, proof_info, treasury_info, treasury_tokens_info, token_program] =
         accounts
     else {
         return Err(ProgramError::NotEnoughAccountKeys);
     };
     load_signer(signer)?;
     load_token_account(beneficiary_info, None, &MINT_ADDRESS, true)?;
-    load_mint(mint_info, MINT_ADDRESS, true)?;
     load_proof(proof_info, signer.key, true)?;
     load_treasury(treasury_info, false)?;
-    load_token_account(
-        treasury_tokens_info,
-        Some(treasury_info.key),
-        &MINT_ADDRESS,
-        true,
-    )?;
+    load_treasury_tokens(treasury_tokens_info, true)?;
     load_program(token_program, spl_token::id())?;
 
     // Update miner balance
diff --git a/program/src/reset.rs b/program/src/reset.rs
index 0e78401..716190e 100644
--- a/program/src/reset.rs
+++ b/program/src/reset.rs
@@ -54,12 +54,7 @@ pub fn process_reset<'a, 'info>(
     load_config(config_info, true)?;
     load_mint(mint_info, MINT_ADDRESS, true)?;
     load_treasury(treasury_info, true)?;
-    load_token_account(
-        treasury_tokens_info,
-        Some(treasury_info.key),
-        mint_info.key,
-        true,
-    )?;
+    load_treasury_tokens(treasury_tokens_info, true)?;
     load_program(token_program, spl_token::id())?;
     let busses: [&AccountInfo; BUS_COUNT] = [
         bus_0_info, bus_1_info, bus_2_info, bus_3_info, bus_4_info, bus_5_info, bus_6_info,
diff --git a/program/src/stake.rs b/program/src/stake.rs
index dbd4f0a..48cfb6f 100644
--- a/program/src/stake.rs
+++ b/program/src/stake.rs
@@ -30,12 +30,7 @@ pub fn process_stake<'a, 'info>(
     load_signer(signer)?;
     load_proof(proof_info, signer.key, true)?;
     load_token_account(sender_info, Some(signer.key), &MINT_ADDRESS, true)?;
-    load_token_account(
-        treasury_tokens_info,
-        Some(&TREASURY_ADDRESS),
-        &MINT_ADDRESS,
-        true,
-    )?;
+    load_treasury_tokens(treasury_tokens_info, true)?;
     load_program(token_program, spl_token::id())?;
 
     // Update proof balance