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

@@ -15,8 +15,8 @@ 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,
+    GetSurfnetInfoResponse, Idl, ResetAccountConfig, RpcCheatcodes, RpcProfileResultConfig,
+    Scenario, SimnetCommand, SimnetEvent, StreamAccountConfig, UiKeyedProfileResult,
     types::{AccountUpdate, SetSomeAccount, SupplyUpdate, TokenAccountUpdate, UuidOrSignature},
 };
 
@@ -179,6 +179,167 @@ 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: Vec<String>,
+    ) -> 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: Vec<String>,
+    ) -> Result<RpcResponse<()>>;
+
+    /// Enables "lockout" mode so that `surfnet_enableCheatcode` and `surfnet_disableCheatcode` can be disabled.
+    ///
+    /// Once lockout is enabled, those two meta-cheatcode methods may be disabled via `surfnet_disableCheatcode`,
+    /// allowing a session to lock down further modification of the cheatcode set.
+    ///
+    /// ## Parameters
+    /// None.
+    ///
+    /// ## Returns
+    /// A `RpcResponse<()>` indicating success.
+    ///
+    /// ## Example Request
+    /// ```json
+    /// {
+    ///   "jsonrpc": "2.0",
+    ///   "id": 1,
+    ///   "method": "surfnet_lockout",
+    ///   "params": []
+    /// }
+    /// ```
+    ///
+    /// ## Example Response
+    /// ```json
+    /// {
+    ///   "jsonrpc": "2.0",
+    ///   "result": {},
+    ///   "id": 1
+    /// }
+    /// ```
+    ///
+    /// # Notes
+    /// Lockout cannot be turned off once enabled for the current session.
+    ///
+    /// # See Also
+    /// - `surfnet_enableCheatcode`, `surfnet_disableCheatcode`, `surfnet_disableAllCheatcodes`
+    #[rpc(meta, name = "surfnet_lockout")]
+    fn lockout(&self, meta: Self::Metadata) -> Result<RpcResponse<()>>;
+
+    /// Disables all Surfpool cheatcode RPC methods for the current session.
+    ///
+    /// This method turns off every cheatcode at once. After calling it, no `surfnet_*` cheatcode methods
+    /// are callable until re-enabled via `surfnet_enableCheatcode`.
+    ///
+    /// ## Parameters
+    /// None.
+    ///
+    /// ## Returns
+    /// A `RpcResponse<()>` indicating success.
+    ///
+    /// ## Example Request
+    /// ```json
+    /// {
+    ///   "jsonrpc": "2.0",
+    ///   "id": 1,
+    ///   "method": "surfnet_disableAllCheatcodes",
+    ///   "params": []
+    /// }
+    /// ```
+    ///
+    /// ## Example Response
+    /// ```json
+    /// {
+    ///   "jsonrpc": "2.0",
+    ///   "result": {},
+    ///   "id": 1
+    /// }
+    /// ```
+    ///
+    /// # Notes
+    /// Does not require lockout to be enabled; it disables all cheatcodes including
+    /// `surfnet_enableCheatcode` and `surfnet_disableCheatcode`.
+    ///
+    /// # See Also
+    /// - `surfnet_enableCheatcode`, `surfnet_disableCheatcode`, `surfnet_lockout`
+    #[rpc(meta, name = "surfnet_disableAllCheatcodes")]
+    fn disable_all_cheatcodes(&self, meta: Self::Metadata) -> 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 +1362,107 @@ impl SurfnetCheatcodes for SurfnetCheatcodesRpc {
         })
     }
 
+    fn disable_cheatcode(
+        &self,
+        meta: Self::Metadata,
+        cheatcodes: Vec<String>,
+    ) -> 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 mut cheatcode_ctx = runloop_ctx.cheatcode_config.lock().unwrap();
+
+            for cheatcode in cheatcodes {
+                if !cheatcode_ctx.lockout
+                    && (cheatcode.eq(&String::from(RpcCheatcodes::EnableCheatcode))
+                        || cheatcode.eq(&String::from(RpcCheatcodes::DisableCheatcode)))
+                {
+                    return Err(SurfpoolError::disable_cheatcode(
+                        "Cannot disable surfnet_enableCheatcode or surfnet_disableCheatcode 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: Vec<String>,
+    ) -> 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;
+
+            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 disable_all_cheatcodes(&self, meta: Self::Metadata) -> Result<RpcResponse<()>> {
+        if let Some(runloop_ctx) = meta {
+            let cheatcode_ctx = runloop_ctx.cheatcode_config;
+
+            cheatcode_ctx.lock().unwrap().disable_all();
+        }
+
+        Ok(RpcResponse {
+            value: (),
+            context: RpcResponseContext::new(0),
+        })
+    }
+
+    fn lockout(&self, meta: Self::Metadata) -> Result<RpcResponse<()>> {
+        if let Some(runloop_ctx) = meta {
+            let cheatcode_ctx = runloop_ctx.cheatcode_config;
+
+            cheatcode_ctx.lock().unwrap().lockout();
+        }
+
+        Ok(RpcResponse {
+            value: (),
+            context: RpcResponseContext::new(0),
+        })
+    }
+
     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,