diff --git a/crates/core/src/rpc/surfnet_cheatcodes.rs b/crates/core/src/rpc/surfnet_cheatcodes.rs index 930c83c2..c758bab7 100644 --- a/crates/core/src/rpc/surfnet_cheatcodes.rs +++ b/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, BlockAccountDownloadConfig, ClockCommand, ExportSnapshotConfig, + GetStreamedAccountsResponse, GetSurfnetInfoResponse, Idl, ResetAccountConfig, + RpcProfileResultConfig, Scenario, SimnetCommand, SimnetEvent, StreamAccountConfig, + UiKeyedProfileResult, types::{AccountUpdate, SetSomeAccount, SupplyUpdate, TokenAccountUpdate, UuidOrSignature}, }; @@ -780,6 +781,48 @@ pub trait SurfnetCheatcodes { #[rpc(meta, name = "surfnet_resetNetwork")] fn reset_network(&self, meta: Self::Metadata) -> BoxFuture>>; + /// A cheat code to prevent an account from being downloaded from the remote RPC. + /// + /// ## Parameters + /// - `pubkey_str`: The base-58 encoded public key of the account/program to block. + /// - `config`: A `BlockAccountDownloadConfig` specifying whether to also block accounts + /// owned by this pubkey. If omitted, only the account itself is blocked. + /// + /// ## Returns + /// An `RpcResponse<()>` indicating whether the download block registration was successful. + /// + /// ## Example Request + /// ```json + /// { + /// "jsonrpc": "2.0", + /// "id": 1, + /// "method": "surfnet_blockAccountDownload", + /// "params": [ "4EXSeLGxVBpAZwq7vm6evLdewpcvE2H56fpqL2pPiLFa", { "includeOwnedAccounts": true } ] + /// } + /// ``` + /// + /// ## Example Response + /// ```json + /// { + /// "jsonrpc": "2.0", + /// "result": { + /// "context": { + /// "slot": 123456789, + /// "apiVersion": "2.3.8" + /// }, + /// "value": null + /// }, + /// "id": 1 + /// } + /// ``` + #[rpc(meta, name = "surfnet_blockAccountDownload")] + fn block_account_download( + &self, + meta: Self::Metadata, + pubkey_str: String, + config: Option, + ) -> BoxFuture>>; + /// A cheat code to export a snapshot of all accounts in the Surfnet SVM. /// /// This method retrieves the current state of all accounts stored in the Surfnet Virtual Machine (SVM) @@ -1708,6 +1751,35 @@ impl SurfnetCheatcodes for SurfnetCheatcodesRpc { }) } + fn block_account_download( + &self, + meta: Self::Metadata, + pubkey_str: String, + config: Option, + ) -> BoxFuture>> { + let SurfnetRpcContext { svm_locker, .. } = + match meta.get_rpc_context(CommitmentConfig::confirmed()) { + Ok(res) => res, + Err(e) => return e.into(), + }; + let pubkey = match verify_pubkey(&pubkey_str) { + Ok(res) => res, + Err(e) => return e.into(), + }; + let config = config.unwrap_or_default(); + let include_owned_accounts = config.include_owned_accounts.unwrap_or_default(); + + Box::pin(async move { + svm_locker + .block_account_download(pubkey, include_owned_accounts) + .await?; + Ok(RpcResponse { + context: RpcResponseContext::new(svm_locker.get_latest_absolute_slot()), + value: (), + }) + }) + } + fn stream_account( &self, meta: Self::Metadata, diff --git a/crates/core/src/surfnet/locker.rs b/crates/core/src/surfnet/locker.rs index db7714fa..3f90e006 100644 --- a/crates/core/src/surfnet/locker.rs +++ b/crates/core/src/surfnet/locker.rs @@ -71,8 +71,8 @@ use crate::{ rpc::utils::{convert_transaction_metadata_from_canonical, verify_pubkey}, surfnet::{FINALIZATION_SLOT_THRESHOLD, SLOTS_PER_EPOCH}, types::{ - GeyserAccountUpdate, RemoteRpcResult, SurfnetTransactionStatus, TimeTravelConfig, - TokenAccount, TransactionLoadedAddresses, TransactionWithStatusMeta, + BlockedAccountConfig, GeyserAccountUpdate, RemoteRpcResult, SurfnetTransactionStatus, + TimeTravelConfig, TokenAccount, TransactionLoadedAddresses, TransactionWithStatusMeta, }, }; @@ -244,6 +244,34 @@ impl SurfnetSvmLocker { /// Functions for getting accounts from the underlying SurfnetSvm instance or remote client impl SurfnetSvmLocker { + /// Filters the downloaded account result to remove accounts that are owned by blocked owners. + fn filter_downloaded_account_result( + &self, + requested_pubkey: &Pubkey, + result: GetAccountResult, + ) -> GetAccountResult { + let blocked_owners = self.get_blocked_account_owners(); + match result { + GetAccountResult::FoundAccount(_, account, _) + | GetAccountResult::FoundProgramAccount((_, account), _) + | GetAccountResult::FoundTokenAccount((_, account), _) + if blocked_owners.contains(&account.owner) => + { + let blocked_pubkey = *requested_pubkey; + self.with_svm_writer(move |svm_writer| { + let _ = svm_writer.blocked_accounts.store( + blocked_pubkey.to_string(), + BlockedAccountConfig { + include_owned_accounts: false, + }, + ); + }); + GetAccountResult::None(*requested_pubkey) + } + other => other, + } + } + /// Retrieves a local account from the SVM cache, returning a contextualized result. pub fn get_account_local(&self, pubkey: &Pubkey) -> SvmAccessContext { self.with_contextualized_svm_reader(|svm_reader| { @@ -266,7 +294,7 @@ impl SurfnetSvmLocker { /// Attempts local retrieval, then fetches from remote if missing, returning a contextualized result. /// - /// Does not fetch from remote if the account has been explicitly closed by the user. + /// Does not fetch from remote if the account has been explicitly blocked from remote downloads. pub async fn get_account_local_then_remote( &self, client: &SurfnetRemoteClient, @@ -276,12 +304,12 @@ impl SurfnetSvmLocker { let result = self.get_account_local(pubkey); if result.inner.is_none() { - // Check if the account has been explicitly closed - if so, don't fetch from remote - let is_closed = self.get_closed_accounts().contains(pubkey); + let is_blocked = self.is_account_blocked(pubkey); - if !is_closed { + if !is_blocked { let remote_account = client.get_account(pubkey, commitment_config).await?; - Ok(result.with_new_value(remote_account)) + Ok(result + .with_new_value(self.filter_downloaded_account_result(pubkey, remote_account))) } else { Ok(result) } @@ -342,7 +370,7 @@ impl SurfnetSvmLocker { /// /// Returns accounts in the same order as the input `pubkeys` array. Accounts found locally /// are returned as-is; accounts not found locally are fetched from the remote RPC client. - /// Accounts that have been explicitly closed are not fetched from remote. + /// Accounts that have been explicitly blocked from remote downloads are not fetched from remote. pub async fn get_multiple_accounts_with_remote_fallback( &self, client: &SurfnetRemoteClient, @@ -356,23 +384,17 @@ impl SurfnetSvmLocker { inner: local_results, } = self.get_multiple_accounts_local(pubkeys); - // Get the closed accounts set - let closed_accounts = self.get_closed_accounts(); - - // Collect missing pubkeys that are NOT closed (local_results is already in correct order from pubkeys) - let missing_accounts: Vec = local_results - .iter() - .filter_map(|result| match result { - GetAccountResult::None(pubkey) => { - if !closed_accounts.contains(pubkey) { - Some(*pubkey) - } else { - None - } - } - _ => None, - }) - .collect(); + // Collect missing pubkeys that are not blocked (local_results is already in correct order from pubkeys). + let mut missing_accounts = Vec::new(); + for result in &local_results { + let GetAccountResult::None(pubkey) = result else { + continue; + }; + if self.is_account_blocked(pubkey) { + continue; + } + missing_accounts.push(*pubkey); + } if missing_accounts.is_empty() { // All accounts found locally, already in correct order @@ -395,8 +417,15 @@ impl SurfnetSvmLocker { // Build map of pubkey -> remote result for O(1) lookup let remote_map: HashMap = missing_accounts - .into_iter() + .iter() + .copied() .zip(remote_results.into_iter()) + .map(|(requested_pubkey, result)| { + ( + requested_pubkey, + self.filter_downloaded_account_result(&requested_pubkey, result), + ) + }) .collect(); // Replace None entries with remote results while preserving order @@ -407,8 +436,7 @@ impl SurfnetSvmLocker { .map(|(pubkey, local_result)| { match local_result { GetAccountResult::None(_) => { - // Replace with remote result if available and not closed - if closed_accounts.contains(pubkey) { + if self.is_account_blocked(pubkey) { GetAccountResult::None(*pubkey) } else { remote_map @@ -1917,9 +1945,6 @@ impl SurfnetSvmLocker { /// allowing them to be fetched fresh from mainnet on the next access. /// It handles program accounts (including their program data accounts) and can optionally /// cascade the reset to all accounts owned by a program. - /// - /// This is different from `close_account()` which marks an account as permanently closed - /// and prevents it from being fetched from mainnet. pub fn reset_account( &self, pubkey: Pubkey, @@ -1930,17 +1955,17 @@ impl SurfnetSvmLocker { "Account {} will be reset", pubkey ))); - // Unclose the account so it can be fetched from mainnet again - self.unclose_account(pubkey)?; + // Unblock the account so it can be fetched from mainnet again. + self.unblock_account_download(pubkey)?; self.with_svm_writer(move |svm_writer| { svm_writer.reset_account(&pubkey, include_owned_accounts) }) } - /// Resets SVM state and clears all closed accounts. + /// Resets SVM state and clears all blocked account download entries. /// /// This function coordinates the reset of the entire network state. - /// It also clears the closed_accounts set so all accounts can be fetched from mainnet again. + /// It also clears the blocked account set so all accounts can be fetched from mainnet again. pub async fn reset_network( &self, remote_ctx: &Option, @@ -1965,11 +1990,39 @@ impl SurfnetSvmLocker { self.with_svm_writer(move |svm_writer| { let _ = svm_writer.reset_network(epoch_info); - svm_writer.closed_accounts.clear(); + let _ = svm_writer.blocked_accounts.clear(); }); Ok(()) } + /// Blocks an account from being downloaded from the remote RPC. + /// + /// When `include_owned_accounts` is enabled, this also blocks accounts already known locally. + /// Accounts discovered later through direct remote fetches are rejected lazily if they are + /// owned by a blocked owner. + pub async fn block_account_download( + &self, + pubkey: Pubkey, + include_owned_accounts: bool, + ) -> SurfpoolResult<()> { + let simnet_events_tx = self.simnet_events_tx(); + let _ = simnet_events_tx.send(SimnetEvent::info(format!( + "Account {} will be blocked from remote downloads", + pubkey + ))); + + self.with_svm_writer(move |svm_writer| { + let _ = svm_writer.blocked_accounts.store( + pubkey.to_string(), + BlockedAccountConfig { + include_owned_accounts, + }, + ); + }); + + Ok(()) + } + /// Streams an account by its pubkey. pub fn stream_account( &self, @@ -1999,20 +2052,51 @@ impl SurfnetSvmLocker { }) } - /// Removes an account from the closed accounts set. + /// Removes an account from the blocked account download set. /// /// This allows the account to be fetched from mainnet again if requested. /// This is useful when resetting an account for a refresh/stream operation. - pub fn unclose_account(&self, pubkey: Pubkey) -> SurfpoolResult<()> { + pub fn unblock_account_download(&self, pubkey: Pubkey) -> SurfpoolResult<()> { self.with_svm_writer(move |svm_writer| { - svm_writer.closed_accounts.remove(&pubkey); + let _ = svm_writer.blocked_accounts.take(&pubkey.to_string()); }); Ok(()) } - /// Gets all currently closed accounts. - pub fn get_closed_accounts(&self) -> Vec { - self.with_svm_reader(|svm_reader| svm_reader.closed_accounts.iter().copied().collect()) + /// Returns true if the given pubkey is blocked from remote download. + pub fn is_account_blocked(&self, pubkey: &Pubkey) -> bool { + self.with_svm_reader(|svm_reader| { + svm_reader + .blocked_accounts + .contains_key(&pubkey.to_string()) + .unwrap_or(false) + }) + } + + /// Gets all currently blocked account downloads. + pub fn get_blocked_accounts(&self) -> Vec { + self.with_svm_reader(|svm_reader| { + svm_reader + .blocked_accounts + .keys() + .unwrap_or_default() + .iter() + .filter_map(|k| k.parse().ok()) + .collect() + }) + } + + /// Gets all owners whose accounts are blocked from remote download. + pub fn get_blocked_account_owners(&self) -> Vec { + self.with_svm_reader(|svm_reader| { + svm_reader + .blocked_accounts + .into_iter() + .unwrap_or_else(|_| Box::new(std::iter::empty())) + .filter(|(_, config)| config.include_owned_accounts) + .filter_map(|(k, _)| k.parse().ok()) + .collect() + }) } /// Registers a scenario for execution diff --git a/crates/core/src/surfnet/svm.rs b/crates/core/src/surfnet/svm.rs index bbcb9638..db74a639 100644 --- a/crates/core/src/surfnet/svm.rs +++ b/crates/core/src/surfnet/svm.rs @@ -102,7 +102,7 @@ use crate::{ LogsSubscriptionData, locker::is_supported_token_program, surfnet_lite_svm::SurfnetLiteSvm, }, types::{ - GeyserAccountUpdate, MintAccount, SerializableAccountAdditionalData, + BlockedAccountConfig, GeyserAccountUpdate, MintAccount, SerializableAccountAdditionalData, SurfnetTransactionStatus, SyntheticBlockhash, TokenAccount, TransactionWithStatusMeta, }, }; @@ -295,9 +295,11 @@ pub struct SurfnetSvm { pub streamed_accounts: Box>, pub recent_blockhashes: VecDeque<(SyntheticBlockhash, i64)>, pub scheduled_overrides: Box>>, - /// Tracks accounts that have been explicitly closed by the user. - /// These accounts will not be fetched from mainnet even if they don't exist in the local cache. - pub closed_accounts: HashSet, + /// Tracks accounts that should not be downloaded from the remote RPC. + /// This includes accounts explicitly closed locally and accounts blocked via cheatcodes. + /// The key is the account pubkey as a string. If `include_owned_accounts` is true, + /// accounts owned by this pubkey are also blocked from remote download. + pub blocked_accounts: Box>, /// The slot at which this surfnet instance started (may be non-zero when connected to remote). /// Used as the lower bound for block reconstruction. pub genesis_slot: Slot, @@ -419,7 +421,7 @@ impl SurfnetSvm { runbook_executions: self.runbook_executions.clone(), account_update_slots: self.account_update_slots.clone(), recent_blockhashes: self.recent_blockhashes.clone(), - closed_accounts: self.closed_accounts.clone(), + blocked_accounts: OverlayStorage::wrap(self.blocked_accounts.clone_box()), genesis_slot: self.genesis_slot, genesis_updated_at: self.genesis_updated_at, slot_checkpoint: OverlayStorage::wrap(self.slot_checkpoint.clone_box()), @@ -508,6 +510,8 @@ impl SurfnetSvm { new_kv_store(&database_url, "streamed_accounts", surfnet_id)?; let scheduled_overrides_db: Box>> = new_kv_store(&database_url, "scheduled_overrides", surfnet_id)?; + let blocked_accounts_db: Box> = + new_kv_store(&database_url, "blocked_accounts", surfnet_id)?; let registered_idls_db: Box>> = new_kv_store(&database_url, "registered_idls", surfnet_id)?; let profile_tag_map_db: Box>> = @@ -619,7 +623,7 @@ impl SurfnetSvm { streamed_accounts: streamed_accounts_db, recent_blockhashes: VecDeque::new(), scheduled_overrides: scheduled_overrides_db, - closed_accounts: HashSet::new(), + blocked_accounts: blocked_accounts_db, genesis_slot: 0, // Will be updated when connecting to remote network genesis_updated_at: Utc::now().timestamp_millis() as u64, slot_checkpoint: slot_checkpoint_db, @@ -1314,7 +1318,12 @@ impl SurfnetSvm { } if is_deleted_account { - self.closed_accounts.insert(*pubkey); + self.blocked_accounts.store( + pubkey.to_string(), + BlockedAccountConfig { + include_owned_accounts: false, + }, + )?; if let Some(old_account) = self.get_account(pubkey)? { self.remove_from_indexes(pubkey, &old_account)?; } @@ -4554,14 +4563,22 @@ mod tests { svm.set_account(&account_pubkey, account.clone()).unwrap(); assert!(svm.get_account(&account_pubkey).unwrap().is_some()); - assert!(!svm.closed_accounts.contains(&account_pubkey)); + assert!( + !svm.blocked_accounts + .contains_key(&account_pubkey.to_string()) + .unwrap() + ); assert_eq!(svm.get_account_owned_by(&owner).unwrap().len(), 1); let empty_account = Account::default(); svm.update_account_registries(&account_pubkey, &empty_account) .unwrap(); - assert!(svm.closed_accounts.contains(&account_pubkey)); + assert!( + svm.blocked_accounts + .contains_key(&account_pubkey.to_string()) + .unwrap() + ); assert_eq!(svm.get_account_owned_by(&owner).unwrap().len(), 0); @@ -4609,13 +4626,21 @@ mod tests { 1 ); assert_eq!(svm.get_token_accounts_by_delegate(&delegate).len(), 1); - assert!(!svm.closed_accounts.contains(&token_account_pubkey)); + assert!( + !svm.blocked_accounts + .contains_key(&token_account_pubkey.to_string()) + .unwrap() + ); let empty_account = Account::default(); svm.update_account_registries(&token_account_pubkey, &empty_account) .unwrap(); - assert!(svm.closed_accounts.contains(&token_account_pubkey)); + assert!( + svm.blocked_accounts + .contains_key(&token_account_pubkey.to_string()) + .unwrap() + ); assert_eq!( svm.get_token_accounts_by_owner(&token_owner).unwrap().len(), diff --git a/crates/core/src/tests/integration.rs b/crates/core/src/tests/integration.rs index 3bfa0149..bdd3b9b0 100644 --- a/crates/core/src/tests/integration.rs +++ b/crates/core/src/tests/integration.rs @@ -4742,6 +4742,106 @@ async fn test_closed_accounts(test_type: TestType) { } } +#[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"))] +#[cfg_attr(feature = "ignore_tests_ci", ignore = "flaky CI tests")] +#[tokio::test(flavor = "multi_thread")] +async fn test_block_account_download_including_owned_accounts(test_type: TestType) { + let owner = Pubkey::new_unique(); + let owned = Pubkey::new_unique(); + let another_test_type = match &test_type { + TestType::OnDiskSqlite(_) => TestType::sqlite(), + TestType::InMemorySqlite => TestType::in_memory(), + TestType::NoDb => TestType::no_db(), + #[cfg(feature = "postgres")] + TestType::Postgres { url, .. } => TestType::Postgres { + url: url.clone(), + surfnet_id: crate::storage::tests::random_surfnet_id(), + }, + }; + + let (datasource_surfnet_url, datasource_svm_locker) = + start_surfnet(vec![], None, test_type).expect("Failed to start datasource surfnet"); + + datasource_svm_locker + .with_svm_writer(|svm_writer| { + svm_writer + .set_account( + &owner, + Account { + lamports: LAMPORTS_PER_SOL, + data: vec![1, 2, 3], + owner: system_program::id(), + executable: false, + rent_epoch: 0, + }, + ) + .unwrap(); + svm_writer + .set_account( + &owned, + Account { + lamports: LAMPORTS_PER_SOL / 2, + data: vec![4, 5, 6], + owner, + executable: false, + rent_epoch: 0, + }, + ) + .unwrap(); + Ok::<(), SurfpoolError>(()) + }) + .expect("Failed to seed datasource accounts"); + + let (surfnet_url, surfnet_svm_locker) = + start_surfnet(vec![], Some(datasource_surfnet_url), another_test_type) + .expect("Failed to start surfnet"); + let rpc_client = RpcClient::new(surfnet_url); + + let _: serde_json::Value = rpc_client + .send( + solana_client::rpc_request::RpcRequest::Custom { + method: "surfnet_blockAccountDownload", + }, + serde_json::json!([owner.to_string(), { "includeOwnedAccounts": true }]), + ) + .await + .expect("Failed to block account download"); + + assert!( + surfnet_svm_locker.is_account_blocked(&owner), + "Owner should be recorded as blocked" + ); + assert!( + surfnet_svm_locker + .get_blocked_account_owners() + .contains(&owner), + "Owner should be recorded as a blocked account owner" + ); + + let owner_result = rpc_client.get_account(&owner).await; + assert!( + owner_result.is_err(), + "Blocked owner account should not be fetched from remote" + ); + assert!( + surfnet_svm_locker.get_account_local(&owner).inner.is_none(), + "Blocked owner account should remain absent locally" + ); + + let owned_result = rpc_client.get_account(&owned).await; + assert!( + owned_result.is_err(), + "Owned account should not be fetched from remote once blocked" + ); + assert!( + surfnet_svm_locker.get_account_local(&owned).inner.is_none(), + "Blocked owned account should remain absent locally" + ); +} + #[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")] diff --git a/crates/core/src/types.rs b/crates/core/src/types.rs index d8a7c1d5..435a85cb 100644 --- a/crates/core/src/types.rs +++ b/crates/core/src/types.rs @@ -1068,6 +1068,11 @@ impl<'de> Deserialize<'de> for TokenAccount { } } +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct BlockedAccountConfig { + pub include_owned_accounts: bool, +} + #[derive(Debug, Clone)] pub enum MintAccount { SplToken2022(spl_token_2022_interface::state::Mint), diff --git a/crates/types/src/types.rs b/crates/types/src/types.rs index 7cdc3f8a..f98fda6d 100644 --- a/crates/types/src/types.rs +++ b/crates/types/src/types.rs @@ -778,6 +778,7 @@ impl CloudSurfnetRpcGating { "surfnet_resetAccount".into(), "surfnet_resetNetwork".into(), "surfnet_exportSnapshot".into(), + "surfnet_blockAccountDownload".into(), "surfnet_streamAccount".into(), "surfnet_getStreamedAccounts".into(), ], @@ -1170,6 +1171,7 @@ pub struct ExportSnapshotFilter { } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] pub struct ResetAccountConfig { pub include_owned_accounts: Option, } @@ -1183,6 +1185,7 @@ impl Default for ResetAccountConfig { } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] pub struct StreamAccountConfig { pub include_owned_accounts: Option, } @@ -1195,6 +1198,20 @@ impl Default for StreamAccountConfig { } } +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct BlockAccountDownloadConfig { + pub include_owned_accounts: Option, +} + +impl Default for BlockAccountDownloadConfig { + fn default() -> Self { + Self { + include_owned_accounts: Some(false), + } + } +} + #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct StreamedAccountInfo {