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/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,9 @@ use jsonrpc_core::{
 };
 use jsonrpc_pubsub::{PubSubMetadata, Session};
 use solana_clock::Slot;
-use surfpool_types::{SimnetCommand, SimnetEvent, types::RpcConfig};
+use surfpool_types::{
+    CheatcodeConfig, RpcCheatcodes, SimnetCommand, SimnetEvent, types::RpcConfig,
+};
 
 use crate::{
     PluginManagerCommand,
@@ -48,6 +53,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 +119,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 +136,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 +179,30 @@ 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(),
         });
+
+        if RpcCheatcodes::try_from(method_name.as_str()).is_ok()
+            && let Some(meta_val) = meta.clone()
+            && meta_val
+                .cheatcode_config
+                .lock()
+                .unwrap() // this is okay since only on_request, disable_cheatcode and enable_cheatcode only use it, the rpc method being called after on_request
+                .is_cheatcode_disabled(&method_name)
+        {
+            let error = Response::from(
+                Error {
+                    code: ErrorCode::InvalidRequest,
+                    message: format!("Cheatsheet rpc method: {method_name} is currently disabled"),
+                    data: None,
+                },
+                None,
+            );
+            warn!("Request rejected due to cheatsheet 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 +253,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()

crates/core/src/rpc/surfnet_cheatcodes.rs

@@ -14,9 +14,10 @@ use solana_system_interface::program as system_program;
 use solana_transaction::versioned::VersionedTransaction;
 use spl_associated_token_account_interface::address::get_associated_token_address_with_program_id;
 use surfpool_types::{
-    AccountSnapshot, ClockCommand, ExportSnapshotConfig, GetStreamedAccountsResponse,
-    GetSurfnetInfoResponse, Idl, ResetAccountConfig, RpcProfileResultConfig, Scenario,
-    SimnetCommand, SimnetEvent, StreamAccountConfig, UiKeyedProfileResult,
+    AccountSnapshot, CheatcodeControlConfig, CheatcodeFilter, ClockCommand, ExportSnapshotConfig,
+    GetStreamedAccountsResponse, GetSurfnetInfoResponse, Idl, ResetAccountConfig, RpcCheatcodes,
+    RpcProfileResultConfig, Scenario, SimnetCommand, SimnetEvent, StreamAccountConfig,
+    UiKeyedProfileResult,
     types::{AccountUpdate, SetSomeAccount, SupplyUpdate, TokenAccountUpdate, UuidOrSignature},
 };
 
@@ -179,6 +180,92 @@ pub trait SurfnetCheatcodes {
         update: AccountUpdate,
     ) -> BoxFuture<Result<RpcResponse<()>>>;
 
+    /// Enables one or more Surfpool cheatcode RPC methods for the current session.
+    ///
+    /// This method allows developers to re-enable cheatcode methods that were previously disabled.
+    /// Each cheatcode name must match a valid `surfnet_*` RPC method (e.g. `surfnet_setAccount`, `surfnet_timeTravel`).
+    ///
+    /// ## Parameters
+    /// - `cheatcodes`: A list of cheatcode method names to enable, as strings (e.g. `["surfnet_setAccount", "surfnet_timeTravel"]`).
+    ///
+    /// ## Returns
+    /// A `RpcResponse<()>` indicating whether the enable operation was successful.
+    ///
+    /// ## Example Request
+    /// ```json
+    /// {
+    ///   "jsonrpc": "2.0",
+    ///   "id": 1,
+    ///   "method": "surfnet_enableCheatcode",
+    ///   "params": [["surfnet_setAccount", "surfnet_timeTravel"]]
+    /// }
+    /// ```
+    ///
+    /// ## Example Response
+    /// ```json
+    /// {
+    ///   "jsonrpc": "2.0",
+    ///   "result": {},
+    ///   "id": 1
+    /// }
+    /// ```
+    ///
+    /// # Notes
+    /// Invalid cheatcode names return an error. Use `surfnet_disableCheatcode` to disable methods and `surfnet_lockout` to allow disabling `surfnet_enableCheatcode` and `surfnet_disableCheatcode` themselves.
+    ///
+    /// # See Also
+    /// - `surfnet_disableCheatcode`, `surfnet_disableAllCheatcodes`, `surfnet_lockout`
+    #[rpc(meta, name = "surfnet_enableCheatcode")]
+    fn enable_cheatcode(
+        &self,
+        meta: Self::Metadata,
+        cheatcodes_filter: CheatcodeFilter,
+    ) -> Result<RpcResponse<()>>;
+
+    /// Disables one or more Surfpool cheatcode RPC methods for the current session.
+    ///
+    /// This method allows developers to turn off specific cheatcode methods so they are no longer callable.
+    /// Each cheatcode name must match a valid `surfnet_*` RPC method. When lockout is not enabled,
+    /// `surfnet_enableCheatcode` and `surfnet_disableCheatcode` cannot be disabled.
+    ///
+    /// ## Parameters
+    /// - `cheatcodes`: A list of cheatcode method names to disable, as strings (e.g. `["surfnet_setAccount"]`).
+    ///
+    /// ## Returns
+    /// A `RpcResponse<()>` indicating whether the disable operation was successful.
+    ///
+    /// ## Example Request
+    /// ```json
+    /// {
+    ///   "jsonrpc": "2.0",
+    ///   "id": 1,
+    ///   "method": "surfnet_disableCheatcode",
+    ///   "params": [["surfnet_setAccount", "surfnet_timeTravel"]]
+    /// }
+    /// ```
+    ///
+    /// ## Example Response
+    /// ```json
+    /// {
+    ///   "jsonrpc": "2.0",
+    ///   "result": {},
+    ///   "id": 1
+    /// }
+    /// ```
+    ///
+    /// # Notes
+    /// Call `surfnet_lockout` first if you need to disable `surfnet_enableCheatcode` or `surfnet_disableCheatcode`.
+    ///
+    /// # See Also
+    /// - `surfnet_enableCheatcode`, `surfnet_disableAllCheatcodes`, `surfnet_lockout`
+    #[rpc(meta, name = "surfnet_disableCheatcode")]
+    fn disable_cheatcode(
+        &self,
+        meta: Self::Metadata,
+        cheatcodes_filter: CheatcodeFilter,
+        lockout: Option<CheatcodeControlConfig>,
+    ) -> Result<RpcResponse<()>>;
+
     /// A "cheat code" method for developers to set or update a token account in Surfpool.
     ///
     /// This method allows developers to set or update various properties of a token account,
@@ -1201,6 +1288,111 @@ impl SurfnetCheatcodes for SurfnetCheatcodesRpc {
         })
     }
 
+    fn disable_cheatcode(
+        &self,
+        meta: Self::Metadata,
+        cheatcodes_filter: CheatcodeFilter,
+        control_config: Option<CheatcodeControlConfig>,
+    ) -> Result<RpcResponse<()>> {
+        let svm_locker = match meta.get_svm_locker() {
+            Ok(locker) => locker,
+            Err(e) => return Err(e.into()),
+        };
+
+        let CheatcodeControlConfig { lockout } = control_config.unwrap_or_default();
+        let lockout = lockout.unwrap_or_default();
+
+        if let Some(runloop_ctx) = meta {
+            let mut cheatcode_ctx = runloop_ctx.cheatcode_config.lock().unwrap();
+
+            match cheatcodes_filter {
+                CheatcodeFilter::All(all) => {
+                    if all.ne("all") {
+                        return Err(SurfpoolError::disable_cheatcode(
+                            "Invalid optioin provided for disabling all cheatcodes. Try using \"all\"".to_string(),
+                        )
+                        .into());
+                    }
+
+                    cheatcode_ctx.disable_all(lockout);
+                }
+                CheatcodeFilter::List(cheatdcodes) => {
+                    for cheatcode in cheatdcodes {
+                        if !lockout && cheatcode.eq(&String::from(RpcCheatcodes::EnableCheatcode)) {
+                            return Err(SurfpoolError::disable_cheatcode(
+                                "Cannot disable surfnet_enableCheatcode rpc method when lockout is not enabledd".to_string(),
+                            )
+                            .into());
+                        }
+                        debug!("disabling cheatcode: {cheatcode}");
+                        if RpcCheatcodes::try_from(cheatcode.as_str()).is_err() {
+                            return Err(SurfpoolError::disable_cheatcode(
+                                "Invalid cheatcode rpc method".to_string(),
+                            )
+                            .into());
+                        }
+
+                        if let Err(e) = cheatcode_ctx.disable_cheatcode(&cheatcode) {
+                            return Err(SurfpoolError::disable_cheatcode(e).into());
+                        }
+                    }
+                }
+            }
+        }
+
+        Ok(RpcResponse {
+            value: (),
+            context: RpcResponseContext::new(svm_locker.get_latest_absolute_slot()),
+        })
+    }
+
+    fn enable_cheatcode(
+        &self,
+        meta: Self::Metadata,
+        cheatcodes_filter: CheatcodeFilter,
+    ) -> Result<RpcResponse<()>> {
+        let svm_locker = match meta.get_svm_locker() {
+            Ok(locker) => locker,
+            Err(e) => return Err(e.into()),
+        };
+        if let Some(runloop_ctx) = meta {
+            let cheatcode_ctx = runloop_ctx.cheatcode_config;
+            match cheatcodes_filter {
+                CheatcodeFilter::All(all) => {
+                    if all.ne("all") {
+                        return Err(SurfpoolError::enable_cheatcode(
+                            "Invalid optioin provided for enabling all cheatcodes. Try using \"all\"".to_string(),
+                        )
+                        .into());
+                    }
+
+                    // we probably don't need to check whether lockout == true because surfnet_enableCheatcode won't be called if it's disabled
+                    cheatcode_ctx.lock().unwrap().filter = CheatcodeFilter::List(vec![]);
+                }
+                CheatcodeFilter::List(cheatcodes) => {
+                    for ref cheatcode in cheatcodes {
+                        debug!("enabling cheatcode: {cheatcode}");
+                        if RpcCheatcodes::try_from(cheatcode.as_str()).is_err() {
+                            return Err(SurfpoolError::enable_cheatcode(
+                                "Invalid cheatcode rpc method".to_string(),
+                            )
+                            .into());
+                        }
+
+                        if let Err(e) = cheatcode_ctx.lock().unwrap().enable_cheatcode(cheatcode) {
+                            return Err(SurfpoolError::enable_cheatcode(e).into());
+                        }
+                    }
+                }
+            }
+        }
+
+        Ok(RpcResponse {
+            value: (),
+            context: RpcResponseContext::new(svm_locker.get_latest_absolute_slot()),
+        })
+    }
+
     fn set_token_account(
         &self,
         meta: Self::Metadata,

crates/core/src/runloops/mod.rs

@@ -1140,6 +1140,7 @@ async fn start_ws_rpc_server_runloop(
                                 .clone(),
                             remote_rpc_client: middleware.remote_rpc_client.clone(),
                             rpc_config: middleware.config.clone(),
+                            cheatcode_config: middleware.cheatcode_config.clone(),
                         };
                         Some(SurfpoolWebsocketMeta::new(
                             runloop_context,

crates/core/src/tests/helpers.rs

@@ -5,7 +5,7 @@ use crossbeam_channel::Sender;
 use solana_clock::Clock;
 use solana_epoch_info::EpochInfo;
 use solana_transaction::versioned::VersionedTransaction;
-use surfpool_types::{RpcConfig, SimnetCommand};
+use surfpool_types::{CheatcodeConfig, RpcConfig, SimnetCommand};
 
 use crate::{
     rpc::RunloopContext,
@@ -67,6 +67,7 @@ where
                 svm_locker: SurfnetSvmLocker::new(surfnet_svm),
                 remote_rpc_client: None,
                 rpc_config: RpcConfig::default(),
+                cheatcode_config: CheatcodeConfig::new(),
             },
             rpc,
         }