From 24aace2c24a0b3595e08237988e7e9274be65112 Mon Sep 17 00:00:00 2001
From: Hardhat Chad <hardhatchad@gmail.com>
Date: Fri, 11 Oct 2024 20:58:56 +0000
Subject: [PATCH] add stricter address check

---
 api/src/loaders.rs  | 8 ++++++++
 program/src/mine.rs | 2 +-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/api/src/loaders.rs b/api/src/loaders.rs
index ebc7383..de380e2 100644
--- a/api/src/loaders.rs
+++ b/api/src/loaders.rs
@@ -6,12 +6,20 @@ use crate::{
 };
 
 pub trait OreAccountInfoValidation {
+    fn is_bus(&self) -> Result<&Self, ProgramError>;
     fn is_config(&self) -> Result<&Self, ProgramError>;
     fn is_treasury(&self) -> Result<&Self, ProgramError>;
     fn is_treasury_tokens(&self) -> Result<&Self, ProgramError>;
 }
 
 impl OreAccountInfoValidation for AccountInfo<'_> {
+    fn is_bus(&self) -> Result<&Self, ProgramError> {
+        if !BUS_ADDRESSES.contains(self.key) {
+            return Err(ProgramError::InvalidSeeds);
+        }
+        Ok(self)
+    }
+
     fn is_config(&self) -> Result<&Self, ProgramError> {
         self.has_address(&CONFIG_ADDRESS)?
             .is_type::<Config>(&crate::ID)
diff --git a/program/src/mine.rs b/program/src/mine.rs
index be4c55b..486fe32 100644
--- a/program/src/mine.rs
+++ b/program/src/mine.rs
@@ -29,7 +29,7 @@ pub fn process_mine(accounts: &[AccountInfo], data: &[u8]) -> ProgramResult {
         return Err(ProgramError::NotEnoughAccountKeys);
     };
     signer_info.is_signer()?;
-    let bus = bus_info.to_account_mut::<Bus>(&ore_api::ID)?;
+    let bus = bus_info.is_bus()?.to_account_mut::<Bus>(&ore_api::ID)?;
     let config = config_info
         .is_config()?
         .to_account::<Config>(&ore_api::ID)?;