diff --git a/crates/core/src/rpc/surfnet_cheatcodes.rs b/crates/core/src/rpc/surfnet_cheatcodes.rs index d02942d0..65334c13 100644 --- a/crates/core/src/rpc/surfnet_cheatcodes.rs +++ b/crates/core/src/rpc/surfnet_cheatcodes.rs @@ -1,3 +1,5 @@ +use std::collections::BTreeMap; + use base64::{Engine as _, engine::general_purpose::STANDARD}; use jsonrpc_core::{BoxFuture, Error, Result, futures::future}; use jsonrpc_derive::rpc; @@ -12,8 +14,9 @@ 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::{ - ClockCommand, GetSurfnetInfoResponse, Idl, ResetAccountConfig, RpcProfileResultConfig, - SimnetCommand, SimnetEvent, StreamAccountConfig, UiKeyedProfileResult, + AccountSnapshot, ClockCommand, ExportSnapshotConfig, GetSurfnetInfoResponse, Idl, + ResetAccountConfig, RpcProfileResultConfig, SimnetCommand, SimnetEvent, StreamAccountConfig, + UiKeyedProfileResult, types::{AccountUpdate, SetSomeAccount, SupplyUpdate, TokenAccountUpdate, UuidOrSignature}, }; @@ -746,6 +749,66 @@ pub trait SurfnetCheatcodes { config: Option, ) -> Result>; + /// 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) + /// and returns them as a mapping of account public keys to their respective account snapshots. + /// + /// ## Parameters + /// - `config`: An optional `ExportSnapshotConfig` to customize the export behavior. The config fields are: + /// - `includeParsedAccounts`: If true, includes parsed account data in the snapshot. + /// - `filter`: An optional filter config to limit which accounts are included in the snapshot. Fields include: + /// - `includeProgramAccounts`: A list of program IDs to include accounts for. + /// - `includeAccounts`: A list of specific account public keys to include. + /// - `excludeAccounts`: A list of specific account public keys to exclude. + /// + /// + /// ## Returns + /// An `RpcResponse>` containing the exported account snapshots. + /// + /// The keys of the map are the base-58 encoded public keys of the accounts, + /// and the values are the corresponding `AccountSnapshot` objects. + /// + /// ## Example Request + /// ```json + /// { + /// "jsonrpc": "2.0", + /// "id": 1, + /// "method": "surfnet_exportSnapshot" + /// } + /// ``` + /// + /// ## Example Response + /// ```json + /// { + /// "jsonrpc": "2.0", + /// "result": { + /// "4EXSeLGxVBpAZwq7vm6evLdewpcvE2H56fpqL2pPiLFa": { + /// "lamports": 1000000, + /// "owner": "11111111111111111111111111111111", + /// "executable": false, + /// "rent_epoch": 0, + /// "data": "base64_encoded_data_string" + /// }, + /// "AnotherAccountPubkeyBase58": { + /// "lamports": 500000, + /// "owner": "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA", + /// "executable": false, + /// "rent_epoch": 0, + /// "data": "base64_encoded_data_string" + /// } + /// }, + /// "id": 1 + /// } + /// ``` + /// + #[rpc(meta, name = "surfnet_exportSnapshot")] + fn export_snapshot( + &self, + meta: Self::Metadata, + config: Option, + ) -> Result>>; + /// A cheat code to simulate account streaming. /// When a transaction is processed, the accounts that are accessed are downloaded from the datasource and cached in the SVM. /// With this method, you can simulate the streaming of accounts by providing a pubkey. @@ -1387,6 +1450,20 @@ impl SurfnetCheatcodes for SurfnetCheatcodesRpc { value: GetSurfnetInfoResponse::new(runbook_executions), }) } + + fn export_snapshot( + &self, + meta: Self::Metadata, + config: Option, + ) -> Result>> { + let config = config.unwrap_or_default(); + let svm_locker = meta.get_svm_locker()?; + let snapshot = svm_locker.export_snapshot(config); + Ok(RpcResponse { + context: RpcResponseContext::new(svm_locker.get_latest_absolute_slot()), + value: snapshot, + }) + } } #[cfg(test)] @@ -1396,6 +1473,7 @@ mod tests { }; use solana_keypair::Keypair; use solana_program_pack::Pack; + use solana_pubkey::Pubkey; use solana_signer::Signer; use solana_system_interface::instruction::create_account; use solana_transaction::Transaction; @@ -1405,7 +1483,9 @@ mod tests { }; use spl_token_2022_interface::instruction::{initialize_mint2, mint_to, transfer_checked}; use spl_token_interface::state::Mint; - use surfpool_types::{RpcProfileDepth, UiAccountChange, UiAccountProfileState}; + use surfpool_types::{ + ExportSnapshotFilter, RpcProfileDepth, UiAccountChange, UiAccountProfileState, + }; use super::*; use crate::{rpc::surfnet_cheatcodes::SurfnetCheatcodesRpc, tests::helpers::TestSetup}; @@ -2210,4 +2290,270 @@ mod tests { ); } } + + fn set_account(client: &TestSetup, pubkey: &Pubkey, account: &Account) { + client + .context + .svm_locker + .with_svm_writer(|svm| svm.inner.set_account(*pubkey, account.clone())) + .expect("Failed to set account"); + } + + fn verify_snapshot_account( + snapshot: &BTreeMap, + expected_account_pubkey: &Pubkey, + expected_account: &Account, + ) { + let account = snapshot + .get(&expected_account_pubkey.to_string()) + .unwrap_or_else(|| { + panic!( + "Account fixture not found for pubkey {}", + expected_account_pubkey + ) + }); + assert_eq!(expected_account.lamports, account.lamports); + assert_eq!( + base64::engine::general_purpose::STANDARD.encode(&expected_account.data), + account.data + ); + assert_eq!(expected_account.owner.to_string(), account.owner); + assert_eq!(expected_account.executable, account.executable); + assert_eq!(expected_account.rent_epoch, account.rent_epoch); + } + + #[test] + fn test_export_snapshot() { + let client = TestSetup::new(SurfnetCheatcodesRpc); + + let pubkey1 = Pubkey::new_unique(); + let account1 = Account { + lamports: 1_000_000, + data: vec![1, 2, 3, 4], + owner: system_program::id(), + executable: false, + rent_epoch: 0, + }; + + set_account(&client, &pubkey1, &account1); + + let pubkey2 = Pubkey::new_unique(); + let account2 = Account { + lamports: 2_000_000, + data: vec![5, 6, 7, 8, 9], + owner: system_program::id(), + executable: false, + rent_epoch: 0, + }; + + set_account(&client, &pubkey2, &account2); + + let snapshot = client + .rpc + .export_snapshot(Some(client.context.clone()), None) + .expect("Failed to export snapshot") + .value; + + verify_snapshot_account(&snapshot, &pubkey1, &account1); + verify_snapshot_account(&snapshot, &pubkey2, &account2); + } + + #[test] + fn test_export_snapshot_json_parsed() { + let client = TestSetup::new(SurfnetCheatcodesRpc); + + let pubkey1 = Pubkey::new_unique(); + println!("Pubkey1: {}", pubkey1); + let account1 = Account { + lamports: 1_000_000, + data: vec![1, 2, 3, 4], + owner: system_program::id(), + executable: false, + rent_epoch: 0, + }; + + set_account(&client, &pubkey1, &account1); + + let mint_pubkey = Pubkey::new_unique(); + println!("Mint Pubkey: {}", mint_pubkey); + let mint_authority = Pubkey::new_unique(); + + let mut mint_data = [0u8; Mint::LEN]; + let mint = Mint { + mint_authority: COption::Some(mint_authority), + supply: 1000, + decimals: 6, + is_initialized: true, + freeze_authority: COption::None, + }; + mint.pack_into_slice(&mut mint_data); + + let mint_account = Account { + lamports: 1_000_000, + data: mint_data.to_vec(), + owner: spl_token_interface::id(), + executable: false, + rent_epoch: 0, + }; + + set_account(&client, &mint_pubkey, &mint_account); + + let snapshot = client + .rpc + .export_snapshot( + Some(client.context.clone()), + Some(ExportSnapshotConfig { + include_parsed_accounts: Some(true), + filter: None, + }), + ) + .expect("Failed to export snapshot") + .value; + + verify_snapshot_account(&snapshot, &pubkey1, &account1); + let actual_account1 = snapshot + .get(&pubkey1.to_string()) + .expect("Account fixture not found"); + assert!( + actual_account1.parsed_data.is_none(), + "Account1 should not have parsed data" + ); + + verify_snapshot_account(&snapshot, &mint_pubkey, &mint_account); + let mint_snapshot = snapshot + .get(&mint_pubkey.to_string()) + .expect("Mint account snapshot not found"); + let parsed = mint_snapshot + .parsed_data + .as_ref() + .expect("Parsed data should be present"); + + assert_eq!(parsed.program, "spl-token"); + assert_eq!(parsed.space, Mint::LEN as u64); + + let parsed_info = parsed + .parsed + .as_object() + .expect("Parsed data should be an object"); + let info = parsed_info + .get("info") + .expect("Parsed data should have info field") + .as_object() + .expect("Info field should be an object"); + assert_eq!( + info.get("mintAuthority") + .and_then(|v| v.as_str()) + .expect("mintAuthority should be a string"), + mint_authority.to_string() + ); + } + + #[test] + fn test_export_snapshot_filtering() { + let system_account_pubkey = Pubkey::new_unique(); + println!("System Account Pubkey: {}", system_account_pubkey); + let excluded_system_account_pubkey = Pubkey::new_unique(); + println!( + "Excluded System Account Pubkey: {}", + excluded_system_account_pubkey + ); + let program_account_pubkey = Pubkey::new_unique(); + println!("Program Account Pubkey: {}", program_account_pubkey); + let included_program_account_pubkey = Pubkey::new_unique(); + println!( + "Included Program Account Pubkey: {}", + included_program_account_pubkey + ); + + let client = TestSetup::new(SurfnetCheatcodesRpc); + + let system_account = Account { + lamports: 1_000_000, + data: vec![1, 2, 3, 4], + owner: system_program::id(), + executable: false, + rent_epoch: 0, + }; + set_account(&client, &system_account_pubkey, &system_account); + set_account(&client, &excluded_system_account_pubkey, &system_account); + + let program_account = Account { + lamports: 2_000_000, + data: vec![5, 6, 7, 8, 9], + owner: solana_sdk_ids::bpf_loader_upgradeable::id(), + executable: false, + rent_epoch: 0, + }; + set_account(&client, &program_account_pubkey, &program_account); + set_account(&client, &included_program_account_pubkey, &program_account); + + let snapshot = client + .rpc + .export_snapshot(Some(client.context.clone()), None) + .expect("Failed to export snapshot") + .value; + assert!( + !snapshot.contains_key(&program_account_pubkey.to_string()), + "Program account should be excluded by default" + ); + assert!( + !snapshot.contains_key(&included_program_account_pubkey.to_string()), + "Program account should be excluded by default" + ); + let snapshot = client + .rpc + .export_snapshot( + Some(client.context.clone()), + Some(ExportSnapshotConfig { + filter: Some(ExportSnapshotFilter { + include_accounts: Some(vec![included_program_account_pubkey.to_string()]), + ..Default::default() + }), + ..Default::default() + }), + ) + .expect("Failed to export snapshot") + .value; + assert!( + !snapshot.contains_key(&program_account_pubkey.to_string()), + "Program account should be excluded by default" + ); + assert!( + snapshot.contains_key(&included_program_account_pubkey.to_string()), + "Program account should be included when explicitly listed" + ); + + let snapshot = client + .rpc + .export_snapshot( + Some(client.context.clone()), + Some(ExportSnapshotConfig { + filter: Some(ExportSnapshotFilter { + include_program_accounts: Some(true), + exclude_accounts: Some(vec![excluded_system_account_pubkey.to_string()]), + ..Default::default() + }), + ..Default::default() + }), + ) + .expect("Failed to export snapshot") + .value; + + assert!( + snapshot.contains_key(&program_account_pubkey.to_string()), + "Program account should be included when filter is set" + ); + assert!( + snapshot.contains_key(&included_program_account_pubkey.to_string()), + "Included program account should be present" + ); + assert!( + snapshot.contains_key(&system_account_pubkey.to_string()), + "System account should be present" + ); + assert!( + !snapshot.contains_key(&excluded_system_account_pubkey.to_string()), + "Excluded system account should not be present" + ); + } } diff --git a/crates/core/src/surfnet/locker.rs b/crates/core/src/surfnet/locker.rs index c536fb36..bcc1fce4 100644 --- a/crates/core/src/surfnet/locker.rs +++ b/crates/core/src/surfnet/locker.rs @@ -53,10 +53,10 @@ use solana_transaction_status::{ UiTransactionEncoding, }; use surfpool_types::{ - ComputeUnitsEstimationResult, ExecutionCapture, Idl, KeyedProfileResult, ProfileResult, - RpcProfileResultConfig, RunbookExecutionStatusReport, SimnetCommand, SimnetEvent, - TransactionConfirmationStatus, TransactionStatusEvent, UiKeyedProfileResult, UuidOrSignature, - VersionedIdl, + AccountSnapshot, ComputeUnitsEstimationResult, ExecutionCapture, ExportSnapshotConfig, Idl, + KeyedProfileResult, ProfileResult, RpcProfileResultConfig, RunbookExecutionStatusReport, + SimnetCommand, SimnetEvent, TransactionConfirmationStatus, TransactionStatusEvent, + UiKeyedProfileResult, UuidOrSignature, VersionedIdl, }; use tokio::sync::RwLock; use txtx_addon_kit::indexmap::IndexSet; @@ -3009,6 +3009,13 @@ impl SurfnetSvmLocker { } }); } + + pub fn export_snapshot( + &self, + config: ExportSnapshotConfig, + ) -> BTreeMap { + self.with_svm_reader(|svm_reader| svm_reader.export_snapshot(config)) + } } // Helper function to apply filters diff --git a/crates/core/src/surfnet/svm.rs b/crates/core/src/surfnet/svm.rs index b941bb23..835a47bd 100644 --- a/crates/core/src/surfnet/svm.rs +++ b/crates/core/src/surfnet/svm.rs @@ -1,10 +1,11 @@ use std::{ cmp::max, - collections::{BinaryHeap, HashMap, HashSet, VecDeque}, + collections::{BTreeMap, BinaryHeap, HashMap, HashSet, VecDeque}, str::FromStr, }; use agave_feature_set::{FeatureSet, enable_extend_program_checked}; +use base64::{Engine, prelude::BASE64_STANDARD}; use chrono::Utc; use convert_case::Casing; use crossbeam_channel::{Receiver, Sender, unbounded}; @@ -37,7 +38,7 @@ use solana_message::{Message, VersionedMessage, v0::LoadedAddresses}; use solana_program_option::COption; use solana_pubkey::Pubkey; use solana_rpc_client_api::response::SlotInfo; -use solana_sdk_ids::system_program; +use solana_sdk_ids::{bpf_loader, system_program}; use solana_signature::Signature; use solana_signer::Signer; use solana_system_interface::instruction as system_instruction; @@ -49,10 +50,11 @@ use spl_token_2022_interface::extension::{ scaled_ui_amount::ScaledUiAmountConfig, }; use surfpool_types::{ - AccountChange, AccountProfileState, DEFAULT_PROFILING_MAP_CAPACITY, DEFAULT_SLOT_TIME_MS, - FifoMap, Idl, ProfileResult, RpcProfileDepth, RpcProfileResultConfig, - RunbookExecutionStatusReport, SimnetEvent, TransactionConfirmationStatus, - TransactionStatusEvent, UiAccountChange, UiAccountProfileState, UiProfileResult, VersionedIdl, + AccountChange, AccountProfileState, AccountSnapshot, DEFAULT_PROFILING_MAP_CAPACITY, + DEFAULT_SLOT_TIME_MS, ExportSnapshotConfig, FifoMap, Idl, ProfileResult, RpcProfileDepth, + RpcProfileResultConfig, RunbookExecutionStatusReport, SimnetEvent, + TransactionConfirmationStatus, TransactionStatusEvent, UiAccountChange, UiAccountProfileState, + UiProfileResult, VersionedIdl, types::{ ComputeUnitsEstimationResult, KeyedProfileResult, UiKeyedProfileResult, UuidOrSignature, }, @@ -146,6 +148,7 @@ pub struct SurfnetSvm { pub instruction_profiling_enabled: bool, pub max_profiles: usize, pub runbook_executions: Vec, + pub account_update_slots: HashMap, pub streamed_accounts: HashMap, } @@ -235,6 +238,7 @@ impl SurfnetSvm { instruction_profiling_enabled: true, max_profiles: DEFAULT_PROFILING_MAP_CAPACITY, runbook_executions: Vec::new(), + account_update_slots: HashMap::new(), streamed_accounts: HashMap::new(), }; @@ -514,6 +518,9 @@ impl SurfnetSvm { .set_account(*pubkey, account.clone()) .map_err(|e| SurfpoolError::set_account(*pubkey, e))?; + self.account_update_slots + .insert(*pubkey, self.get_latest_absolute_slot()); + // Update the account registries and indexes self.update_account_registries(pubkey, &account)?; @@ -848,6 +855,7 @@ impl SurfnetSvm { fn confirm_transactions(&mut self) -> Result<(Vec, HashSet), SurfpoolError> { let mut confirmed_transactions = vec![]; let slot = self.latest_epoch_info.slot_index; + let current_slot = self.latest_epoch_info.absolute_slot; let mut all_mutated_account_keys = HashSet::new(); @@ -881,6 +889,10 @@ impl SurfnetSvm { let (tx_with_status_meta, mutated_account_keys) = tx_data.as_ref(); all_mutated_account_keys.extend(mutated_account_keys); + for pubkey in mutated_account_keys { + self.account_update_slots.insert(*pubkey, current_slot); + } + self.notify_logs_subscribers( &signature, None, @@ -1725,6 +1737,84 @@ impl SurfnetSvm { execution.mark_completed(error); } } + + /// Export all accounts to a JSON file suitable for test fixtures + /// + /// # Arguments + /// * `encoding` - The encoding to use for account data (Base64, JsonParsed, etc.) + /// + /// # Returns + /// A BTreeMap of pubkey -> AccountFixture that can be serialized to JSON. + pub fn export_snapshot( + &self, + config: ExportSnapshotConfig, + ) -> BTreeMap { + let mut fixtures = BTreeMap::new(); + let encoding = if config.include_parsed_accounts.unwrap_or_default() { + UiAccountEncoding::JsonParsed + } else { + UiAccountEncoding::Base64 + }; + let filter = config.filter.unwrap_or_default(); + let include_program_accounts = filter.include_program_accounts.unwrap_or(false); + let include_accounts = filter.include_accounts.unwrap_or_default(); + let exclude_accounts = filter.exclude_accounts.unwrap_or_default(); + + fn is_program_account(pubkey: &Pubkey) -> bool { + pubkey == &bpf_loader::id() + || pubkey == &solana_sdk_ids::bpf_loader_deprecated::id() + || pubkey == &solana_sdk_ids::bpf_loader_upgradeable::id() + } + for (pubkey, account_shared_data) in self.iter_accounts() { + let is_include_account = include_accounts.iter().any(|k| k.eq(&pubkey.to_string())); + let is_exclude_account = exclude_accounts.iter().any(|k| k.eq(&pubkey.to_string())); + let is_program_account = is_program_account(account_shared_data.owner()); + if is_exclude_account + || ((is_program_account && !include_program_accounts) && !is_include_account) + { + continue; + } + let account = Account::from(account_shared_data.clone()); + + // For token accounts, we need to provide the mint additional data + let additional_data = if account.owner == spl_token_interface::id() + || account.owner == spl_token_2022_interface::id() + { + if let Ok(token_account) = TokenAccount::unpack(&account.data) { + self.account_associated_data + .get(&token_account.mint()) + .cloned() + } else { + self.account_associated_data.get(pubkey).cloned() + } + } else { + self.account_associated_data.get(pubkey).cloned() + }; + + let ui_account = + self.encode_ui_account(pubkey, &account, encoding, additional_data, None); + + let (base64, parsed_data) = match ui_account.data { + UiAccountData::Json(parsed_account) => { + (BASE64_STANDARD.encode(account.data()), Some(parsed_account)) + } + UiAccountData::Binary(base64, _) => (base64, None), + UiAccountData::LegacyBinary(_) => unreachable!(), + }; + + let account_snapshot = AccountSnapshot::new( + account.lamports, + account.owner.to_string(), + account.executable, + account.rent_epoch, + base64, + parsed_data, + ); + + fixtures.insert(pubkey.to_string(), account_snapshot); + } + fixtures + } } #[cfg(test)] diff --git a/crates/types/src/types.rs b/crates/types/src/types.rs index a155ab90..daea7e50 100644 --- a/crates/types/src/types.rs +++ b/crates/types/src/types.rs @@ -12,7 +12,7 @@ use crossbeam_channel::{Receiver, Sender}; use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Visitor}; use serde_with::{BytesOrString, serde_as}; use solana_account::Account; -use solana_account_decoder_client_types::{UiAccount, UiAccountEncoding}; +use solana_account_decoder_client_types::{ParsedAccount, UiAccount, UiAccountEncoding}; use solana_clock::{Clock, Epoch, Slot}; use solana_epoch_info::EpochInfo; use solana_message::inner_instruction::InnerInstructionsList; @@ -955,6 +955,54 @@ impl FifoMap { } } +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountSnapshot { + pub lamports: u64, + pub owner: String, + pub executable: bool, + pub rent_epoch: u64, + /// Base64 encoded data + pub data: String, + /// Parsed account data if available + pub parsed_data: Option, +} + +impl AccountSnapshot { + pub fn new( + lamports: u64, + owner: String, + executable: bool, + rent_epoch: u64, + data: String, + parsed_data: Option, + ) -> Self { + Self { + lamports, + owner, + executable, + rent_epoch, + data, + parsed_data, + } + } +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ExportSnapshotConfig { + pub include_parsed_accounts: Option, + pub filter: Option, +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ExportSnapshotFilter { + pub include_program_accounts: Option, + pub include_accounts: Option>, + pub exclude_accounts: Option>, +} + #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct ResetAccountConfig { pub include_owned_accounts: Option,