diff --git a/crates/core/src/error.rs b/crates/core/src/error.rs index 60de4047..4cb4561a 100644 --- a/crates/core/src/error.rs +++ b/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(pubkey: Pubkey, e: T) -> Self where T: ToString, diff --git a/crates/core/src/rpc/accounts_scan.rs b/crates/core/src/rpc/accounts_scan.rs index 07761e11..d2241489 100644 --- a/crates/core/src/rpc/accounts_scan.rs +++ b/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 = (0..100) .map(|i| match i % 10 { diff --git a/crates/core/src/rpc/mod.rs b/crates/core/src/rpc/mod.rs index ccc2d468..04dbcfc5 100644 --- a/crates/core/src/rpc/mod.rs +++ b/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, pub remote_rpc_client: Option, pub rpc_config: RpcConfig, + pub cheatcode_config: Arc>, } pub struct SurfnetRpcContext { @@ -113,6 +117,7 @@ pub struct SurfpoolMiddleware { pub plugin_manager_commands_tx: Sender, pub config: RpcConfig, pub remote_rpc_client: Option, + pub cheatcode_config: Arc>, } 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> 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> 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() diff --git a/crates/core/src/rpc/surfnet_cheatcodes.rs b/crates/core/src/rpc/surfnet_cheatcodes.rs index 930c83c2..4bd97c21 100644 --- a/crates/core/src/rpc/surfnet_cheatcodes.rs +++ b/crates/core/src/rpc/surfnet_cheatcodes.rs @@ -1,4 +1,7 @@ -use std::collections::BTreeMap; +use std::{ + collections::BTreeMap, + sync::{Arc, RwLock}, +}; use base64::{Engine as _, engine::general_purpose::STANDARD}; use jsonrpc_core::{BoxFuture, Error, Result, futures::future}; @@ -14,9 +17,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, + RpcProfileResultConfig, Scenario, SimnetCommand, SimnetEvent, StreamAccountConfig, + UiKeyedProfileResult, types::{AccountUpdate, SetSomeAccount, SupplyUpdate, TokenAccountUpdate, UuidOrSignature}, }; @@ -179,6 +183,92 @@ pub trait SurfnetCheatcodes { update: AccountUpdate, ) -> BoxFuture>>; + /// 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>; + + /// 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, + ) -> Result>; + /// 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, @@ -1131,7 +1221,23 @@ pub trait SurfnetCheatcodes { } #[derive(Clone)] -pub struct SurfnetCheatcodesRpc; +pub struct SurfnetCheatcodesRpc { + pub registered_methods: Arc>>, +} +impl SurfnetCheatcodesRpc { + pub fn empty() -> Self { + Self { + registered_methods: Arc::new(RwLock::new(vec![])), + } + } + pub fn is_available_cheatcode(&self, cheatcode: &String) -> bool { + let Ok(methods) = self.registered_methods.read() else { + return false; + }; + methods.contains(cheatcode) + } +} + impl SurfnetCheatcodes for SurfnetCheatcodesRpc { type Metadata = Option; @@ -1201,6 +1307,118 @@ impl SurfnetCheatcodes for SurfnetCheatcodesRpc { }) } + fn disable_cheatcode( + &self, + meta: Self::Metadata, + cheatcodes_filter: CheatcodeFilter, + control_config: Option, + ) -> Result> { + 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 Ok(mut cheatcode_ctx) = runloop_ctx.cheatcode_config.lock() else { + return Err(jsonrpc_core::Error::internal_error()); + }; + + match cheatcodes_filter { + CheatcodeFilter::All(all) => { + if all.ne("all") { + return Err(SurfpoolError::disable_cheatcode( + "Invalid option provided for disabling all cheatcodes. Try using 'all' or providing an array of specific cheatcodes".to_string(), + ) + .into()); + } + + let Ok(available_cheatcodes) = self.registered_methods.read() else { + return Err(jsonrpc_core::Error::internal_error()); + }; + cheatcode_ctx.disable_all(lockout, (*available_cheatcodes).clone()); + } + CheatcodeFilter::List(cheatdcodes) => { + for cheatcode in cheatdcodes { + if !lockout && cheatcode.eq("surfnet_enableCheatcode") { + return Err(SurfpoolError::disable_cheatcode( + "Cannot disable surfnet_enableCheatcode rpc method when lockout is not enabled".to_string(), + ) + .into()); + } + debug!("disabling cheatcode: {cheatcode}"); + if !self.is_available_cheatcode(&cheatcode) { + 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> { + let svm_locker = match meta.get_svm_locker() { + Ok(locker) => locker, + Err(e) => return Err(e.into()), + }; + if let Some(runloop_ctx) = meta { + let Ok(mut cheatcode_ctx) = runloop_ctx.cheatcode_config.lock() else { + return Err(jsonrpc_core::Error::internal_error()); + }; + match cheatcodes_filter { + CheatcodeFilter::All(all) => { + if all.ne("all") { + return Err(SurfpoolError::enable_cheatcode( + "Invalid option provided for enabling all cheatcodes. Try using 'all' or providing an array of specific cheatcodes".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.filter = CheatcodeFilter::List(vec![]); + } + CheatcodeFilter::List(cheatcodes) => { + for ref cheatcode in cheatcodes { + debug!("enabling cheatcode: {cheatcode}"); + if !self.is_available_cheatcode(cheatcode) { + return Err(SurfpoolError::enable_cheatcode( + "Invalid cheatcode rpc method".to_string(), + ) + .into()); + } + + if let Err(e) = cheatcode_ctx.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, @@ -1878,7 +2096,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_get_transaction_profile() { // Create connection to local validator - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let recent_blockhash = client .context .svm_locker @@ -2711,7 +2929,7 @@ mod tests { #[test] fn test_export_snapshot() { - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let pubkey1 = Pubkey::new_unique(); let account1 = Account { @@ -2747,7 +2965,7 @@ mod tests { #[test] fn test_export_snapshot_json_parsed() { - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let pubkey1 = Pubkey::new_unique(); println!("Pubkey1: {}", pubkey1); @@ -2843,7 +3061,7 @@ mod tests { use solana_signature::Signature; use surfpool_types::{ProfileResult, types::KeyedProfileResult}; - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); // Create several accounts in the network let account1_pubkey = Pubkey::new_unique(); @@ -2988,7 +3206,7 @@ mod tests { included_program_account_pubkey ); - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let system_account = Account { lamports: 1_000_000, @@ -3083,7 +3301,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_write_program_creates_accounts_automatically() { // Test that both program and program data accounts are created if they don't exist - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let program_id = Keypair::new(); // Verify accounts don't exist initially @@ -3161,7 +3379,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_write_program_single_chunk_small() { // Test writing a small program in a single write - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let program_id = Keypair::new(); let program_data = vec![0x01, 0x02, 0x03, 0x04, 0x05]; @@ -3210,7 +3428,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_write_program_single_chunk_large() { // Test writing a large program (1MB) in a single write - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let program_id = Keypair::new(); let program_data = vec![0xAB; 1024 * 1024]; // 1 MB @@ -3259,7 +3477,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_write_program_multiple_sequential_chunks() { // Test writing a program in multiple sequential chunks - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let program_id = Keypair::new(); let chunks = vec![ @@ -3324,7 +3542,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_write_program_non_sequential_chunks() { // Test writing chunks out of order (backwards) - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let program_id = Keypair::new(); let chunks = vec![ @@ -3390,7 +3608,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_write_program_overlapping_writes() { // Test that overlapping writes correctly overwrite previous data - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let program_id = Keypair::new(); // Write initial data @@ -3458,7 +3676,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_write_program_zero_offset() { // Test writing at offset 0 (should write immediately after metadata) - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let program_id = Keypair::new(); let data = vec![0x42; 128]; @@ -3496,7 +3714,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_write_program_large_offset() { // Test writing at a large offset (account should expand) - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let program_id = Keypair::new(); let large_offset = 1024 * 1024; // 1 MB offset @@ -3545,7 +3763,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_write_program_empty_data() { // Test writing empty data (should succeed but not write anything) - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let program_id = Keypair::new(); let result = client @@ -3567,7 +3785,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_write_program_single_byte() { // Test writing a single byte - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let program_id = Keypair::new(); let data = vec![0x42]; @@ -3605,7 +3823,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_write_program_invalid_program_id() { // Test with invalid program ID - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let result = client .rpc @@ -3630,7 +3848,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_write_program_invalid_hex_data() { // Test with invalid hex encoding - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let program_id = Keypair::new(); let invalid_hex_strings = vec![ @@ -3669,7 +3887,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_write_program_rent_exemption() { // Test that rent exemption is maintained when account expands - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let program_id = Keypair::new(); // Write initial small data @@ -3752,7 +3970,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_write_program_account_ownership() { // Test that created accounts have correct ownership - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let program_id = Keypair::new(); let authority = Keypair::new(); @@ -3832,7 +4050,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_write_program_metadata_preservation() { // Test that program data account metadata is preserved across writes - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let program_id = Keypair::new(); // First write @@ -3899,7 +4117,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_write_program_idempotent() { // Test that writing the same data twice produces the same result - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let program_id = Keypair::new(); let data = vec![0x55; 512]; @@ -3965,7 +4183,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] async fn test_write_program_context_slot() { // Test that response context contains valid slot - let client = TestSetup::new(SurfnetCheatcodesRpc); + let client = TestSetup::new(SurfnetCheatcodesRpc::empty()); let program_id = Keypair::new(); let result = client diff --git a/crates/core/src/runloops/mod.rs b/crates/core/src/runloops/mod.rs index 0a4252cc..4e5c006a 100644 --- a/crates/core/src/runloops/mod.rs +++ b/crates/core/src/runloops/mod.rs @@ -1053,12 +1053,28 @@ async fn start_http_rpc_server_runloop( .map_err(|e| e.to_string())?; let mut io = MetaIoHandler::with_middleware(middleware); + + // Cheatcodes should be added first and are a special case. One of the cheatcode methods needs access + // to the list of all cheatcode methods. The IoHandler allows us to iterate over them, so we're + // initializing and storing that list here. + { + let cheatcode_methods: Arc>> = Arc::new(RwLock::new(vec![])); + let mut cheatcodes_impl = rpc::surfnet_cheatcodes::SurfnetCheatcodesRpc { + registered_methods: Arc::clone(&cheatcode_methods), + }; + io.extend_with(cheatcodes_impl.to_delegate()); + + cheatcode_methods + .write() + .unwrap() + .extend(io.iter().map(|(n, _)| n.clone()).collect::>()); + } + io.extend_with(rpc::minimal::SurfpoolMinimalRpc.to_delegate()); io.extend_with(rpc::full::SurfpoolFullRpc.to_delegate()); io.extend_with(rpc::accounts_data::SurfpoolAccountsDataRpc.to_delegate()); io.extend_with(rpc::accounts_scan::SurfpoolAccountsScanRpc.to_delegate()); io.extend_with(rpc::bank_data::SurfpoolBankDataRpc.to_delegate()); - io.extend_with(rpc::surfnet_cheatcodes::SurfnetCheatcodesRpc.to_delegate()); io.extend_with(rpc::admin::SurfpoolAdminRpc.to_delegate()); if !config.plugin_config_path.is_empty() { @@ -1140,6 +1156,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, diff --git a/crates/core/src/tests/helpers.rs b/crates/core/src/tests/helpers.rs index 8cf3061f..7008c2f7 100644 --- a/crates/core/src/tests/helpers.rs +++ b/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, } diff --git a/crates/core/src/tests/integration.rs b/crates/core/src/tests/integration.rs index c2f635ae..bf60b205 100644 --- a/crates/core/src/tests/integration.rs +++ b/crates/core/src/tests/integration.rs @@ -32,8 +32,9 @@ use solana_system_interface::{ }; use solana_transaction::{Transaction, versioned::VersionedTransaction}; use surfpool_types::{ - DEFAULT_SLOT_TIME_MS, Idl, RpcProfileDepth, RpcProfileResultConfig, SimnetCommand, SimnetEvent, - SurfpoolConfig, UiAccountChange, UiAccountProfileState, UiKeyedProfileResult, + CheatcodeConfig, CheatcodeControlConfig, CheatcodeFilter, DEFAULT_SLOT_TIME_MS, Idl, + RpcProfileDepth, RpcProfileResultConfig, SimnetCommand, SimnetEvent, SurfpoolConfig, + UiAccountChange, UiAccountProfileState, UiKeyedProfileResult, types::{ BlockProductionMode, RpcConfig, SimnetConfig, SubgraphConfig, TransactionStatusEvent, UuidOrSignature, @@ -763,7 +764,7 @@ async fn test_simulate_transaction_no_signers(test_type: TestType) { #[tokio::test(flavor = "multi_thread")] async fn test_surfnet_estimate_compute_units(test_type: TestType) { let (mut svm_instance, _simnet_events_rx, _geyser_events_rx) = test_type.initialize_svm(); - let rpc_server = crate::rpc::surfnet_cheatcodes::SurfnetCheatcodesRpc; + let rpc_server = crate::rpc::surfnet_cheatcodes::SurfnetCheatcodesRpc::empty(); let payer = Keypair::new(); let recipient = Pubkey::new_unique(); @@ -796,6 +797,7 @@ async fn test_surfnet_estimate_compute_units(test_type: TestType) { plugin_manager_commands_tx: plugin_cmd_tx, remote_rpc_client: None, rpc_config: RpcConfig::default(), + cheatcode_config: CheatcodeConfig::new(), }; // Test with None tag @@ -1048,13 +1050,280 @@ async fn test_surfnet_estimate_compute_units(test_type: TestType) { assert!(found_cu_event, "Did not find CU estimation SimnetEvent"); } +#[test_case(TestType::sqlite(); "with on-disk sqlite db")] +#[test_case(TestType::in_memory(); "with in-memory sqlite db")] +#[test_case(TestType::no_db(); "with no db")] +#[cfg_attr(feature = "postgres", test_case(TestType::postgres(); "with postgres db"))] +fn test_enable_and_disable_cheatcodes(test_type: TestType) { + let valid_cheatcode_method = "surfnet_getActiveIdl".to_string(); + let enable_cheatcode_method = "surfnet_enableCheatcode".to_string(); + let disable_cheatcode_method = "surfnet_disableCheatcode".to_string(); + let invalid_cheatcode_method = "surfnet_invalidCheatcode".to_string(); + let available_methods = vec![ + valid_cheatcode_method.clone(), + enable_cheatcode_method.clone(), + disable_cheatcode_method.clone(), + ]; + let rpc_server = SurfnetCheatcodesRpc { + registered_methods: Arc::new(std::sync::RwLock::new(available_methods.clone())), + }; + let (svm_instance, _simnet_events_rx, _geyser_events_rx) = test_type.initialize_svm(); + let svm_locker_for_context = SurfnetSvmLocker::new(svm_instance); + let (simnet_cmd_tx, _simnet_cmd_rx) = crossbeam_unbounded::(); + let (plugin_cmd_tx, _plugin_cmd_rx) = crossbeam_unbounded::(); + + let runloop_context = RunloopContext { + id: None, + svm_locker: svm_locker_for_context.clone(), + simnet_commands_tx: simnet_cmd_tx, + plugin_manager_commands_tx: plugin_cmd_tx, + remote_rpc_client: None, + rpc_config: RpcConfig::default(), + cheatcode_config: CheatcodeConfig::new(), + }; + + // Test disable works properly + let disable_cheatcode_pass_result = rpc_server.disable_cheatcode( + Some(runloop_context.clone()), + CheatcodeFilter::List(vec![valid_cheatcode_method.clone()]), + None, + ); + + assert!( + disable_cheatcode_pass_result.is_ok(), + "Expected surfnet_disableCheatcode to pass" + ); + + assert!( + runloop_context + .cheatcode_config + .lock() + .unwrap() + .is_cheatcode_disabled(&valid_cheatcode_method), + "Expected cheatcode to be disabled" + ); + + let disable_enable_cheatcode_lockout_false_fails = rpc_server.disable_cheatcode( + Some(runloop_context.clone()), + CheatcodeFilter::List(vec![enable_cheatcode_method.clone()]), + None, + ); + + let expected_error: jsonrpc_core::Result<()> = Err(SurfpoolError::disable_cheatcode( + "Cannot disable surfnet_enableCheatcode rpc method when lockout is not enabled".to_string(), + ) + .into()); + + assert!( + disable_enable_cheatcode_lockout_false_fails.is_err(), + "Expected surfnet_disableCheatcode to fail when disabling surfnet_enableCheatcode with lockout == false" + ); + assert_eq!( + disable_enable_cheatcode_lockout_false_fails.err().unwrap(), + expected_error.err().unwrap(), + "Expected error did not match the resulting error" + ); + + let disable_cheatcode_fails_on_invalid_cheatcode = rpc_server.disable_cheatcode( + Some(runloop_context.clone()), + CheatcodeFilter::List(vec![invalid_cheatcode_method.clone()]), + None, + ); + + let expected_error_disable_invalid_cheatcode: jsonrpc_core::Result<()> = + Err(SurfpoolError::disable_cheatcode("Invalid cheatcode rpc method".to_string()).into()); + + assert!( + disable_cheatcode_fails_on_invalid_cheatcode.is_err(), + "Expected surfnet_disableCheatcode to fail on providing invalid cheatcode method" + ); + + assert_eq!( + expected_error_disable_invalid_cheatcode.err().unwrap(), + disable_cheatcode_fails_on_invalid_cheatcode.err().unwrap(), + "Resulting error does not match the expected error" + ); + + let disable_all_cheatcodes_fails_with_invalid_config = rpc_server.disable_cheatcode( + Some(runloop_context.clone()), + CheatcodeFilter::All("not_all".to_string()), + None, + ); + let expected_error_disable_all_with_invalid_config: jsonrpc_core::Result<()> = + Err(SurfpoolError::disable_cheatcode( + "Invalid option provided for disabling all cheatcodes. Try using 'all' or providing an array of specific cheatcodes".to_string(), + ) + .into()); + + assert!( + disable_all_cheatcodes_fails_with_invalid_config.is_err(), + "Expected surfnet_disableCheatcode to fail when config is wrong" + ); + assert_eq!( + disable_all_cheatcodes_fails_with_invalid_config + .err() + .unwrap(), + expected_error_disable_all_with_invalid_config + .err() + .unwrap(), + "Expected error did not match the resulting error" + ); + + let disabled_all_cheatcode_passes_result = rpc_server.disable_cheatcode( + Some(runloop_context.clone()), + CheatcodeFilter::All("all".to_string()), + None, + ); + + assert!( + disabled_all_cheatcode_passes_result.is_ok(), + "Expected surfnet_disableCheatcode to pass with valid config" + ); + + assert_eq!( + runloop_context.cheatcode_config.lock().unwrap().filter, + CheatcodeConfig::filter_all_list(false, available_methods.clone()), + "The disabled cheatcodes list doesn't match the expected one" + ); + + let disable_fails_if_cheatcode_already_disabled_result = rpc_server.disable_cheatcode( + Some(runloop_context.clone()), + CheatcodeFilter::List(vec![valid_cheatcode_method.clone()]), + None, + ); + + let expected_error_if_cheatcode_already_disabled: jsonrpc_core::Result<()> = + Err(SurfpoolError::disable_cheatcode("Cheatcode already disabled".to_string()).into()); + + assert!( + disable_fails_if_cheatcode_already_disabled_result.is_err(), + "Expected surfnet_disableCheatcode to fail if cheatcode already disabled" + ); + + assert_eq!( + disable_fails_if_cheatcode_already_disabled_result + .err() + .unwrap(), + expected_error_if_cheatcode_already_disabled.err().unwrap(), + "The expected error does not match the resulting error" + ); + + // test enable works properly + let enable_cheatcode_works_properly_result = rpc_server.enable_cheatcode( + Some(runloop_context.clone()), + CheatcodeFilter::List(vec![valid_cheatcode_method.clone()]), + ); + + assert!( + enable_cheatcode_works_properly_result.is_ok(), + "expected surfnet_enableCheatcode to pass" + ); + assert!( + !runloop_context + .cheatcode_config + .lock() + .unwrap() + .is_cheatcode_disabled(&valid_cheatcode_method), + "Expected the cheatcode to be enabled after surnet_enableCheatcode rpc method" + ); + + let enable_cheatcode_fails_on_invalid_cheatcode = rpc_server.enable_cheatcode( + Some(runloop_context.clone()), + CheatcodeFilter::List(vec![invalid_cheatcode_method]), + ); + + let expected_error_enable_invalid_cheatcode: jsonrpc_core::Result<()> = + Err(SurfpoolError::enable_cheatcode("Invalid cheatcode rpc method".to_string()).into()); + + assert!( + enable_cheatcode_fails_on_invalid_cheatcode.is_err(), + "Expected surfnet_disableCheatcode to fail on providing invalid cheatcode method" + ); + + assert_eq!( + enable_cheatcode_fails_on_invalid_cheatcode.err().unwrap(), + expected_error_enable_invalid_cheatcode.err().unwrap(), + "Resulting error does not match the expected error" + ); + + let enable_all_cheatcodes_fails_with_invalid_config = rpc_server.enable_cheatcode( + Some(runloop_context.clone()), + CheatcodeFilter::All("not_all".to_string()), + ); + let expected_error_enable_all_with_invalid_config: jsonrpc_core::Result<()> = + Err(SurfpoolError::enable_cheatcode( + "Invalid option provided for enabling all cheatcodes. Try using 'all' or providing an array of specific cheatcodes".to_string(), + ) + .into()); + + assert!( + enable_all_cheatcodes_fails_with_invalid_config.is_err(), + "Expected surfnet_disableCheatcode to fail when config is wrong" + ); + assert_eq!( + enable_all_cheatcodes_fails_with_invalid_config + .err() + .unwrap(), + expected_error_enable_all_with_invalid_config.err().unwrap(), + "Expected error did not match the resulting error" + ); + + let enable_all_cheatcodes_pass_result = rpc_server.enable_cheatcode( + Some(runloop_context.clone()), + CheatcodeFilter::All("all".to_string()), + ); + + assert!( + enable_all_cheatcodes_pass_result.is_ok(), + "Expected surnet_enableCheatcode with correct config to pass" + ); + + assert_eq!( + runloop_context.cheatcode_config.lock().unwrap().filter, + CheatcodeFilter::List(vec![]), + "Expected all the cheatcodes to be enabled" + ); + + let disable_all_with_lockout_passes = rpc_server.disable_cheatcode( + Some(runloop_context.clone()), + CheatcodeFilter::All("all".to_string()), + Some(CheatcodeControlConfig { + lockout: Some(true), + }), + ); + + assert!( + disable_all_with_lockout_passes.is_ok(), + "Expected surfnet_disableCheatcode with lockout == true to pass" + ); + + assert_eq!( + runloop_context.cheatcode_config.lock().unwrap().filter, + CheatcodeConfig::filter_all_list(true, vec![]), + "Expected all the features to be disabled" + ); + + // assert that surfnet_enableCheatcode and surfnet_disableCheatcode both are disabled so that on_request Middleware sucessfully filters them + for ref cheatcode in [enable_cheatcode_method, disable_cheatcode_method] { + assert!( + runloop_context + .cheatcode_config + .lock() + .unwrap() + .is_cheatcode_disabled(cheatcode), + "Expected {} to be disabled", + cheatcode + ); + } +} + #[test_case(TestType::sqlite(); "with on-disk sqlite db")] #[test_case(TestType::in_memory(); "with in-memory sqlite db")] #[test_case(TestType::no_db(); "with no db")] #[cfg_attr(feature = "postgres", test_case(TestType::postgres(); "with postgres db"))] #[tokio::test(flavor = "multi_thread")] async fn test_get_transaction_profile(test_type: TestType) { - let rpc_server = SurfnetCheatcodesRpc; + let rpc_server = SurfnetCheatcodesRpc::empty(); let (mut svm_instance, _simnet_events_rx, _geyser_events_rx) = test_type.initialize_svm(); // Set up test accounts @@ -1090,6 +1359,7 @@ async fn test_get_transaction_profile(test_type: TestType) { plugin_manager_commands_tx: plugin_cmd_tx, remote_rpc_client: None, rpc_config: RpcConfig::default(), + cheatcode_config: CheatcodeConfig::new(), }; // Test 1: Profile a transaction with a tag and retrieve by UUID @@ -1275,7 +1545,7 @@ async fn test_get_transaction_profile(test_type: TestType) { #[cfg_attr(feature = "postgres", test_case(TestType::postgres(); "with postgres db"))] fn test_register_and_get_idl_without_slot(test_type: TestType) { let idl: Idl = serde_json::from_slice(include_bytes!("./assets/idl_v1.json")).unwrap(); - let rpc_server = SurfnetCheatcodesRpc; + let rpc_server = SurfnetCheatcodesRpc::empty(); let (svm_instance, _simnet_events_rx, _geyser_events_rx) = test_type.initialize_svm(); let svm_locker_for_context = SurfnetSvmLocker::new(svm_instance); @@ -1289,6 +1559,7 @@ fn test_register_and_get_idl_without_slot(test_type: TestType) { plugin_manager_commands_tx: plugin_cmd_tx, remote_rpc_client: None, rpc_config: RpcConfig::default(), + cheatcode_config: CheatcodeConfig::new(), }; // Test 1: Register IDL without slot @@ -1330,7 +1601,7 @@ fn test_register_and_get_idl_without_slot(test_type: TestType) { #[cfg_attr(feature = "postgres", test_case(TestType::postgres(); "with postgres db"))] fn test_register_and_get_idl_with_slot(test_type: TestType) { let idl: Idl = serde_json::from_slice(include_bytes!("./assets/idl_v1.json")).unwrap(); - let rpc_server = SurfnetCheatcodesRpc; + let rpc_server = SurfnetCheatcodesRpc::empty(); let (svm_instance, _simnet_events_rx, _geyser_events_rx) = test_type.initialize_svm(); let svm_locker_for_context = SurfnetSvmLocker::new(svm_instance); @@ -1344,6 +1615,7 @@ fn test_register_and_get_idl_with_slot(test_type: TestType) { plugin_manager_commands_tx: plugin_cmd_tx, remote_rpc_client: None, rpc_config: RpcConfig::default(), + cheatcode_config: CheatcodeConfig::new(), }; // Test 1: Register IDL with slot @@ -1398,7 +1670,7 @@ async fn test_register_and_get_same_idl_with_different_slots(test_type: TestType let idl_v1: Idl = serde_json::from_slice(include_bytes!("./assets/idl_v1.json")).unwrap(); let idl_v2: Idl = serde_json::from_slice(include_bytes!("./assets/idl_v2.json")).unwrap(); let idl_v3: Idl = serde_json::from_slice(include_bytes!("./assets/idl_v3.json")).unwrap(); - let rpc_server = SurfnetCheatcodesRpc; + let rpc_server = SurfnetCheatcodesRpc::empty(); let (svm_instance, _simnet_events_rx, _geyser_events_rx) = test_type.initialize_svm(); let svm_locker_for_context = SurfnetSvmLocker::new(svm_instance); @@ -1425,6 +1697,7 @@ async fn test_register_and_get_same_idl_with_different_slots(test_type: TestType plugin_manager_commands_tx: plugin_cmd_tx, remote_rpc_client: None, rpc_config: RpcConfig::default(), + cheatcode_config: CheatcodeConfig::new(), }; // Step 1: Register IDL v1 at slot_1 @@ -3090,7 +3363,7 @@ async fn test_profile_transaction_versioned_message(test_type: TestType) { #[cfg_attr(feature = "postgres", test_case(TestType::postgres(); "with postgres db"))] #[tokio::test(flavor = "multi_thread")] async fn test_get_local_signatures_without_limit(test_type: TestType) { - let rpc_server = SurfnetCheatcodesRpc; + let rpc_server = SurfnetCheatcodesRpc::empty(); let (svm_instance, _simnet_events_rx, _geyser_events_rx) = test_type.initialize_svm(); let svm_locker_for_context = SurfnetSvmLocker::new(svm_instance.clone()); @@ -3105,6 +3378,7 @@ async fn test_get_local_signatures_without_limit(test_type: TestType) { plugin_manager_commands_tx: plugin_cmd_tx, remote_rpc_client: None, rpc_config: RpcConfig::default(), + cheatcode_config: CheatcodeConfig::new(), }; let payer = Keypair::new(); @@ -3195,7 +3469,7 @@ async fn test_get_local_signatures_without_limit(test_type: TestType) { #[cfg_attr(feature = "postgres", test_case(TestType::postgres(); "with postgres db"))] #[tokio::test(flavor = "multi_thread")] async fn test_get_local_signatures_with_limit(test_type: TestType) { - let rpc_server = SurfnetCheatcodesRpc; + let rpc_server = SurfnetCheatcodesRpc::empty(); let (svm_instance, _simnet_events_rx, _geyser_events_rx) = test_type.initialize_svm(); let svm_locker_for_context = SurfnetSvmLocker::new(svm_instance.clone()); @@ -3209,6 +3483,7 @@ async fn test_get_local_signatures_with_limit(test_type: TestType) { plugin_manager_commands_tx: plugin_cmd_tx, remote_rpc_client: None, rpc_config: RpcConfig::default(), + cheatcode_config: CheatcodeConfig::new(), }; let payer = Keypair::new(); @@ -3404,7 +3679,7 @@ fn boot_simnet( #[test_case(TestType::no_db(); "with no db")] #[cfg_attr(feature = "postgres", test_case(TestType::postgres(); "with postgres db"))] fn test_time_travel_resume_paused_clock(test_type: TestType) { - let rpc_server = SurfnetCheatcodesRpc; + let rpc_server = SurfnetCheatcodesRpc::empty(); let (svm_locker, simnet_cmd_tx, _) = boot_simnet(BlockProductionMode::Clock, Some(100), test_type); let (plugin_cmd_tx, _plugin_cmd_rx) = crossbeam_unbounded::(); @@ -3416,6 +3691,7 @@ fn test_time_travel_resume_paused_clock(test_type: TestType) { plugin_manager_commands_tx: plugin_cmd_tx, remote_rpc_client: None, rpc_config: RpcConfig::default(), + cheatcode_config: CheatcodeConfig::new(), }; // Get initial epoch info @@ -3481,7 +3757,7 @@ fn test_time_travel_resume_paused_clock(test_type: TestType) { #[test_case(TestType::no_db(); "with no db")] #[cfg_attr(feature = "postgres", test_case(TestType::postgres(); "with postgres db"))] fn test_time_travel_absolute_timestamp(test_type: TestType) { - let rpc_server = SurfnetCheatcodesRpc; + let rpc_server = SurfnetCheatcodesRpc::empty(); let slot_time = 100; let (svm_locker, simnet_cmd_tx, simnet_events_rx) = boot_simnet( BlockProductionMode::Clock, @@ -3497,6 +3773,7 @@ fn test_time_travel_absolute_timestamp(test_type: TestType) { plugin_manager_commands_tx: plugin_cmd_tx, remote_rpc_client: None, rpc_config: RpcConfig::default(), + cheatcode_config: CheatcodeConfig::new(), }; let clock = Clock { @@ -3567,7 +3844,7 @@ fn test_time_travel_absolute_timestamp(test_type: TestType) { #[test_case(TestType::no_db(); "with no db")] #[cfg_attr(feature = "postgres", test_case(TestType::postgres(); "with postgres db"))] fn test_time_travel_absolute_slot(test_type: TestType) { - let rpc_server = SurfnetCheatcodesRpc; + let rpc_server = SurfnetCheatcodesRpc::empty(); let (svm_locker, simnet_cmd_tx, simnet_events_rx) = boot_simnet(BlockProductionMode::Clock, Some(400), test_type); let (plugin_cmd_tx, _plugin_cmd_rx) = crossbeam_unbounded::(); @@ -3579,6 +3856,7 @@ fn test_time_travel_absolute_slot(test_type: TestType) { plugin_manager_commands_tx: plugin_cmd_tx, remote_rpc_client: None, rpc_config: RpcConfig::default(), + cheatcode_config: CheatcodeConfig::new(), }; let clock = Clock { @@ -3643,7 +3921,7 @@ fn test_time_travel_absolute_slot(test_type: TestType) { #[test_case(TestType::no_db(); "with no db")] #[cfg_attr(feature = "postgres", test_case(TestType::postgres(); "with postgres db"))] fn test_time_travel_absolute_epoch(test_type: TestType) { - let rpc_server = SurfnetCheatcodesRpc; + let rpc_server = SurfnetCheatcodesRpc::empty(); let (svm_locker, simnet_cmd_tx, simnet_events_rx) = boot_simnet(BlockProductionMode::Clock, Some(400), test_type); let (plugin_cmd_tx, _plugin_cmd_rx) = crossbeam_unbounded::(); @@ -3655,6 +3933,7 @@ fn test_time_travel_absolute_epoch(test_type: TestType) { plugin_manager_commands_tx: plugin_cmd_tx, remote_rpc_client: None, rpc_config: RpcConfig::default(), + cheatcode_config: CheatcodeConfig::new(), }; let clock = Clock { @@ -4376,7 +4655,7 @@ fn test_reset_network_keeps_latest_blockhash_valid(test_type: TestType) { #[test_case(TestType::no_db(); "with no db")] #[cfg_attr(feature = "postgres", test_case(TestType::postgres(); "with postgres db"))] fn test_reset_network_time_travel_timestamp(test_type: TestType) { - let rpc_server = SurfnetCheatcodesRpc; + let rpc_server = SurfnetCheatcodesRpc::empty(); let (svm_locker, simnet_cmd_tx, simnet_events_rx) = boot_simnet(BlockProductionMode::Clock, Some(400), test_type); let (plugin_cmd_tx, _plugin_cmd_rx) = crossbeam_unbounded::(); @@ -4388,6 +4667,7 @@ fn test_reset_network_time_travel_timestamp(test_type: TestType) { plugin_manager_commands_tx: plugin_cmd_tx, remote_rpc_client: None, rpc_config: RpcConfig::default(), + cheatcode_config: CheatcodeConfig::new(), }; // Calculate a target timestamp in the future @@ -4428,7 +4708,7 @@ fn test_reset_network_time_travel_timestamp(test_type: TestType) { #[test_case(TestType::no_db(); "with no db")] #[cfg_attr(feature = "postgres", test_case(TestType::postgres(); "with postgres db"))] fn test_reset_network_time_travel_slot(test_type: TestType) { - let rpc_server = SurfnetCheatcodesRpc; + let rpc_server = SurfnetCheatcodesRpc::empty(); let (svm_locker, simnet_cmd_tx, simnet_events_rx) = boot_simnet(BlockProductionMode::Clock, Some(400), test_type); let (plugin_cmd_tx, _plugin_cmd_rx) = crossbeam_unbounded::(); @@ -4440,6 +4720,7 @@ fn test_reset_network_time_travel_slot(test_type: TestType) { plugin_manager_commands_tx: plugin_cmd_tx, remote_rpc_client: None, rpc_config: RpcConfig::default(), + cheatcode_config: CheatcodeConfig::new(), }; // Do an initial reset to ensure we start from slot 0 @@ -4484,7 +4765,7 @@ fn test_reset_network_time_travel_slot(test_type: TestType) { #[test_case(TestType::no_db(); "with no db")] #[cfg_attr(feature = "postgres", test_case(TestType::postgres(); "with postgres db"))] fn test_reset_network_time_travel_epoch(test_type: TestType) { - let rpc_server = SurfnetCheatcodesRpc; + let rpc_server = SurfnetCheatcodesRpc::empty(); let (svm_locker, simnet_cmd_tx, simnet_events_rx) = boot_simnet(BlockProductionMode::Clock, Some(400), test_type); let (plugin_cmd_tx, _plugin_cmd_rx) = crossbeam_unbounded::(); @@ -4496,6 +4777,7 @@ fn test_reset_network_time_travel_epoch(test_type: TestType) { plugin_manager_commands_tx: plugin_cmd_tx, remote_rpc_client: None, rpc_config: RpcConfig::default(), + cheatcode_config: CheatcodeConfig::new(), }; // Do an initial reset to ensure we start from epoch 0 @@ -6939,7 +7221,7 @@ fn test_nonce_accounts() { #[tokio::test(flavor = "multi_thread")] async fn test_profile_transaction_does_not_mutate_state(test_type: TestType) { let (mut svm_instance, _simnet_events_rx, _geyser_events_rx) = test_type.initialize_svm(); - let rpc_server = crate::rpc::surfnet_cheatcodes::SurfnetCheatcodesRpc; + let rpc_server = crate::rpc::surfnet_cheatcodes::SurfnetCheatcodesRpc::empty(); // Setup: Create accounts and fund the payer let payer = Keypair::new(); @@ -6980,6 +7262,7 @@ async fn test_profile_transaction_does_not_mutate_state(test_type: TestType) { plugin_manager_commands_tx: plugin_cmd_tx, remote_rpc_client: None, rpc_config: RpcConfig::default(), + cheatcode_config: CheatcodeConfig::new(), }; // Profile the transaction multiple times to ensure no state leakage diff --git a/crates/types/src/types.rs b/crates/types/src/types.rs index 7cdc3f8a..1e311181 100644 --- a/crates/types/src/types.rs +++ b/crates/types/src/types.rs @@ -4,6 +4,7 @@ use std::{ fmt, path::PathBuf, str::FromStr, + sync::{Arc, Mutex}, }; use blake3::Hash; @@ -1257,6 +1258,96 @@ impl RunbookExecutionStatusReport { } } +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CheatcodeConfig { + pub lockout: bool, // if true, allows disabling even the `surfnet_enableCheatcodes`/`surfnetdisableCheatcodes` methods + pub filter: CheatcodeFilter, +} + +#[derive(Serialize, Deserialize, Default)] +pub struct CheatcodeControlConfig { + pub lockout: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum CheatcodeFilter { + All(String), + List(Vec), // disables cheatcodes in a named list +} + +impl CheatcodeConfig { + pub fn new() -> Arc> { + Arc::new(Mutex::new(CheatcodeConfig { + lockout: false, + filter: CheatcodeFilter::List(vec![]), + })) + } + + pub fn lockout(&mut self) { + self.lockout = true; + } + + pub fn disable_all(&mut self, lockout: bool, available_cheatcodes: Vec) { + self.filter = Self::filter_all_list(lockout, available_cheatcodes); + } + + pub fn disable_cheatcode(&mut self, cheatcode: &String) -> Result<(), String> { + if !self.lockout + && (cheatcode.eq("surfpool_enableCheatcode") + || cheatcode.eq("surfpool_disableCheatcode")) + { + return Err("Cannot disable surfpool_disableCheatcode or surfpool_enableCheatcode while lockout is is false".to_string()); + } + + if let CheatcodeFilter::List(list) = &mut self.filter { + if !list.contains(cheatcode) { + list.push(cheatcode.to_string()); + Ok(()) + } else { + Err("Cheatcode already disabled".to_string()) + } + } else { + Err("All cheatcodes disabled".to_string()) + } + } + pub fn enable_cheatcode(&mut self, cheatcode: &str) -> Result<(), String> { + if let CheatcodeFilter::List(list) = &mut self.filter { + if let Some(pos) = list.iter().position(|c| c == cheatcode) { + list.remove(pos); + Ok(()) + } else { + Err("Cheatcode isn't disabled".to_string()) + } + } else { + Err("All cheatcodes are disabled".to_string()) + } + } + + pub fn is_cheatcode_disabled(&self, cheatcode: &String) -> bool { + match &self.filter { + CheatcodeFilter::List(list) => list.contains(cheatcode), + CheatcodeFilter::All(_) => true, + } + } + + pub fn filter_all_list(lockout: bool, available_cheatcodes: Vec) -> CheatcodeFilter { + // when lockout == true, it's important to disable surfnet_disableCheatcode as well + // since calling surfnet_disableCheatcode with lockout == false will override the current config, which is a bug + if lockout { + CheatcodeFilter::All("all".to_string()) + } else { + // remove `surfnet_disableCheatcode` and `surfnet_enableCheatcode` from the list of available cheatcodes + let filter = available_cheatcodes + .into_iter() + .filter(|c| (c.ne("surfnet_disableCheatcode") && c.ne("surfnet_enableCheatcode"))) + .collect(); + CheatcodeFilter::List(filter) + } + } +} + #[cfg(test)] mod tests { use serde_json::json; diff --git a/fix/get-signatures-ordering b/fix/get-signatures-ordering new file mode 160000 index 00000000..31683fa7 --- /dev/null +++ b/fix/get-signatures-ordering @@ -0,0 +1 @@ +Subproject commit 31683fa71976efaff45a052fd2bb6191a68b2994