feature(rpc): Adds rpc methods to enable/disable other cheatcodes

crates/core/src/error.rs

@@ -104,6 +104,18 @@ impl SurfpoolError {
         Self(error)
     }
 
+    pub fn disable_cheatcode(e: String) -> Self {
+        let mut error = Error::invalid_request();
+        error.data = Some(json!(format!("Unable to disable cheatcode: {}", e)));
+        Self(error)
+    }
+
+    pub fn enable_cheatcode(e: String) -> Self {
+        let mut error = Error::invalid_request();
+        error.data = Some(json!(format!("Unable to enable cheatcode: {}", e)));
+        Self(error)
+    }
+
     pub fn set_account<T>(pubkey: Pubkey, e: T) -> Self
     where
         T: ToString,

crates/core/src/rpc/accounts_scan.rs

@@ -910,7 +910,7 @@ mod tests {
     #[tokio::test(flavor = "multi_thread")]
     async fn test_set_and_get_supply() {
         let setup = TestSetup::new(SurfpoolAccountsScanRpc);
-        let cheatcodes_rpc = SurfnetCheatcodesRpc;
+        let cheatcodes_rpc = SurfnetCheatcodesRpc::empty();
 
         // test initial default values
         let initial_supply = setup
@@ -959,7 +959,7 @@ mod tests {
     #[tokio::test(flavor = "multi_thread")]
     async fn test_get_supply_exclude_accounts() {
         let setup = TestSetup::new(SurfpoolAccountsScanRpc);
-        let cheatcodes_rpc = SurfnetCheatcodesRpc;
+        let cheatcodes_rpc = SurfnetCheatcodesRpc::empty();
 
         // set supply with non-circulating accounts
         let supply_update = SupplyUpdate {
@@ -1007,7 +1007,7 @@ mod tests {
     #[tokio::test(flavor = "multi_thread")]
     async fn test_partial_supply_update() {
         let setup = TestSetup::new(SurfpoolAccountsScanRpc);
-        let cheatcodes_rpc = SurfnetCheatcodesRpc;
+        let cheatcodes_rpc = SurfnetCheatcodesRpc::empty();
 
         // set initial values
         let initial_update = SupplyUpdate {
@@ -1051,7 +1051,7 @@ mod tests {
     #[tokio::test(flavor = "multi_thread")]
     async fn test_set_supply_with_multiple_invalid_pubkeys() {
         let setup = TestSetup::new(SurfpoolAccountsScanRpc);
-        let cheatcodes_rpc = SurfnetCheatcodesRpc;
+        let cheatcodes_rpc = SurfnetCheatcodesRpc::empty();
         let invalid_pubkey = "invalid_pubkey";
 
         // test with multiple invalid pubkeys - should fail on the first one
@@ -1082,7 +1082,7 @@ mod tests {
     #[tokio::test(flavor = "multi_thread")]
     async fn test_set_supply_with_max_values() {
         let setup = TestSetup::new(SurfpoolAccountsScanRpc);
-        let cheatcodes_rpc = SurfnetCheatcodesRpc;
+        let cheatcodes_rpc = SurfnetCheatcodesRpc::empty();
 
         let supply_update = SupplyUpdate {
             total: Some(u64::MAX),
@@ -1112,7 +1112,7 @@ mod tests {
     #[tokio::test(flavor = "multi_thread")]
     async fn test_set_supply_large_valid_account_list() {
         let setup = TestSetup::new(SurfpoolAccountsScanRpc);
-        let cheatcodes_rpc = SurfnetCheatcodesRpc;
+        let cheatcodes_rpc = SurfnetCheatcodesRpc::empty();
 
         let large_account_list: Vec<String> = (0..100)
             .map(|i| match i % 10 {

crates/core/src/rpc/mod.rs

@@ -1,4 +1,7 @@
-use std::{future::Future, sync::Arc};
+use std::{
+    future::Future,
+    sync::{Arc, Mutex},
+};
 
 use blake3::Hash;
 use crossbeam_channel::Sender;
@@ -9,7 +12,7 @@ use jsonrpc_core::{
 };
 use jsonrpc_pubsub::{PubSubMetadata, Session};
 use solana_clock::Slot;
-use surfpool_types::{SimnetCommand, SimnetEvent, types::RpcConfig};
+use surfpool_types::{CheatcodeConfig, SimnetCommand, SimnetEvent, types::RpcConfig};
 
 use crate::{
     PluginManagerCommand,
@@ -48,6 +51,7 @@ pub struct RunloopContext {
     pub plugin_manager_commands_tx: Sender<PluginManagerCommand>,
     pub remote_rpc_client: Option<SurfnetRemoteClient>,
     pub rpc_config: RpcConfig,
+    pub cheatcode_config: Arc<Mutex<CheatcodeConfig>>,
 }
 
 pub struct SurfnetRpcContext<T> {
@@ -113,6 +117,7 @@ pub struct SurfpoolMiddleware {
     pub plugin_manager_commands_tx: Sender<PluginManagerCommand>,
     pub config: RpcConfig,
     pub remote_rpc_client: Option<SurfnetRemoteClient>,
+    pub cheatcode_config: Arc<Mutex<CheatcodeConfig>>,
 }
 
 impl SurfpoolMiddleware {
@@ -129,6 +134,7 @@ impl SurfpoolMiddleware {
             plugin_manager_commands_tx: plugin_manager_commands_tx.clone(),
             config: config.clone(),
             remote_rpc_client: remote_rpc_client.clone(),
+            cheatcode_config: CheatcodeConfig::new(),
         }
     }
 }
@@ -171,7 +177,43 @@ impl Middleware<Option<RunloopContext>> for SurfpoolMiddleware {
             plugin_manager_commands_tx: self.plugin_manager_commands_tx.clone(),
             remote_rpc_client: self.remote_rpc_client.clone(),
             rpc_config: self.config.clone(),
+            cheatcode_config: self.cheatcode_config.clone(),
         });
+
+        // All surfnet cheatcodes will start with surfnet. If the request is a cheatcode, make sure it isn't disabled.
+        if method_name.starts_with("surfnet_")
+            && let Some(meta_val) = meta.clone()
+        {
+            let Ok(meta_val) = meta_val.cheatcode_config.lock() else {
+                let error = Response::from(
+                    Error {
+                        code: ErrorCode::InternalError,
+                        message: "An internal server error occured".to_string(),
+                        data: None,
+                    },
+                    None,
+                );
+                warn!("Request rejected due to cheatsheet being disabled");
+
+                return Either::Left(Box::pin(async move { Some(error) }));
+            };
+            if meta_val.is_cheatcode_disabled(&method_name) {
+                let error = Response::from(
+                    Error {
+                        code: ErrorCode::InvalidRequest,
+                        message: format!(
+                            "Cheatcode rpc method: {method_name} is currently disabled"
+                        ),
+                        data: None,
+                    },
+                    None,
+                );
+                warn!("Request rejected due to cheatcode rpc method being disabled");
+
+                return Either::Left(Box::pin(async move { Some(error) }));
+            }
+        }
+
         Either::Left(Box::pin(next(request, meta).map(move |res| {
             if let Some(Response::Single(output)) = &res {
                 if let jsonrpc_core::Output::Failure(failure) = output {
@@ -222,6 +264,7 @@ impl Middleware<Option<SurfpoolWebsocketMeta>> for SurfpoolWebsocketMiddleware {
             plugin_manager_commands_tx: self.surfpool_middleware.plugin_manager_commands_tx.clone(),
             remote_rpc_client: self.surfpool_middleware.remote_rpc_client.clone(),
             rpc_config: self.surfpool_middleware.config.clone(),
+            cheatcode_config: self.surfpool_middleware.cheatcode_config.clone(),
         };
         let session = meta
             .as_ref()