feat: introduce the surfnet_exportSnapshot cheatcode

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<ResetAccountConfig>,
     ) -> Result<RpcResponse<()>>;
 
+    /// 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<BTreeMap<String, AccountSnapshot>>` 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<ExportSnapshotConfig>,
+    ) -> Result<RpcResponse<BTreeMap<String, AccountSnapshot>>>;
+
     /// 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<ExportSnapshotConfig>,
+    ) -> Result<RpcResponse<BTreeMap<String, AccountSnapshot>>> {
+        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<SurfnetCheatcodesRpc>, 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<String, AccountSnapshot>,
+        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"
+        );
+    }
 }