use strict treasury_tokens loader

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,
+);

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.

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

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,

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