Skip to main content

surfpool_core/surfnet/
locker.rs

1use std::{
2    collections::{BTreeMap, HashMap, HashSet},
3    sync::Arc,
4    time::SystemTime,
5};
6
7use bincode::serialized_size;
8use crossbeam_channel::{Receiver, Sender};
9use itertools::Itertools;
10use litesvm::types::{
11    FailedTransactionMetadata, SimulatedTransactionInfo, TransactionMetadata, TransactionResult,
12};
13use solana_account::{Account, ReadableAccount};
14use solana_account_decoder::{
15    UiAccount, UiAccountEncoding, UiDataSliceConfig,
16    parse_account_data::AccountAdditionalDataV3,
17    parse_bpf_loader::{BpfUpgradeableLoaderAccountType, UiProgram, parse_bpf_upgradeable_loader},
18    parse_token::UiTokenAmount,
19};
20use solana_address_lookup_table_interface::state::AddressLookupTable;
21use solana_client::{
22    rpc_client::SerializableTransaction,
23    rpc_config::{
24        RpcAccountInfoConfig, RpcBlockConfig, RpcLargestAccountsConfig, RpcLargestAccountsFilter,
25        RpcSignaturesForAddressConfig, RpcTransactionConfig, RpcTransactionLogsFilter,
26    },
27    rpc_filter::RpcFilterType,
28    rpc_request::TokenAccountsFilter,
29    rpc_response::{
30        RpcAccountBalance, RpcConfirmedTransactionStatusWithSignature, RpcKeyedAccount,
31        RpcLogsResponse, RpcTokenAccountBalance,
32    },
33};
34use solana_clock::{Clock, Slot, UnixTimestamp};
35use solana_commitment_config::{CommitmentConfig, CommitmentLevel};
36use solana_epoch_info::EpochInfo;
37use solana_hash::Hash;
38use solana_loader_v3_interface::{get_program_data_address, state::UpgradeableLoaderState};
39use solana_message::{
40    Message, SimpleAddressLoader, VersionedMessage,
41    compiled_instruction::CompiledInstruction,
42    v0::{LoadedAddresses, MessageAddressTableLookup},
43};
44use solana_pubkey::Pubkey;
45use solana_rpc_client_api::response::SlotInfo;
46use solana_signature::Signature;
47use solana_transaction::{sanitized::SanitizedTransaction, versioned::VersionedTransaction};
48use solana_transaction_error::TransactionError;
49use solana_transaction_status::{
50    EncodedConfirmedTransactionWithStatusMeta,
51    TransactionConfirmationStatus as SolanaTransactionConfirmationStatus, UiConfirmedBlock,
52    UiTransactionEncoding,
53};
54use surfpool_types::{
55    AccountSnapshot, ComputeUnitsEstimationResult, ExecutionCapture, ExportSnapshotConfig, Idl,
56    KeyedProfileResult, ProfileResult, RpcProfileResultConfig, RunbookExecutionStatusReport,
57    SimnetCommand, SimnetEvent, TransactionConfirmationStatus, TransactionStatusEvent,
58    UiKeyedProfileResult, UuidOrSignature, VersionedIdl,
59};
60use tokio::sync::RwLock;
61use txtx_addon_kit::indexmap::IndexSet;
62use uuid::Uuid;
63
64use super::{
65    AccountFactory, GetAccountResult, GetTransactionResult, GeyserEvent, SignatureSubscriptionType,
66    SurfnetSvm, remote::SurfnetRemoteClient,
67};
68use crate::{
69    error::{SurfpoolError, SurfpoolResult},
70    helpers::time_travel::calculate_time_travel_clock,
71    rpc::utils::{convert_transaction_metadata_from_canonical, verify_pubkey},
72    surfnet::FINALIZATION_SLOT_THRESHOLD,
73    types::{
74        GeyserAccountUpdate, OfflineAccountConfig, RemoteRpcResult, SurfnetTransactionStatus,
75        TimeTravelConfig, TokenAccount, TransactionLoadedAddresses, TransactionWithStatusMeta,
76    },
77};
78
79enum ProcessTransactionResult {
80    Success(TransactionMetadata),
81    SimulationFailure(FailedTransactionMetadata),
82    ExecutionFailure(FailedTransactionMetadata),
83}
84
85pub struct SvmAccessContext<T> {
86    pub slot: Slot,
87    pub latest_epoch_info: EpochInfo,
88    pub latest_blockhash: Hash,
89    pub inner: T,
90}
91
92impl<T> SvmAccessContext<T> {
93    pub fn new(slot: Slot, latest_epoch_info: EpochInfo, latest_blockhash: Hash, inner: T) -> Self {
94        Self {
95            slot,
96            latest_blockhash,
97            latest_epoch_info,
98            inner,
99        }
100    }
101
102    pub fn inner(&self) -> &T {
103        &self.inner
104    }
105
106    pub fn with_new_value<N>(self, inner: N) -> SvmAccessContext<N> {
107        SvmAccessContext {
108            slot: self.slot,
109            latest_blockhash: self.latest_blockhash,
110            latest_epoch_info: self.latest_epoch_info,
111            inner,
112        }
113    }
114}
115
116pub type SurfpoolContextualizedResult<T> = SurfpoolResult<SvmAccessContext<T>>;
117
118/// Helper function to apply an override to a JSON value using dot notation path
119///
120/// # Arguments
121/// * `json` - The JSON value to modify
122/// * `path` - Dot-separated path to the field (e.g., "price_message.price")
123/// * `value` - The new value to set
124///
125/// # Returns
126/// Result indicating success or error
127pub struct SurfnetSvmLocker(pub Arc<RwLock<SurfnetSvm>>);
128
129impl Clone for SurfnetSvmLocker {
130    fn clone(&self) -> Self {
131        Self(self.0.clone())
132    }
133}
134
135/// Functions for reading and writing to the underlying SurfnetSvm instance
136impl SurfnetSvmLocker {
137    /// Explicitly shutdown the SVM, performing cleanup like WAL checkpoint for SQLite.
138    /// This should be called before the application exits to ensure data is persisted.
139    pub fn shutdown(&self) {
140        let read_lock = self.0.clone();
141        tokio::task::block_in_place(move || {
142            let read_guard = read_lock.blocking_read();
143            read_guard.shutdown();
144        });
145    }
146
147    /// Executes a read-only operation on the underlying `SurfnetSvm` by acquiring a blocking read lock.
148    /// Accepts a closure that receives a shared reference to `SurfnetSvm` and returns a value.
149    ///
150    /// # Returns
151    /// The result produced by the closure.
152    pub fn with_svm_reader<T, F>(&self, reader: F) -> T
153    where
154        F: FnOnce(&SurfnetSvm) -> T + Send + Sync,
155    {
156        let read_lock = self.0.clone();
157        tokio::task::block_in_place(move || {
158            let read_guard = read_lock.blocking_read();
159            reader(&read_guard)
160        })
161    }
162
163    /// Executes a read-only operation and wraps the result in `SvmAccessContext`, capturing
164    /// slot, epoch info, and blockhash along with the closure's result.
165    fn with_contextualized_svm_reader<T, F>(&self, reader: F) -> SvmAccessContext<T>
166    where
167        F: Fn(&SurfnetSvm) -> T + Send + Sync,
168        T: Send + 'static,
169    {
170        let read_lock = self.0.clone();
171        tokio::task::block_in_place(move || {
172            let read_guard = read_lock.blocking_read();
173            let res = reader(&read_guard);
174
175            SvmAccessContext::new(
176                read_guard.get_latest_absolute_slot(),
177                read_guard.latest_epoch_info(),
178                read_guard.latest_blockhash(),
179                res,
180            )
181        })
182    }
183
184    /// Executes a write operation on the underlying `SurfnetSvm` by acquiring a blocking write lock.
185    /// Accepts a closure that receives a mutable reference to `SurfnetSvm` and returns a value.
186    ///
187    /// # Returns
188    /// The result produced by the closure.
189    pub fn with_svm_writer<T, F>(&self, writer: F) -> T
190    where
191        F: FnOnce(&mut SurfnetSvm) -> T + Send + Sync,
192        T: Send + 'static,
193    {
194        let write_lock = self.0.clone();
195        tokio::task::block_in_place(move || {
196            let mut write_guard = write_lock.blocking_write();
197            writer(&mut write_guard)
198        })
199    }
200}
201
202/// Functions for creating and initializing the underlying SurfnetSvm instance
203impl SurfnetSvmLocker {
204    /// Constructs a new `SurfnetSvmLocker` wrapping the given `SurfnetSvm` instance.
205    pub fn new(svm: SurfnetSvm) -> Self {
206        Self(Arc::new(RwLock::new(svm)))
207    }
208
209    /// Initializes the locked `SurfnetSvm` with remote-derived startup state when available.
210    pub async fn initialize(&self, remote_ctx: &Option<SurfnetRemoteClient>) -> SurfpoolResult<()> {
211        let Some(remote_client) = remote_ctx else {
212            return Ok(());
213        };
214
215        let (mut epoch_info, epoch_schedule) = {
216            let epoch_info = remote_client.get_epoch_info().await?;
217            let epoch_schedule = remote_client.get_epoch_schedule().await?;
218            (epoch_info, epoch_schedule)
219        };
220        epoch_info.transaction_count = None;
221
222        self.with_svm_writer(move |svm_writer| {
223            svm_writer.initialize(epoch_info, epoch_schedule);
224        });
225        Ok(())
226    }
227}
228
229/// Functions for getting accounts from the underlying SurfnetSvm instance or remote client
230impl SurfnetSvmLocker {
231    /// Filters the downloaded account result to remove accounts owned by offline owners.
232    fn filter_downloaded_account_result(
233        requested_pubkey: &Pubkey,
234        result: GetAccountResult,
235        offline_owners: &[Pubkey],
236    ) -> GetAccountResult {
237        match result {
238            GetAccountResult::FoundAccount(_, account, _)
239            | GetAccountResult::FoundProgramAccount((_, account), _)
240            | GetAccountResult::FoundTokenAccount((_, account), _)
241                if offline_owners.contains(&account.owner) =>
242            {
243                GetAccountResult::None(*requested_pubkey)
244            }
245            other => other,
246        }
247    }
248
249    /// Retrieves a local account from the SVM cache, returning a contextualized result.
250    pub fn get_account_local(&self, pubkey: &Pubkey) -> SvmAccessContext<GetAccountResult> {
251        self.with_contextualized_svm_reader(|svm_reader| {
252            return svm_reader.inner.get_account_result(pubkey).unwrap();
253        })
254    }
255
256    /// Attempts local retrieval, then fetches from remote if missing, returning a contextualized result.
257    ///
258    /// Does not fetch from remote if the account has been explicitly blocked from remote downloads.
259    pub async fn get_account_local_then_remote(
260        &self,
261        client: &SurfnetRemoteClient,
262        pubkey: &Pubkey,
263        commitment_config: CommitmentConfig,
264    ) -> SurfpoolContextualizedResult<GetAccountResult> {
265        let result = self.get_account_local(pubkey);
266
267        if result.inner.is_none() {
268            let is_offline = self.is_account_offline(pubkey);
269
270            if !is_offline {
271                let offline_owners = self.get_offline_account_owners();
272                let remote_account = client.get_account(pubkey, commitment_config).await?;
273                Ok(
274                    result.with_new_value(Self::filter_downloaded_account_result(
275                        pubkey,
276                        remote_account,
277                        &offline_owners,
278                    )),
279                )
280            } else {
281                Ok(result)
282            }
283        } else {
284            Ok(result)
285        }
286    }
287
288    /// Retrieves an account, using local or remote based on context, applying a default factory if provided.
289    pub async fn get_account(
290        &self,
291        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
292        pubkey: &Pubkey,
293        factory: Option<AccountFactory>,
294    ) -> SurfpoolContextualizedResult<GetAccountResult> {
295        let result = if let Some((remote_client, commitment_config)) = remote_ctx {
296            self.get_account_local_then_remote(remote_client, pubkey, *commitment_config)
297                .await?
298        } else {
299            self.get_account_local(pubkey)
300        };
301
302        match (&result.inner, factory) {
303            (&GetAccountResult::None(_), Some(factory)) => {
304                let default = factory(self.clone());
305                Ok(result.with_new_value(default))
306            }
307            _ => Ok(result),
308        }
309    }
310    /// Retrieves multiple accounts from local cache, returning a contextualized result.
311    pub fn get_multiple_accounts_local(
312        &self,
313        pubkeys: &[Pubkey],
314    ) -> SvmAccessContext<Vec<GetAccountResult>> {
315        self.with_contextualized_svm_reader(|svm_reader| {
316            let mut accounts = vec![];
317
318            for pubkey in pubkeys {
319                let result = svm_reader.inner.get_account_result(pubkey).unwrap();
320                if result.is_none() {};
321                accounts.push(result);
322            }
323            accounts
324        })
325    }
326
327    /// Retrieves multiple accounts from local storage, with remote fallback for missing accounts.
328    ///
329    /// Returns accounts in the same order as the input `pubkeys` array. Accounts found locally
330    /// are returned as-is; accounts not found locally are fetched from the remote RPC client.
331    /// Accounts that have been marked offline are not fetched from remote.
332    pub async fn get_multiple_accounts_with_remote_fallback(
333        &self,
334        client: &SurfnetRemoteClient,
335        pubkeys: &[Pubkey],
336        commitment_config: CommitmentConfig,
337    ) -> SurfpoolContextualizedResult<Vec<GetAccountResult>> {
338        let SvmAccessContext {
339            slot,
340            latest_epoch_info,
341            latest_blockhash,
342            inner: local_results,
343        } = self.get_multiple_accounts_local(pubkeys);
344
345        // Collect missing pubkeys that are not offline (local_results is already in correct order from pubkeys).
346        let mut missing_accounts = Vec::new();
347        for result in &local_results {
348            let GetAccountResult::None(pubkey) = result else {
349                continue;
350            };
351            if self.is_account_offline(pubkey) {
352                continue;
353            }
354            missing_accounts.push(*pubkey);
355        }
356
357        if missing_accounts.is_empty() {
358            // All accounts found locally, already in correct order
359            return Ok(SvmAccessContext::new(
360                slot,
361                latest_epoch_info,
362                latest_blockhash,
363                local_results,
364            ));
365        }
366        debug!(
367            "Missing accounts will be fetched: {}",
368            missing_accounts.iter().join(", ")
369        );
370        // Fetch missing accounts from remote
371        let remote_results = client
372            .get_multiple_accounts(&missing_accounts, commitment_config)
373            .await?;
374
375        // Build map of pubkey -> remote result for O(1) lookup
376        let offline_owners = self.get_offline_account_owners();
377        let remote_map: HashMap<Pubkey, GetAccountResult> = missing_accounts
378            .iter()
379            .copied()
380            .zip(remote_results.into_iter())
381            .map(|(requested_pubkey, result)| {
382                (
383                    requested_pubkey,
384                    Self::filter_downloaded_account_result(
385                        &requested_pubkey,
386                        result,
387                        &offline_owners,
388                    ),
389                )
390            })
391            .collect();
392
393        // Replace None entries with remote results while preserving order
394        // We iterate through original pubkeys array to ensure order is explicit
395        let combined_results: Vec<GetAccountResult> = pubkeys
396            .iter()
397            .zip(local_results.into_iter())
398            .map(|(pubkey, local_result)| {
399                match local_result {
400                    GetAccountResult::None(_) => remote_map
401                        .get(pubkey)
402                        .cloned()
403                        .unwrap_or(GetAccountResult::None(*pubkey)),
404                    found => {
405                        debug!("Keeping local account: {}", pubkey);
406                        found
407                    } // Keep found accounts (no clone, just move)
408                }
409            })
410            .collect();
411
412        Ok(SvmAccessContext::new(
413            slot,
414            latest_epoch_info,
415            latest_blockhash,
416            combined_results,
417        ))
418    }
419
420    /// Retrieves multiple accounts, using local or remote context and applying factory defaults if provided.
421    pub async fn get_multiple_accounts(
422        &self,
423        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
424        pubkeys: &[Pubkey],
425        factory: Option<AccountFactory>,
426    ) -> SurfpoolContextualizedResult<Vec<GetAccountResult>> {
427        let results = if let Some((remote_client, commitment_config)) = remote_ctx {
428            self.get_multiple_accounts_with_remote_fallback(
429                remote_client,
430                pubkeys,
431                *commitment_config,
432            )
433            .await?
434        } else {
435            self.get_multiple_accounts_local(pubkeys)
436        };
437
438        let mut combined = Vec::with_capacity(results.inner.len());
439        for result in results.inner.clone() {
440            match (&result, &factory) {
441                (&GetAccountResult::None(_), Some(factory)) => {
442                    let default = factory(self.clone());
443                    combined.push(default);
444                }
445                _ => combined.push(result),
446            }
447        }
448        Ok(results.with_new_value(combined))
449    }
450
451    /// Loads accounts from a snapshot into the SVM.
452    ///
453    /// This method should be called before geyser plugins start to ensure they receive
454    /// the account updates with `is_startup=true`.
455    ///
456    /// # Arguments
457    /// * `snapshot` - A map of pubkey strings to optional account snapshots.
458    ///   - If the value is Some(AccountSnapshot), the account is loaded directly.
459    ///   - If the value is None, the account is fetched from the remote RPC (if available).
460    /// * `remote_client` - Optional remote RPC client to fetch None accounts.
461    /// * `commitment_config` - Commitment level for remote RPC calls.
462    ///
463    /// # Returns
464    /// The number of accounts successfully loaded.
465    pub async fn load_snapshot(
466        &self,
467        snapshot: &BTreeMap<String, Option<AccountSnapshot>>,
468        remote_client: Option<&SurfnetRemoteClient>,
469        commitment_config: CommitmentConfig,
470    ) -> SurfpoolResult<usize> {
471        use std::str::FromStr;
472
473        use base64::{Engine, prelude::BASE64_STANDARD};
474
475        let mut loaded_count = 0;
476
477        // Separate accounts into those with data and those needing remote fetch
478        let mut accounts_to_load: Vec<(Pubkey, Account)> = Vec::new();
479        let mut pubkeys_to_fetch: Vec<Pubkey> = Vec::new();
480
481        for (pubkey_str, account_snapshot_opt) in snapshot.iter() {
482            let pubkey = match Pubkey::from_str(pubkey_str) {
483                Ok(pk) => pk,
484                Err(e) => {
485                    self.with_svm_reader(|svm| {
486                        let _ = svm.simnet_events_tx.send(SimnetEvent::warn(format!(
487                            "Skipping invalid pubkey '{}' in snapshot: {}",
488                            pubkey_str, e
489                        )));
490                    });
491                    continue;
492                }
493            };
494
495            match account_snapshot_opt {
496                Some(account_snapshot) => {
497                    // Decode base64 data
498                    let data = match BASE64_STANDARD.decode(&account_snapshot.data) {
499                        Ok(d) => d,
500                        Err(e) => {
501                            self.with_svm_reader(|svm| {
502                                let _ = svm.simnet_events_tx.send(SimnetEvent::warn(format!(
503                                    "Skipping account '{}': failed to decode base64 data: {}",
504                                    pubkey_str, e
505                                )));
506                            });
507                            continue;
508                        }
509                    };
510
511                    // Parse owner pubkey
512                    let owner = match Pubkey::from_str(&account_snapshot.owner) {
513                        Ok(pk) => pk,
514                        Err(e) => {
515                            self.with_svm_reader(|svm| {
516                                let _ = svm.simnet_events_tx.send(SimnetEvent::warn(format!(
517                                    "Skipping account '{}': invalid owner pubkey: {}",
518                                    pubkey_str, e
519                                )));
520                            });
521                            continue;
522                        }
523                    };
524
525                    // Create the account
526                    let account = Account {
527                        lamports: account_snapshot.lamports,
528                        data,
529                        owner,
530                        executable: account_snapshot.executable,
531                        rent_epoch: account_snapshot.rent_epoch,
532                    };
533
534                    accounts_to_load.push((pubkey, account));
535                }
536                None => {
537                    // Queue for remote fetch if client is available
538                    if remote_client.is_some() {
539                        pubkeys_to_fetch.push(pubkey);
540                    }
541                }
542            }
543        }
544
545        // Fetch None accounts from remote RPC if client is available
546        if let Some(client) = remote_client {
547            if !pubkeys_to_fetch.is_empty() {
548                self.with_svm_reader(|svm| {
549                    let _ = svm.simnet_events_tx.send(SimnetEvent::info(format!(
550                        "Fetching {} accounts from remote RPC for snapshot",
551                        pubkeys_to_fetch.len()
552                    )));
553                });
554
555                match client
556                    .get_multiple_accounts(&pubkeys_to_fetch, commitment_config)
557                    .await
558                {
559                    Ok(remote_results) => {
560                        for (pubkey, result) in pubkeys_to_fetch.iter().zip(remote_results) {
561                            match result {
562                                GetAccountResult::FoundAccount(_, account, _) => {
563                                    accounts_to_load.push((*pubkey, account));
564                                }
565                                GetAccountResult::FoundProgramAccount(
566                                    (program_pubkey, program_account),
567                                    (data_pubkey, data_account_opt),
568                                ) => {
569                                    accounts_to_load.push((program_pubkey, program_account));
570                                    if let Some(data_account) = data_account_opt {
571                                        accounts_to_load.push((data_pubkey, data_account));
572                                    }
573                                }
574                                GetAccountResult::FoundTokenAccount(
575                                    (token_pubkey, token_account),
576                                    (mint_pubkey, mint_account_opt),
577                                ) => {
578                                    accounts_to_load.push((token_pubkey, token_account));
579                                    if let Some(mint_account) = mint_account_opt {
580                                        accounts_to_load.push((mint_pubkey, mint_account));
581                                    }
582                                }
583                                GetAccountResult::None(_) => {
584                                    // Account not found on remote, skip
585                                }
586                            }
587                        }
588                    }
589                    Err(e) => {
590                        self.with_svm_reader(|svm| {
591                            let _ = svm.simnet_events_tx.send(SimnetEvent::warn(format!(
592                                "Failed to fetch some accounts from remote: {}",
593                                e
594                            )));
595                        });
596                    }
597                }
598            }
599        }
600
601        // Load all accounts into the SVM
602        self.with_svm_writer(|svm| {
603            let slot = svm.get_latest_absolute_slot();
604
605            for (pubkey, account) in accounts_to_load {
606                if let Err(e) = svm.set_account(&pubkey, account.clone()) {
607                    let _ = svm.simnet_events_tx.send(SimnetEvent::warn(format!(
608                        "Failed to set account '{}': {}",
609                        pubkey, e
610                    )));
611                    continue;
612                }
613
614                // Send startup account update to geyser
615                let write_version = svm.increment_write_version();
616                let _ = svm.geyser_events_tx.send(GeyserEvent::StartupAccountUpdate(
617                    GeyserAccountUpdate::startup_update(pubkey, account, slot, write_version),
618                ));
619
620                loaded_count += 1;
621            }
622        });
623
624        Ok(loaded_count)
625    }
626
627    /// Retrieves largest accounts from local cache, returning a contextualized result.
628    pub fn get_largest_accounts_local(
629        &self,
630        config: RpcLargestAccountsConfig,
631    ) -> SurfpoolContextualizedResult<Vec<RpcAccountBalance>> {
632        let res: Vec<RpcAccountBalance> = self.with_svm_reader(|svm_reader| {
633            let non_circulating_accounts: Vec<_> = svm_reader
634                .non_circulating_accounts
635                .iter()
636                .flat_map(|acct| verify_pubkey(acct))
637                .collect();
638
639            let ordered_accounts = svm_reader
640                .get_all_accounts()?
641                .into_iter()
642                .sorted_by(|a, b| b.1.lamports().cmp(&a.1.lamports()))
643                .collect::<Vec<_>>();
644            let ordered_filtered_accounts = match config.filter {
645                Some(RpcLargestAccountsFilter::NonCirculating) => ordered_accounts
646                    .into_iter()
647                    .filter(|(pubkey, _)| non_circulating_accounts.contains(pubkey))
648                    .collect::<Vec<_>>(),
649                Some(RpcLargestAccountsFilter::Circulating) => ordered_accounts
650                    .into_iter()
651                    .filter(|(pubkey, _)| !non_circulating_accounts.contains(pubkey))
652                    .collect::<Vec<_>>(),
653                None => ordered_accounts,
654            };
655
656            Ok::<Vec<RpcAccountBalance>, SurfpoolError>(
657                ordered_filtered_accounts
658                    .iter()
659                    .take(20)
660                    .map(|(pubkey, account)| RpcAccountBalance {
661                        address: pubkey.to_string(),
662                        lamports: account.lamports(),
663                    })
664                    .collect(),
665            )
666        })?;
667        Ok(self.with_contextualized_svm_reader(|_| res.to_owned()))
668    }
669
670    pub async fn get_largest_accounts_local_then_remote(
671        &self,
672        client: &SurfnetRemoteClient,
673        config: RpcLargestAccountsConfig,
674        commitment_config: CommitmentConfig,
675    ) -> SurfpoolContextualizedResult<Vec<RpcAccountBalance>> {
676        // get all non-circulating and circulating pubkeys from the remote client first,
677        // and insert them locally
678        {
679            let remote_non_circulating_pubkeys_result = client
680                .get_largest_accounts(Some(RpcLargestAccountsConfig {
681                    filter: Some(RpcLargestAccountsFilter::NonCirculating),
682                    ..config
683                }))
684                .await?;
685
686            let (mut remote_non_circulating_pubkeys, mut remote_circulating_pubkeys) =
687                match remote_non_circulating_pubkeys_result {
688                    RemoteRpcResult::Ok(non_circulating_accounts) => {
689                        let remote_circulating_pubkeys_result = client
690                            .get_largest_accounts(Some(RpcLargestAccountsConfig {
691                                filter: Some(RpcLargestAccountsFilter::Circulating),
692                                ..config
693                            }))
694                            .await?;
695
696                        let remote_circulating_pubkeys = match remote_circulating_pubkeys_result {
697                            RemoteRpcResult::Ok(circulating_accounts) => circulating_accounts,
698                            RemoteRpcResult::MethodNotSupported => {
699                                unreachable!()
700                            }
701                        };
702                        (
703                            non_circulating_accounts
704                                .iter()
705                                .map(|account_balance| verify_pubkey(&account_balance.address))
706                                .collect::<SurfpoolResult<Vec<_>>>()?,
707                            remote_circulating_pubkeys
708                                .iter()
709                                .map(|account_balance| verify_pubkey(&account_balance.address))
710                                .collect::<SurfpoolResult<Vec<_>>>()?,
711                        )
712                    }
713                    RemoteRpcResult::MethodNotSupported => {
714                        let tx = self.simnet_events_tx();
715                        let _ = tx.send(SimnetEvent::warn("The `getLargestAccounts` method was sent to the remote RPC, but this method isn't supported by your RPC provider. Only local accounts will be returned."));
716                        (vec![], vec![])
717                    }
718                };
719
720            let mut combined = Vec::with_capacity(
721                remote_non_circulating_pubkeys.len() + remote_circulating_pubkeys.len(),
722            );
723            combined.append(&mut remote_non_circulating_pubkeys);
724            combined.append(&mut remote_circulating_pubkeys);
725
726            let get_account_results = self
727                .get_multiple_accounts_with_remote_fallback(client, &combined, commitment_config)
728                .await?
729                .inner;
730
731            self.write_multiple_account_updates(&get_account_results);
732        }
733
734        // now that our local cache is aware of all large remote accounts, we can get the largest accounts locally
735        // and filter according to the config
736        self.get_largest_accounts_local(config)
737    }
738
739    pub async fn get_largest_accounts(
740        &self,
741        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
742        config: RpcLargestAccountsConfig,
743    ) -> SurfpoolContextualizedResult<Vec<RpcAccountBalance>> {
744        if let Some((remote_client, commitment_config)) = remote_ctx {
745            self.get_largest_accounts_local_then_remote(remote_client, config, *commitment_config)
746                .await
747        } else {
748            self.get_largest_accounts_local(config)
749        }
750    }
751
752    pub fn account_to_rpc_keyed_account<T: ReadableAccount + Send + Sync>(
753        &self,
754        pubkey: &Pubkey,
755        account: &T,
756        config: &RpcAccountInfoConfig,
757        token_mint: Option<Pubkey>,
758    ) -> RpcKeyedAccount {
759        self.with_svm_reader(|svm_reader| {
760            svm_reader.account_to_rpc_keyed_account(pubkey, account, config, token_mint)
761        })
762    }
763}
764
765/// Get signatures for Addresses
766impl SurfnetSvmLocker {
767    /// Returns local `getSignaturesForAddress` results in the same newest-first order expected by
768    /// the Solana RPC.
769    ///
770    /// The implementation has to do more than filter by slot:
771    /// - transactions are ordered by descending slot
772    /// - transactions within the same slot are ordered by their execution order in the block
773    /// - `before` and `until` are pagination boundaries in that final ordered stream
774    ///
775    /// To preserve those semantics, we first collect matching transactions, reconstruct their
776    /// intra-slot ordering from block headers, sort the full result stream, and only then apply
777    /// the `before` / `until` window followed by `limit`.
778    pub fn get_signatures_for_address_local(
779        &self,
780        pubkey: &Pubkey,
781        config: Option<&RpcSignaturesForAddressConfig>,
782    ) -> SvmAccessContext<Vec<RpcConfirmedTransactionStatusWithSignature>> {
783        let before = config.and_then(|c| c.before.as_ref());
784        let until = config.and_then(|c| c.until.as_ref());
785        let limit = config.and_then(|c| c.limit).unwrap_or(1000);
786        let min_context_slot = config.and_then(|c| c.min_context_slot).unwrap_or_default();
787
788        self.with_contextualized_svm_reader(move |svm_reader| {
789            let current_slot = svm_reader.get_latest_absolute_slot();
790
791            let sigs: Vec<_> = svm_reader
792                .transactions
793                .into_iter()
794                .map(|iter| {
795                    iter.filter_map(|(sig, status)| {
796                        let (
797                            TransactionWithStatusMeta {
798                                slot,
799                                transaction,
800                                meta,
801                            },
802                            _,
803                        ) = status
804                            .as_processed()
805                            .expect("expected processed transaction");
806
807                        if slot < min_context_slot {
808                            return None;
809                        }
810
811                        // Check if the pubkey is a signer
812
813                        if !transaction.message.static_account_keys().contains(pubkey) {
814                            return None;
815                        }
816
817                        // Determine confirmation status
818                        let confirmation_status = match current_slot {
819                            cs if cs == slot => SolanaTransactionConfirmationStatus::Processed,
820                            cs if cs < slot + FINALIZATION_SLOT_THRESHOLD => {
821                                SolanaTransactionConfirmationStatus::Confirmed
822                            }
823                            _ => SolanaTransactionConfirmationStatus::Finalized,
824                        };
825
826                        Some(RpcConfirmedTransactionStatusWithSignature {
827                            err: match meta.status {
828                                Ok(_) => None,
829                                Err(e) => Some(e.into()),
830                            },
831                            slot,
832                            memo: None,
833                            block_time: None,
834                            confirmation_status: Some(confirmation_status),
835                            signature: sig,
836                        })
837                    })
838                    .collect()
839                })
840                .unwrap_or_default();
841
842            // `getSignaturesForAddress` is ordered newest-first, but transactions that share a
843            // slot also need to preserve their execution order within that block.
844            let unique_slots: HashSet<u64> = sigs.iter().map(|s| s.slot).collect();
845            let mut sig_position: HashMap<String, usize> = HashMap::new();
846            for slot in unique_slots {
847                if let Ok(Some(block_header)) = svm_reader.blocks.get(&slot) {
848                    for (idx, block_sig) in block_header.signatures.iter().enumerate() {
849                        sig_position.insert(block_sig.to_string(), idx);
850                    }
851                }
852            }
853
854            let sigs: Vec<_> = sigs
855                .into_iter()
856                // Order from most recent to least recent so pagination boundaries
857                // can be applied against the exact transaction sequence.
858                .sorted_by(|a, b| {
859                    b.slot.cmp(&a.slot).then_with(|| {
860                        let a_pos = sig_position.get(&a.signature).unwrap_or(&usize::MAX);
861                        let b_pos = sig_position.get(&b.signature).unwrap_or(&usize::MAX);
862                        b_pos.cmp(&a_pos)
863                    })
864                })
865                .collect();
866
867            let window = {
868                // `before` and `until` are boundaries in the final ordered result stream, not
869                // just slot filters. We compute a [start..end) index range after sorting so
870                // same-slot pagination behaves correctly and `until` stays exclusive.
871                let start = match before {
872                    // `before` is exclusive, so we start one item after the boundary when it
873                    // exists. If it does not exist locally, the local window is empty.
874                    Some(before) => match sigs.iter().position(|sig| sig.signature.eq(before)) {
875                        Some(idx) => idx + 1,
876                        None => sigs.len(),
877                    },
878                    None => 0,
879                };
880
881                let end = match until {
882                    // `until` is also exclusive, so the boundary itself is not included. We only
883                    // search within `sigs[start..]` so the end boundary is resolved relative to the
884                    // already-trimmed start of the window. If it is missing, we keep the full tail.
885                    Some(until) => {
886                        match sigs[start..].iter().position(|sig| sig.signature.eq(until)) {
887                            Some(offset) => start + offset,
888                            None => sigs.len(),
889                        }
890                    }
891                    None => sigs.len(),
892                };
893                start..end
894            };
895
896            // Apply the pagination window first, then enforce the RPC limit on that slice.
897            sigs[window].iter().take(limit).cloned().collect()
898        })
899    }
900
901    pub async fn get_signatures_for_address_local_then_remote(
902        &self,
903        client: &SurfnetRemoteClient,
904        pubkey: &Pubkey,
905        config: Option<&RpcSignaturesForAddressConfig>,
906    ) -> SurfpoolContextualizedResult<Vec<RpcConfirmedTransactionStatusWithSignature>> {
907        let results = self.get_signatures_for_address_local(pubkey, config);
908        let limit = config.and_then(|c| c.limit).unwrap_or(1000);
909
910        let SvmAccessContext {
911            slot,
912            latest_blockhash,
913            latest_epoch_info,
914            inner: mut combined_results,
915        } = results;
916        if combined_results.len() < limit {
917            let mut remote_results = client.get_signatures_for_address(pubkey, config).await?;
918            combined_results.append(&mut remote_results);
919        }
920
921        Ok(SvmAccessContext::new(
922            slot,
923            latest_epoch_info,
924            latest_blockhash,
925            combined_results,
926        ))
927    }
928
929    pub async fn get_signatures_for_address(
930        &self,
931        remote_ctx: &Option<(SurfnetRemoteClient, ())>,
932        pubkey: &Pubkey,
933        config: Option<&RpcSignaturesForAddressConfig>,
934    ) -> SurfpoolContextualizedResult<Vec<RpcConfirmedTransactionStatusWithSignature>> {
935        let results = if let Some((remote_client, _)) = remote_ctx {
936            self.get_signatures_for_address_local_then_remote(remote_client, pubkey, config)
937                .await?
938        } else {
939            self.get_signatures_for_address_local(pubkey, config)
940        };
941
942        Ok(results)
943    }
944}
945
946/// Functions for getting transactions from the underlying SurfnetSvm instance or remote client
947impl SurfnetSvmLocker {
948    /// Retrieves a transaction by signature, using local or remote based on context.
949    pub async fn get_transaction(
950        &self,
951        remote_ctx: &Option<SurfnetRemoteClient>,
952        signature: &Signature,
953        config: RpcTransactionConfig,
954    ) -> SurfpoolResult<GetTransactionResult> {
955        if let Some(remote_client) = remote_ctx {
956            self.get_transaction_local_then_remote(remote_client, signature, config)
957                .await
958        } else {
959            self.get_transaction_local(signature, &config)
960        }
961    }
962
963    /// Retrieves a transaction from local cache, returning a contextualized result.
964    pub fn get_transaction_local(
965        &self,
966        signature: &Signature,
967        config: &RpcTransactionConfig,
968    ) -> SurfpoolResult<GetTransactionResult> {
969        self.with_svm_reader(|svm_reader| {
970            let latest_absolute_slot = svm_reader.get_latest_absolute_slot();
971
972            let Some(entry) = svm_reader.transactions.get(&signature.to_string())? else {
973                return Ok(GetTransactionResult::None(*signature));
974            };
975
976            let (transaction_with_status_meta, _) = entry.expect_processed();
977            let slot = transaction_with_status_meta.slot;
978            let block_time = svm_reader
979                .blocks
980                .get(&slot)?
981                .map(|b| (b.block_time / 1_000) as UnixTimestamp)
982                .unwrap_or(0);
983            let encoded = transaction_with_status_meta.encode(
984                config.encoding.unwrap_or(UiTransactionEncoding::JsonParsed),
985                config.max_supported_transaction_version,
986                true,
987            )?;
988            Ok(GetTransactionResult::found_transaction(
989                *signature,
990                EncodedConfirmedTransactionWithStatusMeta {
991                    slot,
992                    transaction: encoded,
993                    block_time: Some(block_time),
994                },
995                latest_absolute_slot,
996            ))
997        })
998    }
999
1000    /// Retrieves a transaction locally then from remote if missing, returning a contextualized result.
1001    pub async fn get_transaction_local_then_remote(
1002        &self,
1003        client: &SurfnetRemoteClient,
1004        signature: &Signature,
1005        config: RpcTransactionConfig,
1006    ) -> SurfpoolResult<GetTransactionResult> {
1007        let local_result = self.get_transaction_local(signature, &config)?;
1008        let latest_absolute_slot = self.get_latest_absolute_slot();
1009        if local_result.is_none() {
1010            Ok(client
1011                .get_transaction(*signature, config, latest_absolute_slot)
1012                .await)
1013        } else {
1014            Ok(local_result)
1015        }
1016    }
1017}
1018
1019/// Functions for simulating and processing transactions in the underlying SurfnetSvm instance
1020impl SurfnetSvmLocker {
1021    /// Simulates a transaction on the SVM, returning detailed info or failure metadata.
1022    #[allow(clippy::result_large_err)]
1023    pub fn simulate_transaction(
1024        &self,
1025        transaction: VersionedTransaction,
1026        sigverify: bool,
1027    ) -> Result<SimulatedTransactionInfo, FailedTransactionMetadata> {
1028        self.with_svm_reader(move |svm_reader| {
1029            svm_reader.simulate_transaction(transaction, sigverify)
1030        })
1031    }
1032
1033    pub fn is_instruction_profiling_enabled(&self) -> bool {
1034        self.with_svm_reader(|svm_reader| svm_reader.instruction_profiling_enabled)
1035    }
1036
1037    pub fn get_profiling_map_capacity(&self) -> usize {
1038        self.with_svm_reader(|svm_reader| svm_reader.max_profiles)
1039    }
1040
1041    pub async fn process_transaction(
1042        &self,
1043        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
1044        transaction: VersionedTransaction,
1045        status_tx: Sender<TransactionStatusEvent>,
1046        skip_preflight: bool,
1047        sigverify: bool,
1048    ) -> SurfpoolResult<()> {
1049        let do_propagate_status_updates = true;
1050        let signature = transaction.signatures[0];
1051        let profile_result = match self
1052            .fetch_all_tx_accounts_then_process_tx_returning_profile_res(
1053                remote_ctx,
1054                transaction,
1055                &status_tx,
1056                skip_preflight,
1057                sigverify,
1058                do_propagate_status_updates,
1059            )
1060            .await
1061        {
1062            Ok(result) => result,
1063            Err(e) => {
1064                // Ensure the status channel always receives a response to prevent
1065                // the RPC handler from hanging on recv() when errors occur during
1066                // account fetching, ALT resolution, or other pre-processing steps.
1067                // This is critical for issue #454 where program close stops block production.
1068                //
1069                // AccountLoadedTwice errors should go through SimulationFailure to produce
1070                // Agave-compatible JSON-RPC error format with structured `err` and `data` fields.
1071                let err_str = e.to_string();
1072                if err_str.contains("Account loaded twice") {
1073                    let _ = status_tx.try_send(TransactionStatusEvent::SimulationFailure((
1074                        TransactionError::AccountLoadedTwice,
1075                        surfpool_types::TransactionMetadata::default(),
1076                    )));
1077                } else {
1078                    let _ =
1079                        status_tx.try_send(TransactionStatusEvent::VerificationFailure(err_str));
1080                }
1081                return Err(e);
1082            }
1083        };
1084
1085        self.with_svm_writer(|svm_writer| {
1086            svm_writer.write_executed_profile_result(signature, profile_result)
1087        })?;
1088
1089        Ok(())
1090    }
1091
1092    pub async fn profile_transaction(
1093        &self,
1094        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
1095        transaction: VersionedTransaction,
1096        tag: Option<String>,
1097    ) -> SurfpoolContextualizedResult<Uuid> {
1098        // Use clone_for_profiling to wrap all storage fields with overlay storage,
1099        // ensuring mutations during profiling don't affect the underlying database
1100        let svm_clone = self.with_svm_reader(|svm_reader| svm_reader.clone_for_profiling());
1101
1102        let svm_locker = SurfnetSvmLocker::new(svm_clone);
1103
1104        let (status_tx, _) = crossbeam_channel::unbounded();
1105
1106        let skip_preflight = true; // skip preflight checks during transaction profiling
1107        let sigverify = true; // do verify signatures during transaction profiling
1108        let do_propagate_status_updates = false; // don't propagate status updates during transaction profiling
1109        let mut profile_result = svm_locker
1110            .fetch_all_tx_accounts_then_process_tx_returning_profile_res(
1111                remote_ctx,
1112                transaction,
1113                &status_tx,
1114                skip_preflight,
1115                sigverify,
1116                do_propagate_status_updates,
1117            )
1118            .await?;
1119
1120        let uuid = Uuid::new_v4();
1121        profile_result.key = UuidOrSignature::Uuid(uuid);
1122
1123        self.with_svm_writer(|svm_writer| {
1124            svm_writer.write_simulated_profile_result(uuid, tag, profile_result)
1125        })?;
1126
1127        Ok(self.with_contextualized_svm_reader(|_| uuid))
1128    }
1129
1130    async fn fetch_all_tx_accounts_then_process_tx_returning_profile_res(
1131        &self,
1132        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
1133        transaction: VersionedTransaction,
1134        status_tx: &Sender<TransactionStatusEvent>,
1135        skip_preflight: bool,
1136        sigverify: bool,
1137        do_propagate: bool,
1138    ) -> SurfpoolResult<KeyedProfileResult> {
1139        let signature = transaction.signatures[0];
1140
1141        // Sigverify the transaction upfront before doing any account fetching or other pre-processing.
1142        if sigverify {
1143            self.with_svm_reader(|svm_reader| svm_reader.sigverify(&transaction))
1144                .map_err(|e| Into::<SurfpoolError>::into(e.err))?;
1145        }
1146
1147        let latest_absolute_slot = self.with_svm_writer(|svm_writer| {
1148            let latest_absolute_slot = svm_writer.get_latest_absolute_slot();
1149            svm_writer.notify_signature_subscribers(
1150                SignatureSubscriptionType::received(),
1151                &signature,
1152                latest_absolute_slot,
1153                None,
1154            );
1155
1156            latest_absolute_slot
1157        });
1158
1159        // find accounts that are needed for this transaction but are missing from the local
1160        // svm cache, fetch them from the RPC, and insert them locally
1161        let tx_loaded_addresses = self
1162            .get_loaded_addresses(remote_ctx, &transaction.message)
1163            .await?;
1164
1165        // Check for duplicate accounts between static keys and ALT-loaded addresses.
1166        // Agave rejects such transactions pre-execution with AccountLoadedTwice.
1167        if let Some(ref loaded) = tx_loaded_addresses {
1168            let static_keys: HashSet<&Pubkey> =
1169                transaction.message.static_account_keys().iter().collect();
1170            for loaded_key in loaded.all_loaded_addresses() {
1171                if static_keys.contains(loaded_key) {
1172                    return Err(TransactionError::AccountLoadedTwice.into());
1173                }
1174            }
1175        }
1176
1177        // we don't want the pubkeys of the address lookup tables to be included in the transaction accounts,
1178        // but we do want the pubkeys of the accounts _loaded_ by the ALT to be in the transaction accounts.
1179        let transaction_accounts = self.get_pubkeys_from_message(
1180            &transaction.message,
1181            tx_loaded_addresses
1182                .as_ref()
1183                .map(|l| l.all_loaded_addresses()),
1184        );
1185        debug!(
1186            "Transaction {} accounts inputs: {}",
1187            transaction.get_signature(),
1188            transaction_accounts.iter().join(", ")
1189        );
1190
1191        let account_updates = self
1192            .get_multiple_accounts(remote_ctx, &transaction_accounts, None)
1193            .await?
1194            .inner;
1195
1196        // We also need the pubkeys of the ALTs to be pulled from the remote, so we'll do a fetch for them
1197        let alt_account_updates = self
1198            .get_multiple_accounts(
1199                remote_ctx,
1200                &tx_loaded_addresses
1201                    .as_ref()
1202                    .map(|l| l.alt_addresses())
1203                    .unwrap_or_default(),
1204                None,
1205            )
1206            .await?
1207            .inner;
1208
1209        let readonly_account_states = transaction_accounts
1210            .iter()
1211            .enumerate()
1212            .filter_map(|(i, pubkey)| {
1213                if transaction.message.is_maybe_writable(i, None) {
1214                    None
1215                } else {
1216                    self.get_account_local(pubkey)
1217                        .inner
1218                        .map_account()
1219                        .ok()
1220                        .map(|a| (*pubkey, a))
1221                }
1222            })
1223            .collect::<HashMap<_, _>>();
1224
1225        self.with_svm_writer(|svm_writer| {
1226            for update in &account_updates {
1227                svm_writer.write_account_update(update.clone());
1228            }
1229            for update in alt_account_updates {
1230                svm_writer.write_account_update(update);
1231            }
1232        });
1233
1234        let pre_execution_capture = {
1235            let mut capture = ExecutionCapture::new();
1236            for account_update in account_updates.into_iter() {
1237                match account_update {
1238                    GetAccountResult::None(pubkey) => {
1239                        capture.insert(pubkey, None);
1240                    }
1241                    GetAccountResult::FoundAccount(pubkey, account, _)
1242                    | GetAccountResult::FoundProgramAccount((pubkey, account), _)
1243                    | GetAccountResult::FoundTokenAccount((pubkey, account), _) => {
1244                        capture.insert(pubkey, Some(account));
1245                    }
1246                }
1247            }
1248            capture
1249        };
1250
1251        let (accounts_before, token_accounts_before, token_programs) =
1252            self.with_svm_reader(|svm_reader| {
1253                let accounts_before = transaction_accounts
1254                    .iter()
1255                    .map(|p| svm_reader.inner.get_account(p))
1256                    .collect::<Result<Vec<Option<Account>>, SurfpoolError>>()?;
1257
1258                let token_accounts_before = transaction_accounts
1259                    .iter()
1260                    .enumerate()
1261                    .filter_map(|(i, p)| {
1262                        svm_reader
1263                            .token_accounts
1264                            .get(&p.to_string())
1265                            .ok()
1266                            .flatten()
1267                            .map(|a| (i, a))
1268                    })
1269                    .collect::<Vec<_>>();
1270
1271                let token_programs = token_accounts_before
1272                    .iter()
1273                    .map(|(i, ta)| {
1274                        svm_reader
1275                            .get_account(&transaction_accounts[*i])
1276                            .map(|res| res.map(|a| a.owner).unwrap_or(ta.token_program_id()))
1277                    })
1278                    .collect::<Result<Vec<_>, SurfpoolError>>()?;
1279
1280                Ok::<
1281                    (
1282                        Vec<Option<Account>>,
1283                        Vec<(usize, TokenAccount)>,
1284                        Vec<Pubkey>,
1285                    ),
1286                    SurfpoolError,
1287                >((accounts_before, token_accounts_before, token_programs))
1288            })?;
1289
1290        let loaded_addresses = tx_loaded_addresses.as_ref().map(|l| l.loaded_addresses());
1291
1292        let ix_profiles = if self.is_instruction_profiling_enabled() {
1293            match self
1294                .generate_instruction_profiles(
1295                    &transaction,
1296                    &transaction_accounts,
1297                    &tx_loaded_addresses,
1298                    &accounts_before,
1299                    &token_accounts_before,
1300                    &token_programs,
1301                    pre_execution_capture.clone(),
1302                    &status_tx,
1303                )
1304                .await
1305            {
1306                Ok(profiles) => profiles,
1307                Err(e) => {
1308                    let _ = self.simnet_events_tx().try_send(SimnetEvent::error(format!(
1309                        "Failed to generate instruction profiles: {}",
1310                        e
1311                    )));
1312                    None
1313                }
1314            }
1315        } else {
1316            None
1317        };
1318
1319        let profile_result = self
1320            .process_transaction_internal(
1321                transaction,
1322                skip_preflight,
1323                sigverify,
1324                &transaction_accounts,
1325                &loaded_addresses,
1326                &accounts_before,
1327                &token_accounts_before,
1328                &token_programs,
1329                pre_execution_capture,
1330                &status_tx,
1331                do_propagate,
1332            )
1333            .await?;
1334
1335        Ok(KeyedProfileResult::new(
1336            latest_absolute_slot,
1337            UuidOrSignature::Signature(signature),
1338            ix_profiles,
1339            profile_result,
1340            readonly_account_states,
1341        ))
1342    }
1343
1344    #[allow(clippy::too_many_arguments)]
1345    async fn generate_instruction_profiles(
1346        &self,
1347        transaction: &VersionedTransaction,
1348        transaction_accounts: &[Pubkey],
1349        loaded_addresses: &Option<TransactionLoadedAddresses>,
1350        accounts_before: &[Option<Account>],
1351        token_accounts_before: &[(usize, TokenAccount)],
1352        token_programs: &[Pubkey],
1353        pre_execution_capture: ExecutionCapture,
1354        status_tx: &Sender<TransactionStatusEvent>,
1355    ) -> SurfpoolResult<Option<Vec<ProfileResult>>> {
1356        let instructions = transaction.message.instructions();
1357        let ix_count = instructions.len();
1358        if ix_count == 0 {
1359            return Ok(None);
1360        }
1361        // Extract account categories from original transaction
1362
1363        let mut ix_profile_results: Vec<ProfileResult> = vec![];
1364
1365        for idx in 1..=ix_count {
1366            let partial_transaction_res = self.create_partial_transaction(
1367                instructions,
1368                transaction_accounts,
1369                transaction,
1370                idx,
1371                loaded_addresses,
1372            );
1373
1374            let mut ix_required_accounts = IndexSet::new();
1375            for &account_idx in &instructions[idx - 1].accounts {
1376                ix_required_accounts.insert(transaction_accounts[account_idx as usize]);
1377            }
1378            ix_required_accounts
1379                .insert(transaction_accounts[instructions[idx - 1].program_id_index as usize]);
1380
1381            let Some(partial_tx) = partial_transaction_res else {
1382                debug!("Unable to create partial transaction");
1383                return Ok(None);
1384            };
1385
1386            let mut previous_execution_captures = ExecutionCapture::new();
1387            let mut previous_cus = 0;
1388            let mut previous_log_count = 0;
1389            for result in ix_profile_results.iter() {
1390                previous_execution_captures.extend(result.post_execution_capture.clone());
1391                previous_cus += result.compute_units_consumed;
1392                previous_log_count += result.log_messages.as_ref().map(|m| m.len()).unwrap_or(0);
1393            }
1394
1395            let skip_preflight = true;
1396            let sigverify = false;
1397            let do_propagate = false;
1398
1399            let mut pre_execution_capture_cursor = pre_execution_capture.clone();
1400            // If a pre-execution capture was provided, any pubkeys that are in the capture
1401            // that we just took should be replaced with those from the pre-execution capture.
1402            let capture_keys: Vec<_> = pre_execution_capture_cursor.keys().cloned().collect();
1403            for pubkey in capture_keys.into_iter() {
1404                if let Some(pre_account) = previous_execution_captures.remove(&pubkey) {
1405                    // Replace the account with the pre-execution one
1406                    pre_execution_capture_cursor.insert(pubkey, pre_account);
1407                }
1408            }
1409            let mut svm_clone = self.with_svm_reader(|svm_reader| svm_reader.clone_for_profiling());
1410
1411            let (dummy_simnet_tx, _) = crossbeam_channel::bounded(1);
1412            let (dummy_geyser_tx, _) = crossbeam_channel::bounded(1);
1413            svm_clone.simnet_events_tx = dummy_simnet_tx;
1414            svm_clone.geyser_events_tx = dummy_geyser_tx;
1415
1416            let svm_locker = SurfnetSvmLocker::new(svm_clone);
1417            let mut profile_result = svm_locker
1418                .process_transaction_internal(
1419                    partial_tx,
1420                    skip_preflight,
1421                    sigverify,
1422                    transaction_accounts,
1423                    &loaded_addresses.as_ref().map(|l| l.loaded_addresses()),
1424                    accounts_before,
1425                    token_accounts_before,
1426                    token_programs,
1427                    pre_execution_capture_cursor,
1428                    status_tx,
1429                    do_propagate,
1430                )
1431                .await?;
1432
1433            profile_result
1434                .pre_execution_capture
1435                .retain(|pubkey, _| ix_required_accounts.contains(pubkey));
1436            profile_result
1437                .post_execution_capture
1438                .retain(|pubkey, _| ix_required_accounts.contains(pubkey));
1439
1440            profile_result.compute_units_consumed = profile_result
1441                .compute_units_consumed
1442                .saturating_sub(previous_cus);
1443            profile_result.log_messages = profile_result.log_messages.map(|logs| {
1444                logs.into_iter()
1445                    .skip(previous_log_count)
1446                    .collect::<Vec<_>>()
1447            });
1448
1449            ix_profile_results.push(profile_result);
1450        }
1451
1452        Ok(Some(ix_profile_results))
1453    }
1454
1455    fn handle_simulation_failure(
1456        &self,
1457        signature: Signature,
1458        failed_transaction_metadata: FailedTransactionMetadata,
1459        pre_execution_capture: ExecutionCapture,
1460        simulated_slot: Slot,
1461        status_tx: Sender<TransactionStatusEvent>,
1462        do_propagate: bool,
1463    ) -> ProfileResult {
1464        let FailedTransactionMetadata { err, meta } = failed_transaction_metadata;
1465
1466        let cus = meta.compute_units_consumed;
1467        let log_messages = meta.logs.clone();
1468        let err_string = err.to_string();
1469
1470        if do_propagate {
1471            let meta = convert_transaction_metadata_from_canonical(&meta);
1472            let simnet_events_tx = self.simnet_events_tx();
1473            let _ = simnet_events_tx.try_send(SimnetEvent::error(format!(
1474                "Transaction simulation failed: {}",
1475                err
1476            )));
1477
1478            self.with_svm_writer(|svm_writer| {
1479                svm_writer.notify_signature_subscribers(
1480                    SignatureSubscriptionType::processed(),
1481                    &signature,
1482                    simulated_slot,
1483                    Some(err.clone()),
1484                );
1485                svm_writer.notify_logs_subscribers(
1486                    &signature,
1487                    Some(err.clone()),
1488                    log_messages.clone(),
1489                    CommitmentLevel::Processed,
1490                );
1491            });
1492            let _ = status_tx.try_send(TransactionStatusEvent::SimulationFailure((err, meta)));
1493        }
1494        ProfileResult::new(
1495            pre_execution_capture,
1496            BTreeMap::new(),
1497            cus,
1498            Some(log_messages),
1499            Some(err_string),
1500        )
1501    }
1502
1503    #[allow(clippy::too_many_arguments)]
1504    fn handle_execution_failure(
1505        &self,
1506        failed_transaction_metadata: FailedTransactionMetadata,
1507        transaction: VersionedTransaction,
1508        simulated_slot: Slot,
1509        pubkeys_from_message: &[Pubkey],
1510        accounts_before: &[Option<Account>],
1511        token_accounts_before: &[(usize, TokenAccount)],
1512        token_programs: &[Pubkey],
1513        loaded_addresses: &Option<LoadedAddresses>,
1514        pre_execution_capture: ExecutionCapture,
1515        status_tx: Sender<TransactionStatusEvent>,
1516        do_propagate: bool,
1517    ) -> SurfpoolResult<ProfileResult> {
1518        let FailedTransactionMetadata { err, meta } = failed_transaction_metadata;
1519
1520        let cus = meta.compute_units_consumed;
1521        let log_messages = meta.logs.clone();
1522        let err_string = err.to_string();
1523        let signature = meta.signature;
1524
1525        let accounts_after = pubkeys_from_message
1526            .iter()
1527            .map(|p| self.with_svm_reader(|svm_reader| svm_reader.inner.get_account(p)))
1528            .collect::<SurfpoolResult<Vec<Option<Account>>>>()?;
1529
1530        for (pubkey, (before, after)) in pubkeys_from_message
1531            .iter()
1532            .zip(accounts_before.iter().zip(accounts_after.iter()))
1533        {
1534            if before.ne(&after) {
1535                self.with_svm_writer(|svm_writer| {
1536                    if let Some(after) = &after {
1537                        let _ = svm_writer.update_account_registries(pubkey, after);
1538                        svm_writer.notify_account_subscribers(pubkey, &after);
1539                        svm_writer.notify_program_subscribers(pubkey, &after);
1540                    } else {
1541                        svm_writer.notify_account_subscribers(pubkey, &Account::default());
1542                        svm_writer.notify_program_subscribers(pubkey, &Account::default());
1543                    }
1544                });
1545            }
1546        }
1547
1548        let token_mints = self
1549            .with_svm_reader(|svm_reader| {
1550                token_accounts_before
1551                    .iter()
1552                    .map(|(_, a)| {
1553                        svm_reader
1554                            .token_mints
1555                            .get(&a.mint().to_string())
1556                            .ok()
1557                            .flatten()
1558                            .ok_or(SurfpoolError::token_mint_not_found(a.mint()))
1559                    })
1560                    .collect::<Result<Vec<_>, SurfpoolError>>()
1561            })
1562            .unwrap_or_default();
1563
1564        if do_propagate {
1565            let meta_canonical = convert_transaction_metadata_from_canonical(&meta);
1566            let simnet_events_tx = self.simnet_events_tx();
1567            let _ = simnet_events_tx.try_send(SimnetEvent::error(format!(
1568                "Transaction execution failed: {}",
1569                err
1570            )));
1571            let _ = status_tx.try_send(TransactionStatusEvent::ExecutionFailure((
1572                err.clone(),
1573                meta_canonical.clone(),
1574            )));
1575
1576            self.with_svm_writer(|svm_writer| {
1577                let transaction_with_status_meta = TransactionWithStatusMeta::from_failure(
1578                    simulated_slot,
1579                    transaction.clone(),
1580                    &FailedTransactionMetadata {
1581                        err: err.clone(),
1582                        meta: meta.clone(),
1583                    },
1584                    accounts_before,
1585                    &accounts_after,
1586                    token_accounts_before,
1587                    token_mints,
1588                    token_programs,
1589                    loaded_addresses.clone().unwrap_or_default(),
1590                );
1591                svm_writer.transactions.store(
1592                    signature.to_string(),
1593                    SurfnetTransactionStatus::processed(
1594                        transaction_with_status_meta.clone(),
1595                        HashSet::new(),
1596                    ),
1597                )?;
1598
1599                let _ = svm_writer
1600                    .geyser_events_tx
1601                    .send(GeyserEvent::NotifyTransaction(
1602                        transaction_with_status_meta,
1603                        Some(transaction.clone()),
1604                    ));
1605
1606                svm_writer.transactions_queued_for_confirmation.push_back((
1607                    transaction.clone(),
1608                    status_tx.clone(),
1609                    Some(err.clone()),
1610                ));
1611
1612                svm_writer.notify_signature_subscribers(
1613                    SignatureSubscriptionType::processed(),
1614                    &signature,
1615                    simulated_slot,
1616                    Some(err.clone()),
1617                );
1618                svm_writer.notify_logs_subscribers(
1619                    &signature,
1620                    Some(err.clone()),
1621                    log_messages.clone(),
1622                    CommitmentLevel::Processed,
1623                );
1624                let _ = svm_writer
1625                    .simnet_events_tx
1626                    .try_send(SimnetEvent::transaction_processed(
1627                        meta_canonical,
1628                        Some(err),
1629                    ));
1630                Ok::<(), SurfpoolError>(())
1631            })?;
1632        }
1633        Ok(ProfileResult::new(
1634            pre_execution_capture,
1635            BTreeMap::new(),
1636            cus,
1637            Some(log_messages),
1638            Some(err_string),
1639        ))
1640    }
1641
1642    #[allow(clippy::too_many_arguments)]
1643    fn handle_execution_success(
1644        &self,
1645        transaction_metadata: TransactionMetadata,
1646        transaction: VersionedTransaction,
1647        simulated_slot: Slot,
1648        pubkeys_from_message: &[Pubkey],
1649        loaded_addresses: &Option<LoadedAddresses>,
1650        accounts_before: &[Option<Account>],
1651        token_accounts_before: &[(usize, TokenAccount)],
1652        token_programs: &[Pubkey],
1653        pre_execution_capture: ExecutionCapture,
1654        status_tx: &Sender<TransactionStatusEvent>,
1655        do_propagate: bool,
1656    ) -> SurfpoolResult<ProfileResult> {
1657        let cus = transaction_metadata.compute_units_consumed;
1658        let logs = transaction_metadata.logs.clone();
1659        let signature = transaction.signatures[0];
1660
1661        let post_execution_capture = self.with_svm_writer(|svm_writer| {
1662            let accounts_after = pubkeys_from_message
1663                .iter()
1664                .map(|p| svm_writer.inner.get_account_no_db(p))
1665                .collect::<Vec<Option<Account>>>();
1666            let (sanitized_transaction, versioned_transaction) = if do_propagate {
1667                (
1668                    SanitizedTransaction::try_create(
1669                        transaction.clone(),
1670                        transaction.message.hash(),
1671                        Some(false),
1672                        if let Some(loaded_addresses) = &loaded_addresses {
1673                            SimpleAddressLoader::Enabled(loaded_addresses.clone())
1674                        } else {
1675                            SimpleAddressLoader::Disabled
1676                        },
1677                        &HashSet::new(), // todo: provide reserved account keys
1678                    )
1679                    .ok(),
1680                    Some(transaction.clone()),
1681                )
1682            } else {
1683                (None, None)
1684            };
1685
1686            let mut mutated_account_pubkeys = HashSet::new();
1687            for (pubkey, (before, after)) in pubkeys_from_message
1688                .iter()
1689                .zip(accounts_before.iter().zip(accounts_after.clone()))
1690            {
1691                if before.ne(&after) {
1692                    mutated_account_pubkeys.insert(*pubkey);
1693                    let after = after.unwrap_or_default();
1694                    svm_writer.update_account_registries(pubkey, &after)?;
1695                    let write_version = svm_writer.increment_write_version();
1696
1697                    if let Some(sanitized_transaction) = sanitized_transaction.clone() {
1698                        let _ = svm_writer.geyser_events_tx.send(GeyserEvent::UpdateAccount(
1699                            GeyserAccountUpdate::transaction_update(
1700                                *pubkey,
1701                                after.clone(),
1702                                svm_writer.get_latest_absolute_slot(),
1703                                sanitized_transaction.clone(),
1704                                write_version,
1705                            ),
1706                        ));
1707                    }
1708                    svm_writer.notify_account_subscribers(pubkey, &after);
1709                    svm_writer.notify_program_subscribers(pubkey, &after);
1710                }
1711            }
1712
1713            let mut token_accounts_after = vec![];
1714            let mut post_execution_capture = BTreeMap::new();
1715            let mut post_token_program_ids = vec![];
1716
1717            for (i, (pubkey, account)) in pubkeys_from_message
1718                .iter()
1719                .zip(accounts_after.iter())
1720                .enumerate()
1721            {
1722                let token_account = svm_writer
1723                    .token_accounts
1724                    .get(&pubkey.to_string())
1725                    .ok()
1726                    .flatten();
1727                post_execution_capture.insert(*pubkey, account.clone());
1728
1729                if let Some(token_account) = token_account {
1730                    token_accounts_after.push((i, token_account));
1731                    post_token_program_ids.push(
1732                        account
1733                            .as_ref()
1734                            .map(|a| a.owner)
1735                            .unwrap_or(spl_token_interface::id()),
1736                    );
1737                }
1738            }
1739
1740            let token_mints = token_accounts_after
1741                .iter()
1742                .map(|(_, a)| {
1743                    svm_writer
1744                        .token_mints
1745                        .get(&a.mint().to_string())
1746                        .ok()
1747                        .flatten()
1748                        .ok_or(SurfpoolError::token_mint_not_found(a.mint()))
1749                })
1750                .collect::<Result<Vec<_>, SurfpoolError>>()?;
1751
1752            if do_propagate {
1753                let transaction_meta =
1754                    convert_transaction_metadata_from_canonical(&transaction_metadata);
1755                let transaction_with_status_meta = TransactionWithStatusMeta::new(
1756                    svm_writer.get_latest_absolute_slot(),
1757                    transaction.clone(),
1758                    transaction_metadata,
1759                    accounts_before,
1760                    &accounts_after,
1761                    token_accounts_before,
1762                    &token_accounts_after,
1763                    token_mints,
1764                    token_programs,
1765                    &post_token_program_ids,
1766                    loaded_addresses.clone().unwrap_or_default(),
1767                );
1768                svm_writer.transactions.store(
1769                    transaction_meta.signature.to_string(),
1770                    SurfnetTransactionStatus::processed(
1771                        transaction_with_status_meta.clone(),
1772                        mutated_account_pubkeys,
1773                    ),
1774                )?;
1775
1776                let _ = svm_writer
1777                    .simnet_events_tx
1778                    .try_send(SimnetEvent::transaction_processed(transaction_meta, None));
1779
1780                let _ = svm_writer
1781                    .geyser_events_tx
1782                    .send(GeyserEvent::NotifyTransaction(
1783                        transaction_with_status_meta,
1784                        versioned_transaction,
1785                    ));
1786
1787                svm_writer.transactions_queued_for_confirmation.push_back((
1788                    transaction.clone(),
1789                    status_tx.clone(),
1790                    None,
1791                ));
1792
1793                svm_writer.notify_signature_subscribers(
1794                    SignatureSubscriptionType::processed(),
1795                    &signature,
1796                    simulated_slot,
1797                    None,
1798                );
1799                svm_writer.notify_logs_subscribers(
1800                    &signature,
1801                    None,
1802                    logs.clone(),
1803                    CommitmentLevel::Processed,
1804                );
1805                let _ = status_tx.try_send(TransactionStatusEvent::Success(
1806                    TransactionConfirmationStatus::Processed,
1807                ));
1808            }
1809
1810            Ok::<ExecutionCapture, SurfpoolError>(post_execution_capture)
1811        })?;
1812
1813        Ok(ProfileResult::new(
1814            pre_execution_capture,
1815            post_execution_capture,
1816            cus,
1817            Some(logs),
1818            None,
1819        ))
1820    }
1821
1822    #[allow(clippy::too_many_arguments)]
1823    async fn process_transaction_internal(
1824        &self,
1825        transaction: VersionedTransaction,
1826        skip_preflight: bool,
1827        sigverify: bool,
1828        transaction_accounts: &[Pubkey],
1829        loaded_addresses: &Option<LoadedAddresses>,
1830        accounts_before: &[Option<Account>],
1831        token_accounts_before: &[(usize, TokenAccount)],
1832        token_programs: &[Pubkey],
1833        pre_execution_capture: ExecutionCapture,
1834        status_tx: &Sender<TransactionStatusEvent>,
1835        do_propagate: bool,
1836    ) -> SurfpoolResult<ProfileResult> {
1837        let res = match self
1838            .do_process_transaction_internal(transaction.clone(), skip_preflight, sigverify)
1839            .await
1840        {
1841            ProcessTransactionResult::Success(transaction_metadata) => self
1842                .handle_execution_success(
1843                    transaction_metadata,
1844                    transaction,
1845                    self.get_latest_absolute_slot(),
1846                    transaction_accounts,
1847                    loaded_addresses,
1848                    accounts_before,
1849                    token_accounts_before,
1850                    token_programs,
1851                    pre_execution_capture,
1852                    status_tx,
1853                    do_propagate,
1854                )?,
1855            ProcessTransactionResult::SimulationFailure(failed_transaction_metadata) => self
1856                .handle_simulation_failure(
1857                    transaction.signatures[0],
1858                    failed_transaction_metadata,
1859                    pre_execution_capture,
1860                    self.get_latest_absolute_slot(),
1861                    status_tx.clone(),
1862                    do_propagate,
1863                ),
1864            ProcessTransactionResult::ExecutionFailure(failed) => self.handle_execution_failure(
1865                failed,
1866                transaction,
1867                self.get_latest_absolute_slot(),
1868                transaction_accounts,
1869                accounts_before,
1870                token_accounts_before,
1871                token_programs,
1872                loaded_addresses,
1873                pre_execution_capture,
1874                status_tx.clone(),
1875                do_propagate,
1876            )?,
1877        };
1878        Ok(res)
1879    }
1880
1881    async fn do_process_transaction_internal(
1882        &self,
1883        transaction: VersionedTransaction,
1884        skip_preflight: bool,
1885        sigverify: bool,
1886    ) -> ProcessTransactionResult {
1887        // if not skipping preflight, simulate the transaction
1888        if !skip_preflight {
1889            if let Err(e) = self.with_svm_reader(|svm_reader| {
1890                svm_reader
1891                    .simulate_transaction(transaction.clone(), sigverify)
1892                    .map_err(ProcessTransactionResult::SimulationFailure)
1893            }) {
1894                return e;
1895            }
1896        }
1897
1898        match self.with_svm_writer(|svm_writer| {
1899            svm_writer
1900                .send_transaction(transaction, false /* cu_analysis_enabled */, sigverify)
1901                .map_err(|e| {
1902                    debug!("Transaction execution failure: {:?}", e.meta);
1903                    ProcessTransactionResult::ExecutionFailure(e)
1904                })
1905                .map(ProcessTransactionResult::Success)
1906        }) {
1907            Ok(res) => res,
1908            Err(res) => res,
1909        }
1910    }
1911}
1912
1913/// Functions for writing account updates to the underlying SurfnetSvm instance
1914impl SurfnetSvmLocker {
1915    /// Writes a single account update into the SVM state if present.
1916    pub fn write_account_update(&self, account_update: GetAccountResult) {
1917        if !account_update.requires_update() {
1918            return;
1919        }
1920
1921        self.with_svm_writer(move |svm_writer| {
1922            svm_writer.write_account_update(account_update.clone())
1923        })
1924    }
1925
1926    /// Writes multiple account updates into the SVM state when any are present.
1927    pub fn write_multiple_account_updates(&self, account_updates: &[GetAccountResult]) {
1928        if account_updates
1929            .iter()
1930            .all(|update| !update.requires_update())
1931        {
1932            return;
1933        }
1934
1935        self.with_svm_writer(move |svm_writer| {
1936            for update in account_updates {
1937                svm_writer.write_account_update(update.clone());
1938            }
1939        });
1940    }
1941
1942    /// Resets an account in the SVM state for refresh/streaming.
1943    ///
1944    /// This function coordinates the reset of accounts by removing them from the local cache,
1945    /// allowing them to be fetched fresh from mainnet on the next access.
1946    /// It handles program accounts (including their program data accounts) and can optionally
1947    /// cascade the reset to all accounts owned by a program.
1948    pub fn reset_account(
1949        &self,
1950        pubkey: Pubkey,
1951        include_owned_accounts: bool,
1952    ) -> SurfpoolResult<()> {
1953        let simnet_events_tx = self.simnet_events_tx();
1954        let _ = simnet_events_tx.send(SimnetEvent::info(format!(
1955            "Account {} will be reset",
1956            pubkey
1957        )));
1958        // Set the account online so it can be fetched from mainnet again.
1959        self.remove_offline_account(pubkey, include_owned_accounts)?;
1960
1961        self.with_svm_writer(move |svm_writer| {
1962            svm_writer.reset_account(&pubkey, include_owned_accounts)
1963        })
1964    }
1965
1966    /// Resets SVM state and clears all offline account entries.
1967    ///
1968    /// This function coordinates the reset of the entire network state.
1969    /// It also clears the offline account set so all accounts can be fetched from mainnet again.
1970    pub async fn reset_network(
1971        &self,
1972        remote_ctx: &Option<SurfnetRemoteClient>,
1973    ) -> SurfpoolResult<()> {
1974        let simnet_events_tx = self.simnet_events_tx();
1975        let _ = simnet_events_tx.send(SimnetEvent::info("Resetting network..."));
1976
1977        // Fetch epoch info from remote if available (similar to initialize)
1978        let (mut epoch_info, epoch_schedule) = if let Some(remote_client) = remote_ctx {
1979            (
1980                remote_client.get_epoch_info().await?,
1981                remote_client.get_epoch_schedule().await?,
1982            )
1983        } else {
1984            let epoch_schedule = SurfnetSvm::default_epoch_schedule();
1985            let epoch_info = SurfnetSvm::default_epoch_info(&epoch_schedule);
1986            (epoch_info, epoch_schedule)
1987        };
1988        epoch_info.transaction_count = None;
1989
1990        self.with_svm_writer(move |svm_writer| {
1991            let _ = svm_writer.reset_network(epoch_info, epoch_schedule);
1992            let _ = svm_writer.offline_accounts.clear();
1993        });
1994        Ok(())
1995    }
1996
1997    /// Marks an account as offline, preventing it from being downloaded from the remote RPC.
1998    ///
1999    /// When `include_owned_accounts` is enabled, this also marks accounts as offline that are already known locally.
2000    /// Accounts discovered later through direct remote fetches are rejected lazily if they are
2001    /// owned by an offline owner.
2002    pub async fn insert_offline_account(
2003        &self,
2004        pubkey: Pubkey,
2005        include_owned_accounts: bool,
2006    ) -> SurfpoolResult<()> {
2007        let simnet_events_tx = self.simnet_events_tx();
2008        let _ = simnet_events_tx.send(SimnetEvent::info(format!(
2009            "Account {} will be marked offline (excluded from remote downloads)",
2010            pubkey
2011        )));
2012
2013        self.with_svm_writer(move |svm_writer| {
2014            if let Err(e) = svm_writer.offline_accounts.store(
2015                pubkey.to_string(),
2016                OfflineAccountConfig {
2017                    include_owned_accounts,
2018                },
2019            ) {
2020                warn!("Failed to store offline account {}: {}", pubkey, e);
2021            }
2022        });
2023
2024        Ok(())
2025    }
2026
2027    /// Streams an account by its pubkey.
2028    pub fn stream_account(
2029        &self,
2030        pubkey: Pubkey,
2031        include_owned_accounts: bool,
2032    ) -> SurfpoolResult<()> {
2033        let simnet_events_tx = self.simnet_events_tx();
2034        let _ = simnet_events_tx.send(SimnetEvent::info(format!(
2035            "Account {} changes will be streamed",
2036            pubkey
2037        )));
2038        self.with_svm_writer(|svm_writer| {
2039            svm_writer
2040                .streamed_accounts
2041                .store(pubkey.to_string(), include_owned_accounts)
2042        })?;
2043        Ok(())
2044    }
2045
2046    pub fn get_streamed_accounts(&self) -> Vec<(String, bool)> {
2047        self.with_svm_reader(|svm_reader| {
2048            svm_reader
2049                .streamed_accounts
2050                .into_iter()
2051                .map(|iter| iter.collect())
2052                .unwrap_or_default()
2053        })
2054    }
2055
2056    /// Removes an account from the offline account set.
2057    ///
2058    /// This allows the account to be fetched from mainnet again if requested.
2059    /// This is useful when resetting an account for a refresh/stream operation.
2060    pub fn remove_offline_account(
2061        &self,
2062        pubkey: Pubkey,
2063        include_owned_accounts: bool,
2064    ) -> SurfpoolResult<()> {
2065        self.with_svm_writer(move |svm_writer| {
2066            if let Err(e) = svm_writer.offline_accounts.take(&pubkey.to_string()) {
2067                warn!("Failed to set account online {}: {}", pubkey, e);
2068            }
2069
2070            if include_owned_accounts {
2071                // Set online any locally-known accounts owned by this pubkey.
2072                // Uses the accounts_by_owner index as a fast lookup.
2073                let owned_pubkeys: Vec<Pubkey> = svm_writer
2074                    .accounts_by_owner
2075                    .get(&pubkey.to_string())
2076                    .ok()
2077                    .flatten()
2078                    .unwrap_or_default()
2079                    .iter()
2080                    .filter_map(|pk_str| pk_str.parse().ok())
2081                    .collect();
2082                for owned_pk in owned_pubkeys {
2083                    if let Err(e) = svm_writer.offline_accounts.take(&owned_pk.to_string()) {
2084                        warn!("Failed to set account online {}: {}", owned_pk, e);
2085                    }
2086                }
2087            }
2088        });
2089        Ok(())
2090    }
2091
2092    /// Returns true if the given pubkey is marked offline.
2093    pub fn is_account_offline(&self, pubkey: &Pubkey) -> bool {
2094        self.with_svm_reader(|svm_reader| {
2095            svm_reader
2096                .offline_accounts
2097                .contains_key(&pubkey.to_string())
2098                .unwrap_or(false)
2099        })
2100    }
2101
2102    /// Gets all owners whose accounts are marked offline.
2103    pub fn get_offline_account_owners(&self) -> Vec<Pubkey> {
2104        self.with_svm_reader(|svm_reader| {
2105            svm_reader
2106                .offline_accounts
2107                .into_iter()
2108                .unwrap_or_else(|e| {
2109                    warn!("Failed to iterate offline_accounts: {}", e);
2110                    Box::new(std::iter::empty())
2111                })
2112                .filter(|(_, config)| config.include_owned_accounts)
2113                .filter_map(|(k, _)| match k.parse() {
2114                    Ok(pk) => Some(pk),
2115                    Err(e) => {
2116                        warn!("Invalid pubkey in offline_accounts: {}: {}", k, e);
2117                        None
2118                    }
2119                })
2120                .collect()
2121        })
2122    }
2123
2124    /// Registers a scenario for execution
2125    pub fn register_scenario(
2126        &self,
2127        scenario: surfpool_types::Scenario,
2128        slot: Option<Slot>,
2129    ) -> SurfpoolResult<()> {
2130        self.with_svm_writer(move |svm_writer| svm_writer.register_scenario(scenario, slot))
2131    }
2132}
2133
2134/// Token account related functions
2135impl SurfnetSvmLocker {
2136    /// Fetches all token accounts for an owner, returning remote results and missing pubkeys contexts.
2137    pub async fn get_token_accounts_by_owner(
2138        &self,
2139        remote_ctx: &Option<SurfnetRemoteClient>,
2140        owner: Pubkey,
2141        filter: &TokenAccountsFilter,
2142        config: &RpcAccountInfoConfig,
2143    ) -> SurfpoolContextualizedResult<Vec<RpcKeyedAccount>> {
2144        if let Some(remote_client) = remote_ctx {
2145            self.get_token_accounts_by_owner_local_then_remote(owner, filter, remote_client, config)
2146                .await
2147        } else {
2148            self.get_token_accounts_by_owner_local(owner, filter, config)
2149        }
2150    }
2151
2152    pub fn get_token_accounts_by_owner_local(
2153        &self,
2154        owner: Pubkey,
2155        filter: &TokenAccountsFilter,
2156        config: &RpcAccountInfoConfig,
2157    ) -> SurfpoolContextualizedResult<Vec<RpcKeyedAccount>> {
2158        let result = self.with_contextualized_svm_reader(|svm_reader| {
2159            svm_reader
2160                .get_parsed_token_accounts_by_owner(&owner)
2161                .iter()
2162                .filter_map(|(pubkey, token_account)| {
2163                    svm_reader
2164                        .get_account(pubkey)
2165                        .map(|res| {
2166                            let Some(account) = res else {
2167                                return None;
2168                            };
2169                            if match filter {
2170                                TokenAccountsFilter::Mint(mint) => token_account.mint().eq(mint),
2171                                TokenAccountsFilter::ProgramId(program_id) => {
2172                                    account.owner.eq(program_id)
2173                                }
2174                            } {
2175                                Some(svm_reader.account_to_rpc_keyed_account(
2176                                    pubkey,
2177                                    &account,
2178                                    config,
2179                                    Some(token_account.mint()),
2180                                ))
2181                            } else {
2182                                None
2183                            }
2184                        })
2185                        .transpose()
2186                })
2187                .collect::<SurfpoolResult<Vec<_>>>()
2188        });
2189        let SvmAccessContext {
2190            slot,
2191            latest_epoch_info,
2192            latest_blockhash,
2193            inner: accounts,
2194        } = result;
2195        Ok(SvmAccessContext::new(
2196            slot,
2197            latest_epoch_info,
2198            latest_blockhash,
2199            accounts?,
2200        ))
2201    }
2202
2203    pub async fn get_token_accounts_by_owner_local_then_remote(
2204        &self,
2205        owner: Pubkey,
2206        filter: &TokenAccountsFilter,
2207        remote_client: &SurfnetRemoteClient,
2208        config: &RpcAccountInfoConfig,
2209    ) -> SurfpoolContextualizedResult<Vec<RpcKeyedAccount>> {
2210        let SvmAccessContext {
2211            slot,
2212            latest_epoch_info,
2213            latest_blockhash,
2214            inner: local_accounts,
2215        } = self.get_token_accounts_by_owner_local(owner, filter, config)?;
2216
2217        let remote_accounts = remote_client
2218            .get_token_accounts_by_owner(owner, filter, config)
2219            .await?;
2220
2221        let mut combined_accounts = remote_accounts;
2222
2223        for local_account in local_accounts {
2224            if let Some((pos, _)) = combined_accounts
2225                .iter()
2226                .find_position(|RpcKeyedAccount { pubkey, .. }| pubkey.eq(&local_account.pubkey))
2227            {
2228                combined_accounts[pos] = local_account;
2229            } else {
2230                combined_accounts.push(local_account);
2231            }
2232        }
2233
2234        Ok(SvmAccessContext::new(
2235            slot,
2236            latest_epoch_info,
2237            latest_blockhash,
2238            combined_accounts,
2239        ))
2240    }
2241
2242    pub async fn get_token_accounts_by_delegate(
2243        &self,
2244        remote_ctx: &Option<SurfnetRemoteClient>,
2245        delegate: Pubkey,
2246        filter: &TokenAccountsFilter,
2247        config: &RpcAccountInfoConfig,
2248    ) -> SurfpoolContextualizedResult<Vec<RpcKeyedAccount>> {
2249        // Validate that the program is supported if using ProgramId filter
2250        if let TokenAccountsFilter::ProgramId(program_id) = filter {
2251            if !is_supported_token_program(program_id) {
2252                return Err(SurfpoolError::unsupported_token_program(*program_id));
2253            }
2254        }
2255
2256        if let Some(remote_client) = remote_ctx {
2257            self.get_token_accounts_by_delegate_local_then_remote(
2258                delegate,
2259                filter,
2260                remote_client,
2261                config,
2262            )
2263            .await
2264        } else {
2265            self.get_token_accounts_by_delegate_local(delegate, filter, config)
2266        }
2267    }
2268}
2269
2270/// Token account by delegate related functions
2271impl SurfnetSvmLocker {
2272    pub fn get_token_accounts_by_delegate_local(
2273        &self,
2274        delegate: Pubkey,
2275        filter: &TokenAccountsFilter,
2276        config: &RpcAccountInfoConfig,
2277    ) -> SurfpoolContextualizedResult<Vec<RpcKeyedAccount>> {
2278        let result = self.with_contextualized_svm_reader(|svm_reader| {
2279            svm_reader
2280                .get_token_accounts_by_delegate(&delegate)
2281                .iter()
2282                .filter_map(|(pubkey, token_account)| {
2283                    svm_reader
2284                        .get_account(pubkey)
2285                        .map(|res| {
2286                            let Some(account) = res else {
2287                                return None;
2288                            };
2289                            let include = match filter {
2290                                TokenAccountsFilter::Mint(mint) => token_account.mint() == *mint,
2291                                TokenAccountsFilter::ProgramId(program_id) => {
2292                                    account.owner == *program_id
2293                                        && is_supported_token_program(program_id)
2294                                }
2295                            };
2296
2297                            if include {
2298                                Some(svm_reader.account_to_rpc_keyed_account(
2299                                    pubkey,
2300                                    &account,
2301                                    config,
2302                                    Some(token_account.mint()),
2303                                ))
2304                            } else {
2305                                None
2306                            }
2307                        })
2308                        .transpose()
2309                })
2310                .collect::<SurfpoolResult<Vec<_>>>()
2311        });
2312        let SvmAccessContext {
2313            slot,
2314            latest_epoch_info,
2315            latest_blockhash,
2316            inner: accounts,
2317        } = result;
2318        Ok(SvmAccessContext::new(
2319            slot,
2320            latest_epoch_info,
2321            latest_blockhash,
2322            accounts?,
2323        ))
2324    }
2325
2326    pub async fn get_token_accounts_by_delegate_local_then_remote(
2327        &self,
2328        delegate: Pubkey,
2329        filter: &TokenAccountsFilter,
2330        remote_client: &SurfnetRemoteClient,
2331        config: &RpcAccountInfoConfig,
2332    ) -> SurfpoolContextualizedResult<Vec<RpcKeyedAccount>> {
2333        let SvmAccessContext {
2334            slot,
2335            latest_epoch_info,
2336            latest_blockhash,
2337            inner: local_accounts,
2338        } = self.get_token_accounts_by_delegate_local(delegate, filter, config)?;
2339
2340        let remote_accounts = remote_client
2341            .get_token_accounts_by_delegate(delegate, filter, config)
2342            .await?;
2343
2344        let mut combined_accounts = remote_accounts;
2345
2346        for local_account in local_accounts {
2347            if let Some((pos, _)) = combined_accounts
2348                .iter()
2349                .find_position(|RpcKeyedAccount { pubkey, .. }| pubkey.eq(&local_account.pubkey))
2350            {
2351                // Replace remote account with local one (local takes precedence)
2352                combined_accounts[pos] = local_account;
2353            } else {
2354                // Add local account that wasn't found in remote results
2355                combined_accounts.push(local_account);
2356            }
2357        }
2358
2359        Ok(SvmAccessContext::new(
2360            slot,
2361            latest_epoch_info,
2362            latest_blockhash,
2363            combined_accounts,
2364        ))
2365    }
2366}
2367
2368/// Get largest account related account
2369impl SurfnetSvmLocker {
2370    pub fn get_token_largest_accounts_local(
2371        &self,
2372        mint: &Pubkey,
2373    ) -> SvmAccessContext<Vec<RpcTokenAccountBalance>> {
2374        self.with_contextualized_svm_reader(|svm_reader| {
2375            let token_accounts = svm_reader.get_token_accounts_by_mint(mint);
2376
2377            // get mint information to determine decimals
2378            let mint_decimals = if let Some(mint_account) =
2379                svm_reader.token_mints.get(&mint.to_string()).ok().flatten()
2380            {
2381                mint_account.decimals()
2382            } else {
2383                0
2384            };
2385
2386            // convert to RpcTokenAccountBalance and sort by balance
2387            let mut balances: Vec<RpcTokenAccountBalance> = token_accounts
2388                .into_iter()
2389                .map(|(pubkey, token_account)| RpcTokenAccountBalance {
2390                    address: pubkey.to_string(),
2391                    amount: UiTokenAmount {
2392                        amount: token_account.amount().to_string(),
2393                        decimals: mint_decimals,
2394                        ui_amount: Some(format_ui_amount(token_account.amount(), mint_decimals)),
2395                        ui_amount_string: format_ui_amount_string(
2396                            token_account.amount(),
2397                            mint_decimals,
2398                        ),
2399                    },
2400                })
2401                .collect();
2402
2403            // sort by amount in descending order
2404            balances.sort_by(|a, b| {
2405                let amount_a: u64 = a.amount.amount.parse().unwrap_or(0);
2406                let amount_b: u64 = b.amount.amount.parse().unwrap_or(0);
2407                amount_b.cmp(&amount_a)
2408            });
2409
2410            // limit to top 20 accounts
2411            balances.truncate(20);
2412
2413            balances
2414        })
2415    }
2416
2417    pub async fn get_token_largest_accounts_local_then_remote(
2418        &self,
2419        client: &SurfnetRemoteClient,
2420        mint: &Pubkey,
2421        commitment_config: CommitmentConfig,
2422    ) -> SurfpoolContextualizedResult<Vec<RpcTokenAccountBalance>> {
2423        let SvmAccessContext {
2424            slot,
2425            latest_epoch_info,
2426            latest_blockhash,
2427            inner: local_accounts,
2428        } = self.get_token_largest_accounts_local(mint);
2429
2430        let remote_accounts = client
2431            .get_token_largest_accounts(mint, commitment_config)
2432            .await?;
2433
2434        let mut combined_accounts = remote_accounts;
2435
2436        // if the account is in both the local and remote list, add the local one and not the remote
2437        for local_account in local_accounts {
2438            if let Some((pos, _)) = combined_accounts
2439                .iter()
2440                .find_position(|remote_account| remote_account.address == local_account.address)
2441            {
2442                combined_accounts[pos] = local_account;
2443            } else {
2444                combined_accounts.push(local_account);
2445            }
2446        }
2447
2448        // re-sort and limit after combining
2449        combined_accounts.sort_by(|a, b| {
2450            let amount_a: u64 = a.amount.amount.parse().unwrap_or(0);
2451            let amount_b: u64 = b.amount.amount.parse().unwrap_or(0);
2452            amount_b.cmp(&amount_a)
2453        });
2454        combined_accounts.truncate(20);
2455
2456        Ok(SvmAccessContext::new(
2457            slot,
2458            latest_epoch_info,
2459            latest_blockhash,
2460            combined_accounts,
2461        ))
2462    }
2463
2464    /// Fetches the largest token accounts for a specific mint, returning contextualized results.
2465    pub async fn get_token_largest_accounts(
2466        &self,
2467        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
2468        mint: &Pubkey,
2469    ) -> SurfpoolContextualizedResult<Vec<RpcTokenAccountBalance>> {
2470        if let Some((remote_client, commitment_config)) = remote_ctx {
2471            self.get_token_largest_accounts_local_then_remote(
2472                remote_client,
2473                mint,
2474                *commitment_config,
2475            )
2476            .await
2477        } else {
2478            Ok(self.get_token_largest_accounts_local(mint))
2479        }
2480    }
2481}
2482
2483/// Address lookup table related functions
2484impl SurfnetSvmLocker {
2485    /// Extracts pubkeys from a VersionedMessage, resolving address lookup tables as needed.
2486    pub fn get_pubkeys_from_message(
2487        &self,
2488        message: &VersionedMessage,
2489        all_transaction_lookup_table_addresses: Option<Vec<&Pubkey>>,
2490    ) -> Vec<Pubkey> {
2491        match message {
2492            VersionedMessage::Legacy(message) => message.account_keys.clone(),
2493            VersionedMessage::V0(message) => {
2494                let mut acc_keys = message.account_keys.clone();
2495
2496                if let Some(loaded_addresses) = all_transaction_lookup_table_addresses {
2497                    acc_keys.extend(loaded_addresses);
2498                }
2499                acc_keys
2500            }
2501        }
2502    }
2503
2504    /// Gets addresses loaded from on-chain lookup tables from a VersionedMessage.
2505    pub async fn get_loaded_addresses(
2506        &self,
2507        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
2508        message: &VersionedMessage,
2509    ) -> SurfpoolResult<Option<TransactionLoadedAddresses>> {
2510        match message {
2511            VersionedMessage::Legacy(_) => Ok(None),
2512            VersionedMessage::V0(message) => {
2513                if message.address_table_lookups.is_empty() {
2514                    return Ok(None);
2515                }
2516                let mut loaded = TransactionLoadedAddresses::new();
2517                for alt in message.address_table_lookups.iter() {
2518                    self.get_lookup_table_addresses(remote_ctx, alt, &mut loaded)
2519                        .await?;
2520                }
2521
2522                Ok(Some(loaded))
2523            }
2524        }
2525    }
2526
2527    /// Retrieves loaded addresses from a lookup table account, validating owner and indices.
2528    pub async fn get_lookup_table_addresses(
2529        &self,
2530        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
2531        address_table_lookup: &MessageAddressTableLookup,
2532        transaction_loaded_addresses: &mut TransactionLoadedAddresses,
2533    ) -> SurfpoolResult<()> {
2534        let table_account = self
2535            .get_account(remote_ctx, &address_table_lookup.account_key, None)
2536            .await?
2537            .inner
2538            .map_account()?;
2539
2540        if table_account.owner == solana_sdk_ids::address_lookup_table::id() {
2541            let SvmAccessContext {
2542                slot: current_slot,
2543                inner: slot_hashes,
2544                ..
2545            } = self.with_contextualized_svm_reader(|svm_reader| {
2546                svm_reader
2547                    .inner
2548                    .get_sysvar::<solana_slot_hashes::SlotHashes>()
2549            });
2550
2551            //let current_slot = self.get_latest_absolute_slot(); // or should i use this?
2552            let data = &table_account.data.clone();
2553            let lookup_table = AddressLookupTable::deserialize(data).map_err(|_ix_err| {
2554                SurfpoolError::invalid_account_data(
2555                    address_table_lookup.account_key,
2556                    table_account.data,
2557                    Some("Attempted to lookup addresses from an invalid account"),
2558                )
2559            })?;
2560
2561            let writable = lookup_table
2562                .lookup(
2563                    current_slot,
2564                    &address_table_lookup.writable_indexes,
2565                    &slot_hashes,
2566                )
2567                .map_err(|_ix_err| {
2568                    SurfpoolError::invalid_lookup_index(address_table_lookup.account_key)
2569                })?;
2570
2571            let readable = lookup_table
2572                .lookup(
2573                    current_slot,
2574                    &address_table_lookup.readonly_indexes,
2575                    &slot_hashes,
2576                )
2577                .map_err(|_ix_err| {
2578                    SurfpoolError::invalid_lookup_index(address_table_lookup.account_key)
2579                })?;
2580
2581            let MessageAddressTableLookup {
2582                account_key,
2583                writable_indexes,
2584                readonly_indexes,
2585            } = address_table_lookup.to_owned();
2586
2587            transaction_loaded_addresses.insert_members(
2588                account_key,
2589                writable_indexes
2590                    .into_iter()
2591                    .zip(writable.into_iter())
2592                    .collect(),
2593                readonly_indexes
2594                    .into_iter()
2595                    .zip(readable.into_iter())
2596                    .collect(),
2597            );
2598
2599            Ok(())
2600        } else {
2601            Err(SurfpoolError::invalid_account_owner(
2602                table_account.owner,
2603                Some("Attempted to lookup addresses from an account owned by the wrong program"),
2604            ))
2605        }
2606    }
2607}
2608
2609/// Profiling helper functions
2610impl SurfnetSvmLocker {
2611    /// Estimates compute units for a transaction via contextualized simulation.
2612    pub fn estimate_compute_units(
2613        &self,
2614        transaction: &VersionedTransaction,
2615    ) -> SvmAccessContext<ComputeUnitsEstimationResult> {
2616        self.with_contextualized_svm_reader(|svm_reader| {
2617            svm_reader.estimate_compute_units(transaction)
2618        })
2619    }
2620
2621    /// Creates a partial transaction for instruction profiling by extracting and remapping
2622    /// a subset of instructions from the original transaction.
2623    ///
2624    /// This helper function handles the complex logic of:
2625    /// - Collecting all accounts referenced by the instruction subset
2626    /// - Categorizing accounts based on their original roles (signers vs non-signers)
2627    /// - Building a new account key list in the correct order
2628    /// - Remapping instruction account indices to match the new account list
2629    /// - Creating a valid partial transaction with appropriate signatures
2630    ///
2631    /// # Arguments
2632    /// * `instructions` - All instructions from the original transaction
2633    /// * `message_accounts` - Account keys from the original transaction
2634    /// * `mutable_signers` - Mutable signer accounts from original transaction
2635    /// * `readonly_signers` - Readonly signer accounts from original transaction
2636    /// * `mutable_non_signers` - Mutable non-signer accounts from original transaction
2637    /// * `transaction` - The original transaction for reference
2638    /// * `idx` - Number of instructions to include in the partial transaction
2639    ///
2640    /// # Returns
2641    /// A partial transaction containing the first `idx` instructions and the accounts used for
2642    /// the last instruction, or None if creation fails
2643    #[allow(clippy::too_many_arguments)]
2644    fn create_partial_transaction(
2645        &self,
2646        instructions: &[CompiledInstruction],
2647        message_accounts: &[Pubkey],
2648        transaction: &VersionedTransaction,
2649        idx: usize,
2650        loaded_addresses: &Option<TransactionLoadedAddresses>,
2651    ) -> Option<VersionedTransaction> {
2652        // Keep the full account map from the original transaction for every partial pass.
2653        // This simplifies remapping: we only keep the first `idx` instructions, but retain
2654        // the original `message_accounts` ordering and address table lookups.
2655        let ixs_for_tx = instructions[0..idx].to_vec();
2656
2657        // Build a new message that keeps the original account map and address table lookups,
2658        // but only contains the first `idx` instructions.
2659        let new_message = match transaction.message {
2660            VersionedMessage::Legacy(ref message) => VersionedMessage::Legacy(Message {
2661                account_keys: message_accounts[..message.account_keys.len()].to_vec(),
2662                header: message.header,
2663                recent_blockhash: *transaction.message.recent_blockhash(),
2664                instructions: ixs_for_tx.clone(),
2665            }),
2666            VersionedMessage::V0(ref message) => {
2667                VersionedMessage::V0(solana_message::v0::Message {
2668                    account_keys: message_accounts[..message.account_keys.len()].to_vec(),
2669                    header: message.header,
2670                    recent_blockhash: *transaction.message.recent_blockhash(),
2671                    instructions: ixs_for_tx.clone(),
2672                    // Preserve the original address table lookups when available.
2673                    address_table_lookups: loaded_addresses
2674                        .as_ref()
2675                        .map(|l| l.to_address_table_lookups())
2676                        .unwrap_or_default(),
2677                })
2678            }
2679        };
2680
2681        let tx = VersionedTransaction {
2682            signatures: transaction.signatures.clone(),
2683            message: new_message,
2684        };
2685
2686        Some(tx)
2687    }
2688
2689    /// Returns the profile result for a given signature or UUID, and whether it exists in the SVM.
2690    pub fn get_profile_result(
2691        &self,
2692        signature_or_uuid: UuidOrSignature,
2693        config: &RpcProfileResultConfig,
2694    ) -> SurfpoolResult<Option<UiKeyedProfileResult>> {
2695        let result = match &signature_or_uuid {
2696            UuidOrSignature::Signature(signature) => self.with_svm_reader(|svm| {
2697                svm.executed_transaction_profiles
2698                    .get(&signature.to_string())
2699                    .ok()
2700                    .flatten()
2701            }),
2702            UuidOrSignature::Uuid(uuid) => self.with_svm_reader(|svm| {
2703                svm.simulated_transaction_profiles
2704                    .get(&uuid.to_string())
2705                    .ok()
2706                    .flatten()
2707            }),
2708        };
2709        Ok(result.map(|profile| self.encode_ui_keyed_profile_result(profile, config)))
2710    }
2711
2712    pub fn encode_ui_keyed_profile_result(
2713        &self,
2714        profile: KeyedProfileResult,
2715        config: &RpcProfileResultConfig,
2716    ) -> UiKeyedProfileResult {
2717        self.with_svm_reader(|svm_reader| {
2718            svm_reader.encode_ui_keyed_profile_result(profile, config)
2719        })
2720    }
2721
2722    /// Returns the profile results for a given tag.
2723    pub fn get_profile_results_by_tag(
2724        &self,
2725        tag: String,
2726        config: &RpcProfileResultConfig,
2727    ) -> SurfpoolResult<Option<Vec<UiKeyedProfileResult>>> {
2728        let tag_map = self.with_svm_reader(|svm| svm.profile_tag_map.get(&tag).ok().flatten());
2729        match tag_map {
2730            None => Ok(None),
2731            Some(uuids_or_sigs) => {
2732                let mut profiles = Vec::new();
2733                for id in uuids_or_sigs {
2734                    let profile = self.get_profile_result(id, config)?;
2735                    if profile.is_none() {
2736                        return Err(SurfpoolError::tag_not_found(&tag));
2737                    }
2738                    profiles.push(profile.unwrap());
2739                }
2740                Ok(Some(profiles))
2741            }
2742        }
2743    }
2744
2745    pub fn register_idl(&self, idl: Idl, slot: Option<Slot>) -> SurfpoolResult<()> {
2746        self.with_svm_writer(|svm_writer| svm_writer.register_idl(idl, slot))
2747    }
2748
2749    pub fn get_idl(&self, address: &Pubkey, slot: Option<Slot>) -> Option<Idl> {
2750        self.with_svm_reader(|svm_reader| {
2751            let query_slot = slot.unwrap_or_else(|| svm_reader.get_latest_absolute_slot());
2752            // IDLs are stored sorted by slot descending, so the first one that passes the filter is the latest
2753            svm_reader
2754                .registered_idls
2755                .get(&address.to_string())
2756                .ok()
2757                .flatten()
2758                .and_then(|idl_versions| {
2759                    idl_versions
2760                        .iter()
2761                        .filter(|VersionedIdl(s, _)| *s <= query_slot)
2762                        .max()
2763                        .map(|VersionedIdl(_, idl)| idl.clone())
2764                })
2765        })
2766    }
2767
2768    /// Forges account data by decoding with IDL, applying overrides, and re-encoding.
2769    ///
2770    /// # Arguments
2771    /// * `account_pubkey` - The public key of the account (used for error messages)
2772    /// * `account_data` - The raw account data bytes
2773    /// * `idl` - The IDL for decoding/encoding the account data
2774    /// * `overrides` - HashMap of field paths (dot notation) to values to override
2775    ///
2776    /// # Returns
2777    /// The modified account data bytes with discriminator
2778    /// Forges account data by applying overrides to existing account data
2779    ///
2780    /// This delegates to the SurfnetSvm implementation.
2781    ///
2782    /// # Arguments
2783    /// * `account_pubkey` - The account address (for error messages)
2784    /// * `account_data` - The original account data bytes
2785    /// * `idl` - The IDL for the account's program
2786    /// * `overrides` - Map of field paths to new values
2787    ///
2788    /// # Returns
2789    /// The forged account data as bytes, or an error
2790    pub fn get_forged_account_data(
2791        &self,
2792        account_pubkey: &Pubkey,
2793        account_data: &[u8],
2794        idl: &Idl,
2795        overrides: &HashMap<String, serde_json::Value>,
2796    ) -> SurfpoolResult<Vec<u8>> {
2797        self.with_svm_reader(|svm_reader| {
2798            svm_reader.get_forged_account_data(account_pubkey, account_data, idl, overrides)
2799        })
2800    }
2801}
2802/// Program account related functions
2803impl SurfnetSvmLocker {
2804    /// Clones a program account from source to destination, handling upgradeable loader state.
2805    pub async fn clone_program_account(
2806        &self,
2807        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
2808        source_program_id: &Pubkey,
2809        destination_program_id: &Pubkey,
2810    ) -> SurfpoolContextualizedResult<()> {
2811        let expected_source_program_data_address = get_program_data_address(source_program_id);
2812
2813        let result = self
2814            .get_multiple_accounts(
2815                remote_ctx,
2816                &[*source_program_id, expected_source_program_data_address],
2817                None,
2818            )
2819            .await?;
2820
2821        let mut accounts = result
2822            .inner
2823            .clone()
2824            .into_iter()
2825            .map(|a| a.map_account())
2826            .collect::<SurfpoolResult<Vec<Account>>>()?;
2827
2828        let source_program_data_account = accounts.remove(1);
2829        let source_program_account = accounts.remove(0);
2830
2831        let BpfUpgradeableLoaderAccountType::Program(UiProgram {
2832            program_data: source_program_data_address,
2833        }) = parse_bpf_upgradeable_loader(&source_program_account.data).map_err(|e| {
2834            SurfpoolError::invalid_program_account(source_program_id, e.to_string())
2835        })?
2836        else {
2837            return Err(SurfpoolError::expected_program_account(source_program_id));
2838        };
2839
2840        if source_program_data_address.ne(&expected_source_program_data_address.to_string()) {
2841            return Err(SurfpoolError::invalid_program_account(
2842                source_program_id,
2843                format!(
2844                    "Program data address mismatch: expected {}, found {}",
2845                    expected_source_program_data_address, source_program_data_address
2846                ),
2847            ));
2848        }
2849
2850        let destination_program_data_address = get_program_data_address(destination_program_id);
2851
2852        // create a new program account that has the `program_data` field set to the
2853        // destination program data address
2854        let mut new_program_account = source_program_account;
2855        new_program_account.data = bincode::serialize(&UpgradeableLoaderState::Program {
2856            programdata_address: destination_program_data_address,
2857        })
2858        .map_err(|e| SurfpoolError::internal(format!("Failed to serialize program data: {}", e)))?;
2859
2860        self.with_svm_writer(|svm_writer| {
2861            svm_writer.set_account(
2862                &destination_program_data_address,
2863                source_program_data_account.clone(),
2864            )?;
2865
2866            svm_writer.set_account(destination_program_id, new_program_account.clone())?;
2867            Ok::<(), SurfpoolError>(())
2868        })?;
2869
2870        Ok(result.with_new_value(()))
2871    }
2872
2873    pub async fn set_program_authority(
2874        &self,
2875        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
2876        program_id: Pubkey,
2877        new_authority: Option<Pubkey>,
2878    ) -> SurfpoolContextualizedResult<()> {
2879        let SvmAccessContext {
2880            slot,
2881            latest_epoch_info,
2882            latest_blockhash,
2883            inner: mut get_account_result,
2884        } = self.get_account(remote_ctx, &program_id, None).await?;
2885
2886        let original_authority = match &mut get_account_result {
2887            GetAccountResult::None(pubkey) => {
2888                return Err(SurfpoolError::invalid_program_account(
2889                    pubkey,
2890                    "Account not found",
2891                ));
2892            }
2893            GetAccountResult::FoundAccount(pubkey, program_account, _) => {
2894                let programdata_address = get_program_data_address(pubkey);
2895                let mut programdata_account_result = self
2896                    .get_account(remote_ctx, &programdata_address, None)
2897                    .await?
2898                    .inner;
2899                match &mut programdata_account_result {
2900                    GetAccountResult::None(pubkey) => {
2901                        return Err(SurfpoolError::invalid_program_account(
2902                            pubkey,
2903                            "Program data account does not exist",
2904                        ));
2905                    }
2906                    GetAccountResult::FoundAccount(_, programdata_account, _) => {
2907                        let original_authority = update_programdata_account(
2908                            &program_id,
2909                            programdata_account,
2910                            new_authority,
2911                        )?;
2912
2913                        get_account_result = GetAccountResult::FoundProgramAccount(
2914                            (*pubkey, program_account.clone()),
2915                            (programdata_address, Some(programdata_account.clone())),
2916                        );
2917
2918                        original_authority
2919                    }
2920                    GetAccountResult::FoundProgramAccount(_, _)
2921                    | GetAccountResult::FoundTokenAccount(_, _) => {
2922                        return Err(SurfpoolError::invalid_program_account(
2923                            pubkey,
2924                            "Not a program account",
2925                        ));
2926                    }
2927                }
2928            }
2929            GetAccountResult::FoundProgramAccount(_, (_, None)) => {
2930                return Err(SurfpoolError::invalid_program_account(
2931                    program_id,
2932                    "Program data account does not exist",
2933                ));
2934            }
2935            GetAccountResult::FoundProgramAccount(_, (_, Some(programdata_account))) => {
2936                update_programdata_account(&program_id, programdata_account, new_authority)?
2937            }
2938            GetAccountResult::FoundTokenAccount(_, _) => {
2939                return Err(SurfpoolError::invalid_program_account(
2940                    program_id,
2941                    "Not a program account",
2942                ));
2943            }
2944        };
2945
2946        let simnet_events_tx = self.simnet_events_tx();
2947        match (original_authority, new_authority) {
2948            (Some(original), Some(new)) => {
2949                if original != new {
2950                    let _ = simnet_events_tx.send(SimnetEvent::info(format!(
2951                        "Setting new authority for program {}",
2952                        program_id
2953                    )));
2954                    let _ = simnet_events_tx
2955                        .send(SimnetEvent::info(format!("Old Authority: {}", original)));
2956                    let _ =
2957                        simnet_events_tx.send(SimnetEvent::info(format!("New Authority: {}", new)));
2958                } else {
2959                    let _ = simnet_events_tx.send(SimnetEvent::info(format!(
2960                        "No authority change for program {}",
2961                        program_id
2962                    )));
2963                }
2964            }
2965            (Some(original), None) => {
2966                let _ = simnet_events_tx.send(SimnetEvent::info(format!(
2967                    "Removing authority for program {}",
2968                    program_id
2969                )));
2970                let _ = simnet_events_tx
2971                    .send(SimnetEvent::info(format!("Old Authority: {}", original)));
2972            }
2973            (None, Some(new)) => {
2974                let _ = simnet_events_tx.send(SimnetEvent::info(format!(
2975                    "Setting new authority for program {}",
2976                    program_id
2977                )));
2978                let _ = simnet_events_tx.send(SimnetEvent::info("Old Authority: None".to_string()));
2979                let _ = simnet_events_tx.send(SimnetEvent::info(format!("New Authority: {}", new)));
2980            }
2981            (None, None) => {
2982                let _ = simnet_events_tx.send(SimnetEvent::info(format!(
2983                    "No authority change for program {}",
2984                    program_id
2985                )));
2986            }
2987        };
2988
2989        self.write_account_update(get_account_result);
2990
2991        Ok(SvmAccessContext::new(
2992            slot,
2993            latest_epoch_info,
2994            latest_blockhash,
2995            (),
2996        ))
2997    }
2998
2999    pub async fn get_program_accounts(
3000        &self,
3001        remote_ctx: &Option<SurfnetRemoteClient>,
3002        program_id: &Pubkey,
3003        account_config: RpcAccountInfoConfig,
3004        filters: Option<Vec<RpcFilterType>>,
3005    ) -> SurfpoolContextualizedResult<Vec<RpcKeyedAccount>> {
3006        if let Some(remote_client) = remote_ctx {
3007            self.get_program_accounts_local_then_remote(
3008                remote_client,
3009                program_id,
3010                account_config,
3011                filters,
3012            )
3013            .await
3014        } else {
3015            self.get_program_accounts_local(program_id, account_config, filters)
3016        }
3017    }
3018
3019    /// Retrieves program accounts from the local SVM cache, returning a contextualized result.
3020    pub fn get_program_accounts_local(
3021        &self,
3022        program_id: &Pubkey,
3023        account_config: RpcAccountInfoConfig,
3024        filters: Option<Vec<RpcFilterType>>,
3025    ) -> SurfpoolContextualizedResult<Vec<RpcKeyedAccount>> {
3026        let res = self.with_svm_reader(|svm_reader| {
3027            let res = svm_reader.get_account_owned_by(program_id)?;
3028
3029            let mut filtered = vec![];
3030            for (pubkey, account) in &res {
3031                if let Some(ref active_filters) = filters {
3032                    match apply_rpc_filters(&account.data, active_filters) {
3033                        Ok(true) => {}           // Account matches all filters
3034                        Ok(false) => continue,   // Filtered out
3035                        Err(e) => return Err(e), // Error applying filter, already JsonRpcError
3036                    }
3037                }
3038
3039                filtered.push(svm_reader.account_to_rpc_keyed_account(
3040                    pubkey,
3041                    account,
3042                    &account_config,
3043                    None,
3044                ));
3045            }
3046            Ok(filtered)
3047        })?;
3048
3049        Ok(self.with_contextualized_svm_reader(|_| res.clone()))
3050    }
3051
3052    pub fn encode_ui_account(
3053        &self,
3054        pubkey: &Pubkey,
3055        account: &Account,
3056        encoding: UiAccountEncoding,
3057        additional_data: Option<AccountAdditionalDataV3>,
3058        data_slice: Option<UiDataSliceConfig>,
3059    ) -> UiAccount {
3060        self.with_svm_reader(|svm_reader| {
3061            svm_reader.encode_ui_account(pubkey, account, encoding, additional_data, data_slice)
3062        })
3063    }
3064
3065    /// Retrieves program accounts from the local cache and remote client, combining results.
3066    pub async fn get_program_accounts_local_then_remote(
3067        &self,
3068        client: &SurfnetRemoteClient,
3069        program_id: &Pubkey,
3070        account_config: RpcAccountInfoConfig,
3071        filters: Option<Vec<RpcFilterType>>,
3072    ) -> SurfpoolContextualizedResult<Vec<RpcKeyedAccount>> {
3073        let SvmAccessContext {
3074            slot,
3075            latest_epoch_info,
3076            latest_blockhash,
3077            inner: local_accounts,
3078        } = self.get_program_accounts_local(program_id, account_config.clone(), filters.clone())?;
3079
3080        let remote_accounts_result = client
3081            .get_program_accounts(program_id, account_config, filters)
3082            .await?;
3083
3084        let remote_accounts = remote_accounts_result.handle_method_not_supported(|| {
3085            let tx = self.simnet_events_tx();
3086            let _ = tx.send(SimnetEvent::warn("The `getProgramAccounts` method was sent to the remote RPC, but this method isn't supported by your RPC provider. If you need this method, please use a different RPC provider."));
3087            vec![]
3088        });
3089
3090        let mut combined_accounts = remote_accounts
3091            .into_iter()
3092            .map(|(pubkey, account)| RpcKeyedAccount {
3093                pubkey: pubkey.to_string(),
3094                account,
3095            })
3096            .collect::<Vec<RpcKeyedAccount>>();
3097
3098        for local_account in local_accounts {
3099            // if the local account is in the remote set, replace it with the local one
3100            if let Some((pos, _)) = combined_accounts.iter().find_position(
3101                |RpcKeyedAccount {
3102                     pubkey: remote_pubkey,
3103                     ..
3104                 }| remote_pubkey.eq(&local_account.pubkey),
3105            ) {
3106                combined_accounts[pos] = local_account;
3107            } else {
3108                // otherwise, add the local account to the combined list
3109                combined_accounts.push(local_account);
3110            };
3111        }
3112
3113        Ok(SvmAccessContext {
3114            slot,
3115            latest_epoch_info,
3116            latest_blockhash,
3117            inner: combined_accounts,
3118        })
3119    }
3120}
3121
3122impl SurfnetSvmLocker {
3123    /// Returns the first local slot (the genesis_slot when this surfnet started).
3124    /// Since empty blocks can be reconstructed on-the-fly, all slots from genesis_slot onwards are valid.
3125    pub fn get_first_local_slot(&self) -> Option<Slot> {
3126        self.with_svm_reader(|svm| Some(svm.genesis_slot))
3127    }
3128
3129    pub async fn get_block(
3130        &self,
3131        remote_ctx: &Option<SurfnetRemoteClient>,
3132        slot: &Slot,
3133        config: &RpcBlockConfig,
3134    ) -> SurfpoolContextualizedResult<Option<UiConfirmedBlock>> {
3135        let committed_slot = self.get_slot_for_commitment(&config.commitment.unwrap_or_default());
3136        if *slot > committed_slot {
3137            return Ok(SvmAccessContext {
3138                slot: committed_slot,
3139                latest_epoch_info: self.get_epoch_info(),
3140                latest_blockhash: self
3141                    .get_latest_blockhash(&CommitmentConfig::processed())
3142                    .unwrap_or_default(),
3143                inner: None,
3144            });
3145        }
3146
3147        let first_local_slot = self.get_first_local_slot();
3148
3149        let result = if first_local_slot.is_some() && first_local_slot.unwrap() > *slot {
3150            match remote_ctx {
3151                Some(remote_client) => Some(remote_client.get_block(slot, *config).await?),
3152                None => return Err(SurfpoolError::slot_too_old(*slot)),
3153            }
3154        } else {
3155            self.get_block_local(slot, config)?
3156        };
3157
3158        Ok(SvmAccessContext {
3159            slot: *slot,
3160            latest_epoch_info: self.get_epoch_info(),
3161            latest_blockhash: self
3162                .get_latest_blockhash(&CommitmentConfig::processed())
3163                .unwrap_or_default(),
3164            inner: result,
3165        })
3166    }
3167
3168    pub fn get_block_local(
3169        &self,
3170        slot: &Slot,
3171        config: &RpcBlockConfig,
3172    ) -> SurfpoolResult<Option<UiConfirmedBlock>> {
3173        self.with_svm_reader(|svm_reader| svm_reader.get_block_at_slot(*slot, config))
3174    }
3175
3176    pub fn get_genesis_hash_local(&self) -> SvmAccessContext<Hash> {
3177        self.with_contextualized_svm_reader(|svm_reader| svm_reader.genesis_config.hash())
3178    }
3179
3180    pub async fn get_genesis_hash(
3181        &self,
3182        remote_ctx: &Option<SurfnetRemoteClient>,
3183    ) -> SurfpoolContextualizedResult<Hash> {
3184        if let Some(client) = remote_ctx {
3185            let remote_hash = client.get_genesis_hash().await?;
3186            Ok(self.with_contextualized_svm_reader(|_| remote_hash))
3187        } else {
3188            Ok(self.get_genesis_hash_local())
3189        }
3190    }
3191}
3192
3193/// Pass through functions for accessing the underlying SurfnetSvm instance
3194impl SurfnetSvmLocker {
3195    /// Returns a sender for simulation events from the underlying SVM.
3196    pub fn simnet_events_tx(&self) -> Sender<SimnetEvent> {
3197        self.with_svm_reader(|svm_reader| svm_reader.simnet_events_tx.clone())
3198    }
3199
3200    /// Retrieves the latest epoch info from the underlying SVM.
3201    pub fn get_epoch_info(&self) -> EpochInfo {
3202        self.with_svm_reader(|svm_reader| svm_reader.latest_epoch_info.clone())
3203    }
3204
3205    pub fn time_travel(
3206        &self,
3207        key: Option<(blake3::Hash, String)>,
3208        simnet_command_tx: Sender<SimnetCommand>,
3209        config: TimeTravelConfig,
3210    ) -> SurfpoolResult<EpochInfo> {
3211        let (epoch_info, slot_time, updated_at) = self.with_svm_reader(|svm_reader| {
3212            (
3213                svm_reader.latest_epoch_info.clone(),
3214                svm_reader.slot_time,
3215                svm_reader.updated_at,
3216            )
3217        });
3218
3219        let clock_update: Clock =
3220            calculate_time_travel_clock(&config, updated_at, slot_time, &epoch_info)
3221                .map_err(|e| SurfpoolError::internal(e.to_string()))?;
3222
3223        let formated_time = chrono::DateTime::from_timestamp(clock_update.unix_timestamp, 0)
3224            .unwrap_or_else(|| chrono::DateTime::from_timestamp(0, 0).unwrap())
3225            .format("%Y-%m-%d %H:%M:%S")
3226            .to_string();
3227
3228        // Create a channel for confirmation
3229        let (response_tx, response_rx) = crossbeam_channel::bounded(1);
3230
3231        // Send the command with confirmation
3232        let _ = simnet_command_tx.send(SimnetCommand::UpdateInternalClockWithConfirmation(
3233            key,
3234            clock_update,
3235            response_tx,
3236        ));
3237
3238        // Wait for confirmation with timeout
3239        let updated_epoch_info = response_rx
3240            .recv_timeout(std::time::Duration::from_secs(2))
3241            .map_err(|e| {
3242                SurfpoolError::internal(format!("Failed to confirm clock update: {}", e))
3243            })?;
3244
3245        let _ = self.simnet_events_tx().send(SimnetEvent::info(format!(
3246            "Time travel to {} successful (epoch {} / slot {})",
3247            formated_time, updated_epoch_info.epoch, updated_epoch_info.absolute_slot
3248        )));
3249
3250        Ok(updated_epoch_info)
3251    }
3252
3253    /// Retrieves the latest absolute slot from the underlying SVM.
3254    pub fn get_latest_absolute_slot(&self) -> Slot {
3255        self.with_svm_reader(|svm_reader| svm_reader.get_latest_absolute_slot())
3256    }
3257
3258    /// Retrieves the latest blockhash for the given commitment config from the underlying SVM.
3259    pub fn get_latest_blockhash(&self, config: &CommitmentConfig) -> Option<Hash> {
3260        let slot = self.get_slot_for_commitment(config);
3261        self.with_svm_reader(|svm_reader| svm_reader.blockhash_for_slot(slot))
3262    }
3263
3264    pub fn latest_absolute_blockhash(&self) -> Hash {
3265        self.with_svm_reader(|svm_reader| svm_reader.latest_blockhash())
3266    }
3267
3268    pub fn get_slot_for_commitment(&self, commitment: &CommitmentConfig) -> Slot {
3269        self.with_svm_reader(|svm_reader| {
3270            let slot = svm_reader.get_latest_absolute_slot();
3271            match commitment.commitment {
3272                CommitmentLevel::Processed => slot,
3273                CommitmentLevel::Confirmed => slot.saturating_sub(1),
3274                CommitmentLevel::Finalized => slot.saturating_sub(FINALIZATION_SLOT_THRESHOLD),
3275            }
3276        })
3277    }
3278
3279    /// Executes an airdrop via the underlying SVM.
3280    #[allow(clippy::result_large_err)]
3281    pub fn airdrop(&self, pubkey: &Pubkey, lamports: u64) -> SurfpoolResult<TransactionResult> {
3282        self.with_svm_writer(|svm_writer| svm_writer.airdrop(pubkey, lamports))
3283    }
3284
3285    /// Executes a batch airdrop via the underlying SVM.
3286    pub fn airdrop_pubkeys(&self, lamports: u64, addresses: &[Pubkey]) {
3287        self.with_svm_writer(|svm_writer| svm_writer.airdrop_pubkeys(lamports, addresses))
3288    }
3289
3290    /// Confirms the current block on the underlying SVM, returning `Ok(())` or an error.
3291    pub async fn confirm_current_block(
3292        &self,
3293        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
3294    ) -> SurfpoolResult<()> {
3295        // Acquire write lock once and do both operations atomically
3296        // This prevents lock contention and potential deadlocks from mixing blocking and async locks
3297        let mut svm_writer = self.0.write().await;
3298        svm_writer.confirm_current_block()?;
3299        svm_writer.materialize_overrides(remote_ctx).await
3300    }
3301
3302    /// Subscribes for signature updates (confirmed/finalized) and returns a receiver of events.
3303    pub fn subscribe_for_signature_updates(
3304        &self,
3305        signature: &Signature,
3306        subscription_type: SignatureSubscriptionType,
3307    ) -> Receiver<(Slot, Option<TransactionError>)> {
3308        self.with_svm_writer(|svm_writer| {
3309            svm_writer.subscribe_for_signature_updates(signature, subscription_type.clone())
3310        })
3311    }
3312
3313    /// Subscribes for account updates and returns a receiver of account updates.
3314    pub fn subscribe_for_account_updates(
3315        &self,
3316        account_pubkey: &Pubkey,
3317        encoding: Option<UiAccountEncoding>,
3318    ) -> Receiver<UiAccount> {
3319        // Handles the locking/unlocking safely
3320        self.with_svm_writer(|svm_writer| {
3321            svm_writer.subscribe_for_account_updates(account_pubkey, encoding)
3322        })
3323    }
3324
3325    /// Subscribes for program account updates and returns a receiver of keyed account updates.
3326    pub fn subscribe_for_program_updates(
3327        &self,
3328        program_id: &Pubkey,
3329        encoding: Option<UiAccountEncoding>,
3330        filters: Option<Vec<RpcFilterType>>,
3331    ) -> Receiver<RpcKeyedAccount> {
3332        self.with_svm_writer(|svm_writer| {
3333            svm_writer.subscribe_for_program_updates(program_id, encoding, filters)
3334        })
3335    }
3336
3337    /// Subscribes for slot updates and returns a receiver of slot updates.
3338    pub fn subscribe_for_slot_updates(&self) -> Receiver<SlotInfo> {
3339        self.with_svm_writer(|svm_writer| svm_writer.subscribe_for_slot_updates())
3340    }
3341
3342    /// Subscribes for logs updates and returns a receiver of logs updates.
3343    pub fn subscribe_for_logs_updates(
3344        &self,
3345        commitment_level: &CommitmentLevel,
3346        filter: &RpcTransactionLogsFilter,
3347    ) -> Receiver<(Slot, RpcLogsResponse)> {
3348        self.with_svm_writer(|svm_writer| {
3349            svm_writer.subscribe_for_logs_updates(commitment_level, filter)
3350        })
3351    }
3352
3353    /// Subscribes for snapshot import updates and returns a receiver of snapshot import notifications.
3354    /// This method spawns a background task that fetches the snapshot and loads it via `load_snapshot`.
3355    pub fn subscribe_for_snapshot_import_updates(
3356        &self,
3357        snapshot_url: &str,
3358        snapshot_id: &str,
3359    ) -> Receiver<super::SnapshotImportNotification> {
3360        // Register the subscription and get the sender/receiver
3361        let (tx, rx) =
3362            self.with_svm_writer(|svm_writer| svm_writer.register_snapshot_subscription());
3363
3364        // Clone the locker for use in the spawned task
3365        let locker = self.clone();
3366        let snapshot_url = snapshot_url.to_string();
3367        let snapshot_id = snapshot_id.to_string();
3368
3369        tokio::spawn(async move {
3370            // Send initial notification
3371            let _ = tx.send(super::SnapshotImportNotification {
3372                snapshot_id: snapshot_id.clone(),
3373                status: super::SnapshotImportStatus::Started,
3374                accounts_loaded: 0,
3375                total_accounts: 0,
3376                error: None,
3377            });
3378
3379            // Fetch snapshot from URL and parse it
3380            let snapshot_data = match SurfnetSvm::fetch_snapshot_from_url(&snapshot_url).await {
3381                Ok(data) => data,
3382                Err(e) => {
3383                    let _ = tx.send(super::SnapshotImportNotification {
3384                        snapshot_id,
3385                        status: super::SnapshotImportStatus::Failed,
3386                        accounts_loaded: 0,
3387                        total_accounts: 0,
3388                        error: Some(format!("Failed to fetch snapshot: {}", e)),
3389                    });
3390                    return;
3391                }
3392            };
3393
3394            let total_accounts = snapshot_data.len() as u64;
3395
3396            // Send progress notification with total count
3397            let _ = tx.send(super::SnapshotImportNotification {
3398                snapshot_id: snapshot_id.clone(),
3399                status: super::SnapshotImportStatus::InProgress,
3400                accounts_loaded: 0,
3401                total_accounts,
3402                error: None,
3403            });
3404
3405            // Load the snapshot using the load_snapshot method
3406            match locker
3407                .load_snapshot(&snapshot_data, None, CommitmentConfig::processed())
3408                .await
3409            {
3410                Ok(loaded_count) => {
3411                    let _ = tx.send(super::SnapshotImportNotification {
3412                        snapshot_id,
3413                        status: super::SnapshotImportStatus::Completed,
3414                        accounts_loaded: loaded_count as u64,
3415                        total_accounts,
3416                        error: None,
3417                    });
3418                }
3419                Err(e) => {
3420                    let _ = tx.send(super::SnapshotImportNotification {
3421                        snapshot_id,
3422                        status: super::SnapshotImportStatus::Failed,
3423                        accounts_loaded: 0,
3424                        total_accounts,
3425                        error: Some(format!("Failed to load snapshot: {}", e)),
3426                    });
3427                }
3428            }
3429        });
3430
3431        rx
3432    }
3433
3434    pub fn runbook_executions(&self) -> Vec<RunbookExecutionStatusReport> {
3435        self.with_svm_reader(|svm_reader| svm_reader.runbook_executions.clone())
3436    }
3437
3438    pub fn start_runbook_execution(&self, runbook_id: String) {
3439        self.with_svm_writer(|svm_writer| {
3440            svm_writer.instruction_profiling_enabled = false;
3441            svm_writer.start_runbook_execution(runbook_id);
3442        });
3443    }
3444
3445    pub fn complete_runbook_execution(
3446        &self,
3447        runbook_id: String,
3448        error: Option<Vec<String>>,
3449        re_enable_ix_profiling: bool,
3450    ) {
3451        self.with_svm_writer(|svm_writer| {
3452            svm_writer.complete_runbook_execution(&runbook_id, error);
3453            let some_runbook_executing = svm_writer
3454                .runbook_executions
3455                .iter()
3456                .any(|e| e.completed_at.is_none());
3457            if !some_runbook_executing && re_enable_ix_profiling {
3458                svm_writer.instruction_profiling_enabled = true;
3459            }
3460        });
3461    }
3462
3463    pub fn export_snapshot(
3464        &self,
3465        config: ExportSnapshotConfig,
3466    ) -> SurfpoolResult<BTreeMap<String, AccountSnapshot>> {
3467        self.with_svm_reader(|svm_reader| svm_reader.export_snapshot(config))
3468    }
3469
3470    pub fn get_start_time(&self) -> SystemTime {
3471        self.with_svm_reader(|svm_reader| svm_reader.start_time)
3472    }
3473}
3474
3475/// Helpers for writing program accounts
3476impl SurfnetSvmLocker {
3477    pub async fn write_program(
3478        &self,
3479        program_id: Pubkey,
3480        authority: Option<Pubkey>,
3481        offset: usize,
3482        data: &[u8],
3483        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
3484    ) -> SurfpoolResult<()> {
3485        let program_data_address = get_program_data_address(&program_id);
3486
3487        let program_account = self
3488            .get_or_create_program_account(program_id, program_data_address, remote_ctx)
3489            .await?;
3490
3491        let _ = self
3492            .write_program_data_account_with_offset(
3493                program_id,
3494                authority,
3495                program_data_address,
3496                offset,
3497                data,
3498                remote_ctx,
3499            )
3500            .await?;
3501
3502        // Re-set the program account to force LiteSVM to recompile the program
3503        // from the updated programdata. Without this, the program cache retains
3504        // the noop placeholder compiled during initial program account creation.
3505        // Errors are expected for incomplete ELF (multi-chunk writes) and are
3506        // logged but not propagated.
3507        let set_result = self.with_svm_writer(|svm_writer| {
3508            svm_writer.set_account(&program_id, program_account.clone())
3509        });
3510        if let Err(e) = set_result {
3511            let _ = self.simnet_events_tx().send(SimnetEvent::info(format!(
3512                "Program cache update deferred for {}: {}",
3513                program_id, e
3514            )));
3515        }
3516
3517        Ok(())
3518    }
3519
3520    pub async fn get_or_create_program_account(
3521        &self,
3522        program_id: Pubkey,
3523        program_data_address: Pubkey,
3524        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
3525    ) -> SurfpoolResult<Account> {
3526        // Get program account
3527        let SvmAccessContext {
3528            inner: program_account_result,
3529            ..
3530        } = self
3531            .get_account(
3532                remote_ctx,
3533                &program_id,
3534                Some(Box::new(move |svm_locker| {
3535                    // Create default program account if it doesn't exist
3536                    let program_state =
3537                        solana_loader_v3_interface::state::UpgradeableLoaderState::Program {
3538                            programdata_address: program_data_address,
3539                        };
3540                    let program_data = bincode::serialize(&program_state)
3541                        .expect("Failed to serialize program state");
3542                    let program_lamports = svm_locker.with_svm_reader(|svm_reader| {
3543                        svm_reader
3544                            .inner
3545                            .minimum_balance_for_rent_exemption(program_data.len())
3546                    });
3547
3548                    let _ = svm_locker
3549                        .simnet_events_tx()
3550                        .send(SimnetEvent::info(format!(
3551                            "Creating program account {} with program data address {}",
3552                            program_id, program_data_address
3553                        )));
3554
3555                    GetAccountResult::FoundAccount(
3556                        program_id,
3557                        solana_account::Account {
3558                            lamports: program_lamports,
3559                            data: program_data,
3560                            owner: solana_sdk_ids::bpf_loader_upgradeable::id(),
3561                            executable: true,
3562                            rent_epoch: 0,
3563                        },
3564                        true,
3565                    )
3566                })),
3567            )
3568            .await?;
3569
3570        // Check if account was created before consuming it
3571        let was_program_created = matches!(
3572            program_account_result,
3573            GetAccountResult::FoundAccount(_, _, true)
3574        );
3575
3576        // Ensure we have a valid program account
3577        let program_account = program_account_result.map_account()?;
3578
3579        // Validate it's owned by the upgradeable loader
3580        if program_account.owner != solana_sdk_ids::bpf_loader_upgradeable::id() {
3581            return Err(SurfpoolError::invalid_program_account(
3582                &program_id,
3583                "Account not owned by the BPF Upgradeable Loader",
3584            ));
3585        }
3586
3587        // Validate it's an executable program account
3588        if !program_account.executable {
3589            return Err(SurfpoolError::invalid_program_account(
3590                &program_id,
3591                "Account not executable",
3592            ));
3593        }
3594
3595        // Persist the program account if it was newly created
3596        if was_program_created {
3597            self.write_account_update(GetAccountResult::FoundAccount(
3598                program_id,
3599                program_account.clone(),
3600                true,
3601            ));
3602        }
3603        Ok(program_account)
3604    }
3605
3606    pub async fn write_program_data_account_with_offset(
3607        &self,
3608        program_id: Pubkey,
3609        authority: Option<Pubkey>,
3610        program_data_address: Pubkey,
3611        offset: usize,
3612        data: &[u8],
3613        remote_ctx: &Option<(SurfnetRemoteClient, CommitmentConfig)>,
3614    ) -> SurfpoolResult<Account> {
3615        // Get or create program data account
3616        let SvmAccessContext {
3617            inner: program_data_result,
3618            slot,
3619            ..
3620        } = self
3621            .get_account(
3622                &remote_ctx,
3623                &program_data_address,
3624                Some(Box::new(move |svm_locker| {
3625                    // Create default program data account if it doesn't exist
3626                    let programdata_state =
3627                        solana_loader_v3_interface::state::UpgradeableLoaderState::ProgramData {
3628                            slot: svm_locker.get_latest_absolute_slot(),
3629                            // TODO: currently litesvm breaks if you don't provide an authority,
3630                            // but once that's fixed we should remove the default to system program
3631                            upgrade_authority_address: authority
3632                                .or(Some(solana_system_interface::program::id())),
3633                        };
3634                    let mut programdata_data = bincode::serialize(&programdata_state)
3635                        .expect("Failed to serialize program data state");
3636
3637                    // Add empty program data (will be filled by writes)
3638                    programdata_data.extend(vec![0u8; 0]);
3639
3640                    let programdata_lamports = svm_locker.with_svm_reader(|svm_reader| {
3641                        svm_reader
3642                            .inner
3643                            .minimum_balance_for_rent_exemption(programdata_data.len())
3644                    });
3645
3646                    let _ = svm_locker
3647                        .simnet_events_tx()
3648                        .send(SimnetEvent::info(format!(
3649                            "Creating program data account {} for program {}",
3650                            program_data_address, program_id
3651                        )));
3652
3653                    GetAccountResult::FoundAccount(
3654                        program_data_address,
3655                        solana_account::Account {
3656                            lamports: programdata_lamports,
3657                            data: programdata_data,
3658                            owner: solana_sdk_ids::bpf_loader_upgradeable::id(),
3659                            executable: false,
3660                            rent_epoch: 0,
3661                        },
3662                        true,
3663                    )
3664                })),
3665            )
3666            .await?;
3667
3668        // Get mutable program data account
3669        let mut program_data_account = program_data_result.map_account()?;
3670
3671        // Calculate metadata size
3672        let metadata_size =
3673            solana_loader_v3_interface::state::UpgradeableLoaderState::size_of_programdata_metadata(
3674            );
3675
3676        // Verify program data account has valid state
3677        let upgrade_authority_address = match bincode::deserialize::<
3678            solana_loader_v3_interface::state::UpgradeableLoaderState,
3679        >(&program_data_account.data[..metadata_size])
3680        {
3681            Ok(solana_loader_v3_interface::state::UpgradeableLoaderState::ProgramData {
3682                upgrade_authority_address,
3683                ..
3684            }) => upgrade_authority_address,
3685            Ok(_) => {
3686                return Err(SurfpoolError::invalid_program_data_account(
3687                    program_data_address,
3688                    "Account is not a program data account",
3689                ));
3690            }
3691            Err(e) => {
3692                return Err(SurfpoolError::invalid_program_data_account(
3693                    program_data_address,
3694                    format!("Invalid program data account state: {}", e),
3695                ));
3696            }
3697        };
3698
3699        let new_metadata = if upgrade_authority_address.ne(&authority) {
3700            let _ = self.simnet_events_tx().send(SimnetEvent::info(format!(
3701                "Updating program authority of program {} to {}",
3702                program_id,
3703                authority.unwrap_or(solana_system_interface::program::id())
3704            )));
3705            solana_loader_v3_interface::state::UpgradeableLoaderState::ProgramData {
3706                slot,
3707                upgrade_authority_address: authority
3708                    .or(Some(solana_system_interface::program::id())),
3709            }
3710        } else {
3711            solana_loader_v3_interface::state::UpgradeableLoaderState::ProgramData {
3712                slot,
3713                upgrade_authority_address,
3714            }
3715        };
3716
3717        let metadata_bytes = bincode::serialize(&new_metadata).map_err(|e| {
3718            SurfpoolError::internal(format!("Failed to serialize program data metadata: {}", e))
3719        })?;
3720
3721        // Strip the minimum_program.so placeholder if it was pre-filled by
3722        // init_programdata_account during program account creation. This prevents
3723        // leftover placeholder bytes when the actual program is smaller than 3312 bytes.
3724        let minimum_program_bytes = crate::surfnet::noop_program::NOOP_PROGRAM_ELF;
3725        if program_data_account.data.len() == metadata_size + minimum_program_bytes.len()
3726            && program_data_account.data[metadata_size..] == *minimum_program_bytes
3727        {
3728            program_data_account.data.truncate(metadata_size);
3729        }
3730
3731        // Calculate absolute offset in account data (metadata + offset)
3732        let absolute_offset = metadata_size + offset;
3733        let end_offset = absolute_offset + data.len();
3734
3735        // Expand account data if necessary
3736        if end_offset > program_data_account.data.len() {
3737            let new_size = end_offset;
3738            program_data_account.data.resize(new_size, 0);
3739
3740            // Update lamports for rent exemption
3741            let new_lamports = self.with_svm_reader(|svm_reader| {
3742                svm_reader
3743                    .inner
3744                    .minimum_balance_for_rent_exemption(new_size)
3745            });
3746            program_data_account.lamports = new_lamports;
3747
3748            let _ = self.simnet_events_tx().send(SimnetEvent::info(format!(
3749                "Expanding program data account to {} bytes",
3750                new_size
3751            )));
3752        }
3753
3754        // Write the metadata
3755        program_data_account.data[..metadata_size].copy_from_slice(&metadata_bytes);
3756        // Write data at the specified offset
3757        program_data_account.data[absolute_offset..end_offset].copy_from_slice(&data);
3758
3759        // Update the account in SVM
3760        self.with_svm_writer(|svm_writer| {
3761            svm_writer.set_account(&program_data_address, program_data_account.clone())?;
3762            Ok::<(), SurfpoolError>(())
3763        })?;
3764
3765        let _ = self.simnet_events_tx().send(SimnetEvent::info(format!(
3766            "Wrote {} bytes to program {} at offset {}",
3767            data.len(),
3768            program_id,
3769            offset
3770        )));
3771
3772        Ok(program_data_account)
3773    }
3774}
3775
3776// Helper function to apply filters
3777pub(crate) fn apply_rpc_filters(
3778    account_data: &[u8],
3779    filters: &[RpcFilterType],
3780) -> SurfpoolResult<bool> {
3781    for filter in filters {
3782        match filter {
3783            RpcFilterType::DataSize(size) => {
3784                if account_data.len() as u64 != *size {
3785                    return Ok(false);
3786                }
3787            }
3788            RpcFilterType::Memcmp(memcmp_filter) => {
3789                // Use the public bytes_match method from solana_client::rpc_filter::Memcmp
3790                if !memcmp_filter.bytes_match(account_data) {
3791                    return Ok(false); // Content mismatch or out of bounds handled by bytes_match
3792                }
3793            }
3794            RpcFilterType::TokenAccountState => {
3795                return Err(SurfpoolError::internal(
3796                    "TokenAccountState filter is not supported",
3797                ));
3798            }
3799        }
3800    }
3801    Ok(true)
3802}
3803
3804// used in the remote.rs
3805pub fn is_supported_token_program(program_id: &Pubkey) -> bool {
3806    *program_id == spl_token_interface::ID || *program_id == spl_token_2022_interface::ID
3807}
3808
3809fn update_programdata_account(
3810    program_id: &Pubkey,
3811    programdata_account: &mut Account,
3812    new_authority: Option<Pubkey>,
3813) -> SurfpoolResult<Option<Pubkey>> {
3814    let upgradeable_loader_state =
3815        bincode::deserialize::<UpgradeableLoaderState>(&programdata_account.data).map_err(|e| {
3816            SurfpoolError::invalid_program_account(
3817                program_id,
3818                format!("Failed to serialize program data: {}", e),
3819            )
3820        })?;
3821    if let UpgradeableLoaderState::ProgramData {
3822        upgrade_authority_address,
3823        slot,
3824    } = upgradeable_loader_state
3825    {
3826        let offset = if upgrade_authority_address.is_some() {
3827            UpgradeableLoaderState::size_of_programdata_metadata()
3828        } else {
3829            UpgradeableLoaderState::size_of_programdata_metadata()
3830                - serialized_size(&Pubkey::default()).unwrap() as usize
3831        };
3832
3833        let mut data = bincode::serialize(&UpgradeableLoaderState::ProgramData {
3834            upgrade_authority_address: new_authority,
3835            slot,
3836        })
3837        .map_err(|e| {
3838            SurfpoolError::invalid_program_account(
3839                program_id,
3840                format!("Failed to serialize program data: {}", e),
3841            )
3842        })?;
3843
3844        data.append(&mut programdata_account.data[offset..].to_vec());
3845
3846        programdata_account.data = data;
3847
3848        Ok(upgrade_authority_address)
3849    } else {
3850        Err(SurfpoolError::invalid_program_account(
3851            program_id,
3852            "Invalid program data account",
3853        ))
3854    }
3855}
3856
3857pub fn format_ui_amount_string(amount: u64, decimals: u8) -> String {
3858    if decimals > 0 {
3859        let divisor = 10u64.pow(decimals as u32);
3860        format!(
3861            "{:.decimals$}",
3862            amount as f64 / divisor as f64,
3863            decimals = decimals as usize
3864        )
3865    } else {
3866        amount.to_string()
3867    }
3868}
3869
3870pub fn format_ui_amount(amount: u64, decimals: u8) -> f64 {
3871    if decimals > 0 {
3872        let divisor = 10u64.pow(decimals as u32);
3873        amount as f64 / divisor as f64
3874    } else {
3875        amount as f64
3876    }
3877}
3878
3879#[cfg(test)]
3880mod tests {
3881    use std::collections::HashMap;
3882
3883    use solana_account::Account;
3884    use solana_account_decoder::UiAccountEncoding;
3885    use solana_epoch_schedule::EpochSchedule;
3886    use solana_transaction_status::TransactionStatusMeta;
3887
3888    use super::*;
3889    use crate::{
3890        scenarios::registry::PYTH_V2_IDL_CONTENT,
3891        surfnet::{BlockHeader, SurfnetSvm, svm::apply_override_to_decoded_account},
3892    };
3893
3894    #[test]
3895    fn test_get_forged_account_data_with_pyth_fixture() {
3896        use borsh::{BorshDeserialize, BorshSerialize};
3897
3898        // Define local structures matching Pyth IDL
3899        #[derive(BorshSerialize, BorshDeserialize, Debug, Clone, PartialEq)]
3900        pub enum VerificationLevel {
3901            Partial { num_signatures: u8 },
3902            Full,
3903        }
3904
3905        #[derive(BorshSerialize, BorshDeserialize, Debug, Clone, PartialEq)]
3906        pub struct PriceFeedMessage {
3907            pub feed_id: [u8; 32],
3908            pub price: i64,
3909            pub conf: u64,
3910            pub exponent: i32,
3911            pub publish_time: i64,
3912            pub prev_publish_time: i64,
3913            pub ema_price: i64,
3914            pub ema_conf: u64,
3915        }
3916
3917        #[derive(BorshSerialize, BorshDeserialize, Debug, Clone, PartialEq)]
3918        pub struct PriceUpdateV2 {
3919            pub write_authority: Pubkey,
3920            pub verification_level: VerificationLevel,
3921            pub price_message: PriceFeedMessage,
3922            pub posted_slot: u64,
3923        }
3924
3925        // Pyth price feed account data fixture
3926        let account_data_hex = vec![
3927            0x22, 0xf1, 0x23, 0x63, 0x9d, 0x7e, 0xf4, 0xcd, // Discriminator
3928            0x35, 0xa7, 0x0c, 0x11, 0x16, 0x2f, 0xbf, 0x5a, 0x0e, 0x7f, 0x7d, 0x2f, 0x96, 0xe1,
3929            0x9f, 0x97, 0xb0, 0x22, 0x46, 0xa1, 0x56, 0x87, 0xee, 0x67, 0x27, 0x94, 0x89, 0x74,
3930            0x48, 0xe6, 0x58, 0xde, 0x01, 0xe6, 0x2d, 0xf6, 0xc8, 0xb4, 0xa8, 0x5f, 0xe1, 0xa6,
3931            0x7d, 0xb4, 0x4d, 0xc1, 0x2d, 0xe5, 0xdb, 0x33, 0x0f, 0x7a, 0xc6, 0x6b, 0x72, 0xdc,
3932            0x65, 0x8a, 0xfe, 0xdf, 0x0f, 0x4a, 0x41, 0x5b, 0x43, 0xd7, 0x1f, 0x18, 0x64, 0x5f,
3933            0x0a, 0x00, 0x00, 0x96, 0x67, 0xea, 0xc5, 0x00, 0x00, 0x00, 0x00, 0xf8, 0xff, 0xff,
3934            0xff, 0x5f, 0x2b, 0x00, 0x69, 0x00, 0x00, 0x00, 0x00, 0x5e, 0x2b, 0x00, 0x69, 0x00,
3935            0x00, 0x00, 0x00, 0xa0, 0x7c, 0x1a, 0x38, 0x63, 0x0a, 0x00, 0x00, 0x94, 0xa6, 0xb9,
3936            0xb5, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x5e, 0x6d, 0x16, 0x00, 0x00, 0x00, 0x00, 0x00,
3937        ];
3938
3939        // Create a minimal Pyth IDL for testing
3940        let idl: Idl = serde_json::from_str(PYTH_V2_IDL_CONTENT).expect("Failed to load IDL");
3941
3942        // Create overrides - note: this won't actually work with the JSON deserialization
3943        // since the account data is Borsh-encoded, but we're testing the structure
3944        let mut overrides: HashMap<String, serde_json::Value> = HashMap::new();
3945
3946        // Verify IDL has matching discriminator
3947        let account_def = idl
3948            .accounts
3949            .iter()
3950            .find(|acc| acc.discriminator.eq(&account_data_hex[..8]));
3951
3952        assert!(
3953            account_def.is_some(),
3954            "Should find PriceUpdateV2 account by discriminator"
3955        );
3956        assert_eq!(account_def.unwrap().name, "PriceUpdateV2");
3957
3958        // Step 1: Instantiate an offline Svm instance
3959        let (surfnet_svm, _simnet_events_rx, _geyser_events_rx) = SurfnetSvm::default();
3960        let svm_locker = SurfnetSvmLocker::new(surfnet_svm);
3961
3962        // Step 2: Register the IDL for this account
3963        let account_pubkey = Pubkey::from_str_const("rec5EKMGg6MxZYaMdyBfgwp4d5rB9T1VQH5pJv5LtFJ");
3964        svm_locker.register_idl(idl.clone(), None).unwrap();
3965
3966        // Step 3: Create an account with the Pyth data
3967        let pyth_account = Account {
3968            lamports: 1_000_000,
3969            data: account_data_hex.clone(),
3970            owner: account_pubkey,
3971            executable: false,
3972            rent_epoch: 0,
3973        };
3974
3975        // Step 4: Use encode_ui_account to decode/encode the account data
3976        let ui_account = svm_locker.encode_ui_account(
3977            &account_pubkey,
3978            &pyth_account,
3979            UiAccountEncoding::JsonParsed,
3980            None,
3981            None, // data_slice
3982        );
3983
3984        // Step 5: Verify the UI account has parsed data
3985        println!("UI Account lamports: {}", ui_account.lamports);
3986        println!("UI Account owner: {}", ui_account.owner);
3987
3988        // Assert on parsed account data
3989        use solana_account_decoder::UiAccountData;
3990        match &ui_account.data {
3991            UiAccountData::Json(parsed_account) => {
3992                let parsed_obj = &parsed_account.parsed;
3993
3994                // Extract price_message object
3995                let price_message = parsed_obj
3996                    .get("price_message")
3997                    .expect("Should have price_message field")
3998                    .as_object()
3999                    .expect("price_message should be an object");
4000
4001                // Assert on price
4002                let price = price_message
4003                    .get("price")
4004                    .expect("Should have price field")
4005                    .as_i64()
4006                    .expect("price should be a number");
4007                assert_eq!(price, 11404817473495, "Price should match expected value");
4008
4009                // Assert on exponent
4010                let exponent = price_message
4011                    .get("exponent")
4012                    .expect("Should have exponent field")
4013                    .as_i64()
4014                    .expect("exponent should be a number");
4015                assert_eq!(exponent, -8, "Exponent should be -8");
4016
4017                // Assert on ema_price
4018                let ema_price = price_message
4019                    .get("ema_price")
4020                    .expect("Should have ema_price field")
4021                    .as_i64()
4022                    .expect("ema_price should be a number");
4023                assert_eq!(
4024                    ema_price, 11421259300000,
4025                    "EMA price should match expected value"
4026                );
4027
4028                // Assert on publish_time
4029                let publish_time = price_message
4030                    .get("publish_time")
4031                    .expect("Should have publish_time field")
4032                    .as_i64()
4033                    .expect("publish_time should be a number");
4034                assert_eq!(
4035                    publish_time, 1761618783,
4036                    "Publish time should match expected value"
4037                );
4038
4039                println!("✓ All price assertions passed!");
4040            }
4041            _ => panic!("Expected JSON parsed account data"),
4042        }
4043
4044        // Step 6: Test get_forged_account_data without overrides (should return same data)
4045        println!("\n--- Testing get_forged_account_data without overrides ---");
4046        let forged_data_no_overrides = svm_locker.get_forged_account_data(
4047            &account_pubkey,
4048            &account_data_hex,
4049            &idl,
4050            &overrides,
4051        );
4052
4053        match forged_data_no_overrides {
4054            Ok(data) => {
4055                // If it succeeds, verify the data is unchanged
4056                assert_eq!(
4057                    data, account_data_hex,
4058                    "Data without overrides should match original"
4059                );
4060                println!("✓ Forged data without overrides matches original!");
4061            }
4062            Err(e) => {
4063                // If it fails, it's due to Borsh/JSON mismatch (expected for now)
4064                println!("Expected error (Borsh vs JSON): {:?}", e);
4065                println!("Note: This documents the need for proper Borsh implementation");
4066            }
4067        }
4068
4069        // Step 7: Test get_forged_account_data with overrides
4070        println!("\n--- Testing get_forged_account_data with overrides ---");
4071
4072        // Set new values for price and publish_time
4073        let new_price = 999999999999i64;
4074        let new_publish_time = 1234567890i64;
4075        let new_ema_price = 888888888888i64;
4076
4077        overrides.insert("price_message.price".into(), json!(new_price));
4078        overrides.insert("price_message.publish_time".into(), json!(new_publish_time));
4079        overrides.insert("price_message.ema_price".into(), json!(new_ema_price));
4080
4081        let forged_data_with_overrides = svm_locker.get_forged_account_data(
4082            &account_pubkey,
4083            &account_data_hex,
4084            &idl,
4085            &overrides,
4086        );
4087
4088        match forged_data_with_overrides {
4089            Ok(modified_data) => {
4090                // Verify the data is different from original
4091                assert_ne!(
4092                    modified_data, account_data_hex,
4093                    "Modified data should be different from original"
4094                );
4095                println!("✓ Modified data is different from original!");
4096
4097                // Create a modified account to verify the changes
4098                let modified_account = Account {
4099                    lamports: 1_000_000,
4100                    data: modified_data.clone(),
4101                    owner: account_pubkey,
4102                    executable: false,
4103                    rent_epoch: 0,
4104                };
4105
4106                // Re-encode the modified account to verify the changes
4107                let modified_ui_account = svm_locker.encode_ui_account(
4108                    &account_pubkey,
4109                    &modified_account,
4110                    UiAccountEncoding::JsonParsed,
4111                    None,
4112                    None,
4113                );
4114
4115                // Verify the modified values in the re-encoded account
4116                match &modified_ui_account.data {
4117                    UiAccountData::Json(parsed_account) => {
4118                        let parsed_obj = &parsed_account.parsed;
4119                        let price_message = parsed_obj
4120                            .get("price_message")
4121                            .expect("Should have price_message field")
4122                            .as_object()
4123                            .expect("price_message should be an object");
4124
4125                        // Verify new price
4126                        let modified_price = price_message
4127                            .get("price")
4128                            .expect("Should have price field")
4129                            .as_i64()
4130                            .expect("price should be a number");
4131                        assert_eq!(
4132                            modified_price, new_price,
4133                            "Modified price should match override value"
4134                        );
4135
4136                        // Verify new publish_time
4137                        let modified_publish_time = price_message
4138                            .get("publish_time")
4139                            .expect("Should have publish_time field")
4140                            .as_i64()
4141                            .expect("publish_time should be a number");
4142                        assert_eq!(
4143                            modified_publish_time, new_publish_time,
4144                            "Modified publish_time should match override value"
4145                        );
4146
4147                        // Verify new ema_price
4148                        let modified_ema_price = price_message
4149                            .get("ema_price")
4150                            .expect("Should have ema_price field")
4151                            .as_i64()
4152                            .expect("ema_price should be a number");
4153                        assert_eq!(
4154                            modified_ema_price, new_ema_price,
4155                            "Modified ema_price should match override value"
4156                        );
4157
4158                        // Verify exponent is unchanged
4159                        let exponent = price_message
4160                            .get("exponent")
4161                            .expect("Should have exponent field")
4162                            .as_i64()
4163                            .expect("exponent should be a number");
4164                        assert_eq!(exponent, -8, "Exponent should remain unchanged");
4165
4166                        println!("✓ All override assertions passed!");
4167                        println!("  - Price changed: 11404817473495 → {}", new_price);
4168                        println!(
4169                            "  - Publish time changed: 1761618783 → {}",
4170                            new_publish_time
4171                        );
4172                        println!("  - EMA price changed: 11421259300000 → {}", new_ema_price);
4173                        println!("  - Exponent unchanged: -8");
4174                    }
4175                    _ => panic!("Expected JSON parsed account data for modified account"),
4176                }
4177            }
4178            Err(e) => {
4179                // If it fails, it's due to Borsh/JSON mismatch (expected for now)
4180                println!("Expected error (Borsh vs JSON): {:?}", e);
4181                println!("Note: Once Borsh serialization is implemented, this test will:");
4182                println!("  1. Successfully modify the account data");
4183                println!("  2. Verify price changed to: {}", new_price);
4184                println!("  3. Verify publish_time changed to: {}", new_publish_time);
4185                println!("  4. Verify ema_price changed to: {}", new_ema_price);
4186                println!("  5. Verify other fields remain unchanged");
4187            }
4188        }
4189
4190        // Step 8: Demonstrate proper Borsh deserialization/serialization
4191        println!("\n--- Step 8: Testing with Borsh structures ---");
4192
4193        // Deserialize the original account data using Borsh
4194        let account_bytes = &account_data_hex[8..];
4195        println!(
4196            "Account data length (without discriminator): {} bytes",
4197            account_bytes.len()
4198        );
4199
4200        let mut reader = std::io::Cursor::new(account_bytes);
4201        let original_price_update = PriceUpdateV2::deserialize_reader(&mut reader)
4202            .expect("Should deserialize Pyth account data with Borsh");
4203
4204        let bytes_read = reader.position() as usize;
4205        println!("Bytes read by Borsh: {}", bytes_read);
4206        if bytes_read < account_bytes.len() {
4207            println!(
4208                "Note: {} extra bytes at end (likely padding)",
4209                account_bytes.len() - bytes_read
4210            );
4211        }
4212
4213        println!("Original Borsh-deserialized data:");
4214        println!("  - Price: {}", original_price_update.price_message.price);
4215        println!(
4216            "  - Exponent: {}",
4217            original_price_update.price_message.exponent
4218        );
4219        println!(
4220            "  - EMA Price: {}",
4221            original_price_update.price_message.ema_price
4222        );
4223        println!(
4224            "  - Publish time: {}",
4225            original_price_update.price_message.publish_time
4226        );
4227
4228        // Assert original values match what we saw in JSON parsing
4229        assert_eq!(
4230            original_price_update.price_message.price, 11404817473495,
4231            "Borsh price should match JSON parsed value"
4232        );
4233        assert_eq!(
4234            original_price_update.price_message.exponent, -8,
4235            "Borsh exponent should match JSON parsed value"
4236        );
4237        assert_eq!(
4238            original_price_update.price_message.ema_price, 11421259300000,
4239            "Borsh ema_price should match JSON parsed value"
4240        );
4241        assert_eq!(
4242            original_price_update.price_message.publish_time, 1761618783,
4243            "Borsh publish_time should match JSON parsed value"
4244        );
4245
4246        println!("✓ Borsh deserialization matches JSON parsing!");
4247
4248        // Step 9: Modify and re-serialize with Borsh
4249        println!("\n--- Step 9: Modifying account data with Borsh ---");
4250
4251        let mut modified_price_update = original_price_update.clone();
4252        modified_price_update.price_message.price = new_price;
4253        modified_price_update.price_message.publish_time = new_publish_time;
4254        modified_price_update.price_message.ema_price = new_ema_price;
4255
4256        // Serialize back to bytes
4257        let modified_account_data =
4258            borsh::to_vec(&modified_price_update).expect("Should serialize modified data");
4259
4260        // Prepend the discriminator
4261        let mut full_modified_data = account_data_hex[..8].to_vec();
4262        full_modified_data.extend_from_slice(&modified_account_data);
4263
4264        println!("Modified Borsh-serialized data:");
4265        println!(
4266            "  - Price: {} → {}",
4267            original_price_update.price_message.price, new_price
4268        );
4269        println!(
4270            "  - Publish time: {} → {}",
4271            original_price_update.price_message.publish_time, new_publish_time
4272        );
4273        println!(
4274            "  - EMA Price: {} → {}",
4275            original_price_update.price_message.ema_price, new_ema_price
4276        );
4277        println!(
4278            "  - Exponent: {} (unchanged)",
4279            modified_price_update.price_message.exponent
4280        );
4281
4282        // Verify the modified data is different
4283        assert_ne!(
4284            full_modified_data, account_data_hex,
4285            "Modified data should differ from original"
4286        );
4287
4288        // Verify we can deserialize the modified data back
4289        let mut modified_reader = std::io::Cursor::new(&full_modified_data[8..]);
4290        let reloaded_price_update = PriceUpdateV2::deserialize_reader(&mut modified_reader)
4291            .expect("Should deserialize modified data");
4292
4293        assert_eq!(
4294            reloaded_price_update.price_message.price, new_price,
4295            "Reloaded price should match modified value"
4296        );
4297        assert_eq!(
4298            reloaded_price_update.price_message.publish_time, new_publish_time,
4299            "Reloaded publish_time should match modified value"
4300        );
4301        assert_eq!(
4302            reloaded_price_update.price_message.ema_price, new_ema_price,
4303            "Reloaded ema_price should match modified value"
4304        );
4305        assert_eq!(
4306            reloaded_price_update.price_message.exponent,
4307            original_price_update.price_message.exponent,
4308            "Exponent should remain unchanged"
4309        );
4310
4311        println!("✓ Borsh round-trip successful!");
4312
4313        // Step 10: Verify with encode_ui_account
4314        println!("\n--- Step 10: Verify modified data with encode_ui_account ---");
4315
4316        let modified_test_account = Account {
4317            lamports: 1_000_000,
4318            data: full_modified_data,
4319            owner: account_pubkey,
4320            executable: false,
4321            rent_epoch: 0,
4322        };
4323
4324        let modified_ui_account = svm_locker.encode_ui_account(
4325            &account_pubkey,
4326            &modified_test_account,
4327            UiAccountEncoding::JsonParsed,
4328            None,
4329            None,
4330        );
4331
4332        // Verify through JSON encoding as well
4333        match &modified_ui_account.data {
4334            UiAccountData::Json(parsed_account) => {
4335                let parsed_obj = &parsed_account.parsed;
4336                let price_message = parsed_obj
4337                    .get("price_message")
4338                    .expect("Should have price_message")
4339                    .as_object()
4340                    .expect("Should be object");
4341
4342                let final_price = price_message
4343                    .get("price")
4344                    .expect("Should have price")
4345                    .as_i64()
4346                    .expect("Should be i64");
4347                let final_publish_time = price_message
4348                    .get("publish_time")
4349                    .expect("Should have publish_time")
4350                    .as_i64()
4351                    .expect("Should be i64");
4352                let final_ema_price = price_message
4353                    .get("ema_price")
4354                    .expect("Should have ema_price")
4355                    .as_i64()
4356                    .expect("Should be i64");
4357
4358                assert_eq!(
4359                    final_price, new_price,
4360                    "JSON-parsed price should match Borsh value"
4361                );
4362                assert_eq!(
4363                    final_publish_time, new_publish_time,
4364                    "JSON-parsed publish_time should match Borsh value"
4365                );
4366                assert_eq!(
4367                    final_ema_price, new_ema_price,
4368                    "JSON-parsed ema_price should match Borsh value"
4369                );
4370            }
4371            _ => panic!("Expected JSON parsed data"),
4372        }
4373    }
4374
4375    #[test]
4376    fn test_apply_override_to_decoded_account() {
4377        use txtx_addon_kit::{indexmap::IndexMap, types::types::Value};
4378
4379        // Create a txtx Value object
4380        let mut price_message_obj = IndexMap::new();
4381        price_message_obj.insert("price".to_string(), Value::Integer(100));
4382        price_message_obj.insert("publish_time".to_string(), Value::Integer(1234567890));
4383
4384        let mut decoded_value = IndexMap::new();
4385        decoded_value.insert(
4386            "price_message".to_string(),
4387            Value::Object(price_message_obj),
4388        );
4389        decoded_value.insert("expo".to_string(), Value::Integer(-8));
4390
4391        let mut decoded_value = Value::Object(decoded_value);
4392
4393        // Test simple override
4394        let result =
4395            apply_override_to_decoded_account(&mut decoded_value, "expo", &serde_json::json!(-6));
4396        assert!(result.is_ok());
4397        match &decoded_value {
4398            Value::Object(map) => {
4399                assert_eq!(map.get("expo"), Some(&Value::Integer(-6)));
4400            }
4401            _ => panic!("Expected Object"),
4402        }
4403
4404        // Test nested override
4405        let result = apply_override_to_decoded_account(
4406            &mut decoded_value,
4407            "price_message.price",
4408            &serde_json::json!(200),
4409        );
4410        assert!(result.is_ok());
4411        match &decoded_value {
4412            Value::Object(map) => match map.get("price_message") {
4413                Some(Value::Object(price_msg)) => {
4414                    assert_eq!(price_msg.get("price"), Some(&Value::Integer(200)));
4415                }
4416                _ => panic!("Expected price_message to be Object"),
4417            },
4418            _ => panic!("Expected Object"),
4419        }
4420
4421        // Test invalid path
4422        let result = apply_override_to_decoded_account(
4423            &mut decoded_value,
4424            "nonexistent.field",
4425            &serde_json::json!(999),
4426        );
4427        assert!(result.is_err());
4428    }
4429
4430    // Snapshot loading tests
4431
4432    #[tokio::test(flavor = "multi_thread")]
4433    async fn test_load_snapshot_basic() {
4434        use base64::{Engine, engine::general_purpose};
4435
4436        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
4437        let locker = SurfnetSvmLocker::new(svm);
4438
4439        let pubkey = Pubkey::new_unique();
4440        let owner = Pubkey::new_unique();
4441        let data = vec![1, 2, 3, 4, 5];
4442        let data_base64 = general_purpose::STANDARD.encode(&data);
4443
4444        let mut snapshot = BTreeMap::new();
4445        snapshot.insert(
4446            pubkey.to_string(),
4447            Some(AccountSnapshot {
4448                lamports: 1_000_000,
4449                owner: owner.to_string(),
4450                executable: false,
4451                rent_epoch: 0,
4452                data: data_base64,
4453                parsed_data: None,
4454            }),
4455        );
4456
4457        let loaded = locker
4458            .load_snapshot(&snapshot, None, CommitmentConfig::confirmed())
4459            .await
4460            .unwrap();
4461        assert_eq!(loaded, 1);
4462
4463        let account = locker
4464            .with_svm_reader(|svm| svm.get_account(&pubkey))
4465            .unwrap();
4466        assert!(account.is_some());
4467        let account = account.unwrap();
4468        assert_eq!(account.lamports, 1_000_000);
4469        assert_eq!(account.owner, owner);
4470        assert_eq!(account.data, data);
4471        assert!(!account.executable);
4472    }
4473
4474    #[tokio::test(flavor = "multi_thread")]
4475    async fn test_load_snapshot_multiple_accounts() {
4476        use base64::{Engine, engine::general_purpose};
4477
4478        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
4479        let locker = SurfnetSvmLocker::new(svm);
4480
4481        let owner = Pubkey::new_unique();
4482        let mut snapshot = BTreeMap::new();
4483
4484        // Add 5 accounts
4485        let pubkeys: Vec<Pubkey> = (0..5).map(|_| Pubkey::new_unique()).collect();
4486        for (i, pubkey) in pubkeys.iter().enumerate() {
4487            let data = vec![i as u8; 10];
4488            snapshot.insert(
4489                pubkey.to_string(),
4490                Some(AccountSnapshot {
4491                    lamports: (i as u64 + 1) * 1_000_000,
4492                    owner: owner.to_string(),
4493                    executable: false,
4494                    rent_epoch: 0,
4495                    data: general_purpose::STANDARD.encode(&data),
4496                    parsed_data: None,
4497                }),
4498            );
4499        }
4500
4501        let loaded = locker
4502            .load_snapshot(&snapshot, None, CommitmentConfig::confirmed())
4503            .await
4504            .unwrap();
4505        assert_eq!(loaded, 5);
4506
4507        // Verify all accounts were loaded
4508        for (i, pubkey) in pubkeys.iter().enumerate() {
4509            let account = locker
4510                .with_svm_reader(|svm| svm.get_account(pubkey))
4511                .unwrap()
4512                .unwrap();
4513            assert_eq!(account.lamports, (i as u64 + 1) * 1_000_000);
4514            assert_eq!(account.owner, owner);
4515            assert_eq!(account.data, vec![i as u8; 10]);
4516        }
4517    }
4518
4519    #[tokio::test(flavor = "multi_thread")]
4520    async fn test_load_snapshot_skips_none_without_remote() {
4521        use base64::{Engine, engine::general_purpose};
4522
4523        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
4524        let locker = SurfnetSvmLocker::new(svm);
4525
4526        let pubkey1 = Pubkey::new_unique();
4527        let pubkey2 = Pubkey::new_unique();
4528        let owner = Pubkey::new_unique();
4529
4530        let mut snapshot = BTreeMap::new();
4531
4532        // Add one real account
4533        snapshot.insert(
4534            pubkey1.to_string(),
4535            Some(AccountSnapshot {
4536                lamports: 1_000_000,
4537                owner: owner.to_string(),
4538                executable: false,
4539                rent_epoch: 0,
4540                data: general_purpose::STANDARD.encode(&[1, 2, 3]),
4541                parsed_data: None,
4542            }),
4543        );
4544
4545        // Add one None account (should be skipped without remote client)
4546        snapshot.insert(pubkey2.to_string(), None);
4547
4548        let loaded = locker
4549            .load_snapshot(&snapshot, None, CommitmentConfig::confirmed())
4550            .await
4551            .unwrap();
4552        assert_eq!(loaded, 1);
4553
4554        // First account should exist
4555        assert!(
4556            locker
4557                .with_svm_reader(|svm| svm.get_account(&pubkey1))
4558                .unwrap()
4559                .is_some()
4560        );
4561
4562        // Second account should not exist (no remote client to fetch it)
4563        assert!(
4564            locker
4565                .with_svm_reader(|svm| svm.get_account(&pubkey2))
4566                .unwrap()
4567                .is_none()
4568        );
4569    }
4570
4571    #[tokio::test(flavor = "multi_thread")]
4572    async fn test_load_snapshot_invalid_pubkey() {
4573        use base64::{Engine, engine::general_purpose};
4574
4575        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
4576        let locker = SurfnetSvmLocker::new(svm);
4577
4578        let owner = Pubkey::new_unique();
4579        let mut snapshot = BTreeMap::new();
4580
4581        // Add an invalid pubkey
4582        snapshot.insert(
4583            "invalid_pubkey".to_string(),
4584            Some(AccountSnapshot {
4585                lamports: 1_000_000,
4586                owner: owner.to_string(),
4587                executable: false,
4588                rent_epoch: 0,
4589                data: general_purpose::STANDARD.encode(&[1, 2, 3]),
4590                parsed_data: None,
4591            }),
4592        );
4593
4594        // Should succeed but load 0 accounts (invalid pubkey is skipped)
4595        let loaded = locker
4596            .load_snapshot(&snapshot, None, CommitmentConfig::confirmed())
4597            .await
4598            .unwrap();
4599        assert_eq!(loaded, 0);
4600    }
4601
4602    #[tokio::test(flavor = "multi_thread")]
4603    async fn test_load_snapshot_invalid_base64_data() {
4604        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
4605        let locker = SurfnetSvmLocker::new(svm);
4606
4607        let pubkey = Pubkey::new_unique();
4608        let owner = Pubkey::new_unique();
4609        let mut snapshot = BTreeMap::new();
4610
4611        // Add account with invalid base64 data
4612        snapshot.insert(
4613            pubkey.to_string(),
4614            Some(AccountSnapshot {
4615                lamports: 1_000_000,
4616                owner: owner.to_string(),
4617                executable: false,
4618                rent_epoch: 0,
4619                data: "not_valid_base64!!!".to_string(),
4620                parsed_data: None,
4621            }),
4622        );
4623
4624        // Should succeed but load 0 accounts (invalid data is skipped)
4625        let loaded = locker
4626            .load_snapshot(&snapshot, None, CommitmentConfig::confirmed())
4627            .await
4628            .unwrap();
4629        assert_eq!(loaded, 0);
4630
4631        // Account should not exist
4632        assert!(
4633            locker
4634                .with_svm_reader(|svm| svm.get_account(&pubkey))
4635                .unwrap()
4636                .is_none()
4637        );
4638    }
4639
4640    #[tokio::test(flavor = "multi_thread")]
4641    async fn test_load_snapshot_invalid_owner() {
4642        use base64::{Engine, engine::general_purpose};
4643
4644        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
4645        let locker = SurfnetSvmLocker::new(svm);
4646
4647        let pubkey = Pubkey::new_unique();
4648        let mut snapshot = BTreeMap::new();
4649
4650        // Add account with invalid owner pubkey
4651        snapshot.insert(
4652            pubkey.to_string(),
4653            Some(AccountSnapshot {
4654                lamports: 1_000_000,
4655                owner: "invalid_owner".to_string(),
4656                executable: false,
4657                rent_epoch: 0,
4658                data: general_purpose::STANDARD.encode(&[1, 2, 3]),
4659                parsed_data: None,
4660            }),
4661        );
4662
4663        // Should succeed but load 0 accounts (invalid owner is skipped)
4664        let loaded = locker
4665            .load_snapshot(&snapshot, None, CommitmentConfig::confirmed())
4666            .await
4667            .unwrap();
4668        assert_eq!(loaded, 0);
4669    }
4670
4671    #[tokio::test(flavor = "multi_thread")]
4672    async fn test_load_snapshot_empty() {
4673        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
4674        let locker = SurfnetSvmLocker::new(svm);
4675
4676        let snapshot = BTreeMap::new();
4677        let loaded = locker
4678            .load_snapshot(&snapshot, None, CommitmentConfig::confirmed())
4679            .await
4680            .unwrap();
4681        assert_eq!(loaded, 0);
4682    }
4683
4684    #[tokio::test(flavor = "multi_thread")]
4685    async fn test_load_snapshot_updates_account_registries() {
4686        use base64::{Engine, engine::general_purpose};
4687
4688        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
4689        let locker = SurfnetSvmLocker::new(svm);
4690
4691        let pubkey = Pubkey::new_unique();
4692        let owner = Pubkey::new_unique();
4693
4694        let mut snapshot = BTreeMap::new();
4695        snapshot.insert(
4696            pubkey.to_string(),
4697            Some(AccountSnapshot {
4698                lamports: 1_000_000,
4699                owner: owner.to_string(),
4700                executable: false,
4701                rent_epoch: 0,
4702                data: general_purpose::STANDARD.encode(&[1, 2, 3]),
4703                parsed_data: None,
4704            }),
4705        );
4706
4707        locker
4708            .load_snapshot(&snapshot, None, CommitmentConfig::confirmed())
4709            .await
4710            .unwrap();
4711
4712        // Verify account is in the owner index
4713        let owned_accounts = locker
4714            .with_svm_reader(|svm| svm.get_account_owned_by(&owner))
4715            .unwrap();
4716        assert_eq!(owned_accounts.len(), 1);
4717        assert_eq!(owned_accounts[0].0, pubkey);
4718    }
4719
4720    #[tokio::test(flavor = "multi_thread")]
4721    async fn test_load_snapshot_mixed_valid_invalid() {
4722        use base64::{Engine, engine::general_purpose};
4723
4724        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
4725        let locker = SurfnetSvmLocker::new(svm);
4726
4727        let valid_pubkey = Pubkey::new_unique();
4728        let owner = Pubkey::new_unique();
4729
4730        let mut snapshot = BTreeMap::new();
4731
4732        // Valid account
4733        snapshot.insert(
4734            valid_pubkey.to_string(),
4735            Some(AccountSnapshot {
4736                lamports: 1_000_000,
4737                owner: owner.to_string(),
4738                executable: false,
4739                rent_epoch: 0,
4740                data: general_purpose::STANDARD.encode(&[1, 2, 3]),
4741                parsed_data: None,
4742            }),
4743        );
4744
4745        // Invalid pubkey
4746        snapshot.insert(
4747            "bad_pubkey".to_string(),
4748            Some(AccountSnapshot {
4749                lamports: 2_000_000,
4750                owner: owner.to_string(),
4751                executable: false,
4752                rent_epoch: 0,
4753                data: general_purpose::STANDARD.encode(&[4, 5, 6]),
4754                parsed_data: None,
4755            }),
4756        );
4757
4758        // None value (skipped without remote)
4759        snapshot.insert(Pubkey::new_unique().to_string(), None);
4760
4761        let loaded = locker
4762            .load_snapshot(&snapshot, None, CommitmentConfig::confirmed())
4763            .await
4764            .unwrap();
4765        assert_eq!(loaded, 1);
4766
4767        // Only the valid account should exist
4768        assert!(
4769            locker
4770                .with_svm_reader(|svm| svm.get_account(&valid_pubkey))
4771                .unwrap()
4772                .is_some()
4773        );
4774    }
4775
4776    /// Helper: create a VersionedTransaction with a given signature whose account keys contain `pubkey`.
4777    fn make_test_tx(sig: Signature, pubkey: &Pubkey) -> VersionedTransaction {
4778        use solana_system_interface::instruction as system_instruction;
4779        VersionedTransaction {
4780            signatures: vec![sig],
4781            message: VersionedMessage::Legacy(Message::new(
4782                &[system_instruction::transfer(pubkey, pubkey, 1)],
4783                Some(pubkey),
4784            )),
4785        }
4786    }
4787
4788    /// Helper: store a transaction into the SVM at the given slot.
4789    fn store_test_tx(svm: &mut SurfnetSvm, sig: Signature, pubkey: &Pubkey, slot: u64) {
4790        let tx = make_test_tx(sig, pubkey);
4791        svm.transactions
4792            .store(
4793                sig.to_string(),
4794                SurfnetTransactionStatus::processed(
4795                    TransactionWithStatusMeta {
4796                        slot,
4797                        transaction: tx,
4798                        meta: TransactionStatusMeta {
4799                            status: Ok(()),
4800                            fee: 5000,
4801                            pre_balances: vec![0; 3],
4802                            post_balances: vec![0; 3],
4803                            inner_instructions: Some(vec![]),
4804                            log_messages: Some(vec![]),
4805                            pre_token_balances: Some(vec![]),
4806                            post_token_balances: Some(vec![]),
4807                            rewards: Some(vec![]),
4808                            loaded_addresses: LoadedAddresses::default(),
4809                            return_data: None,
4810                            compute_units_consumed: Some(0),
4811                            cost_units: None,
4812                        },
4813                    },
4814                    HashSet::new(),
4815                ),
4816            )
4817            .unwrap();
4818    }
4819
4820    fn seed_signature_history(
4821        locker: &SurfnetSvmLocker,
4822        pubkey: &Pubkey,
4823        blocks: &[(u64, Vec<Signature>)],
4824    ) {
4825        locker.with_svm_writer(|svm| {
4826            for (slot, signatures) in blocks {
4827                for sig in signatures {
4828                    store_test_tx(svm, *sig, pubkey, *slot);
4829                }
4830
4831                svm.blocks
4832                    .store(
4833                        *slot,
4834                        BlockHeader {
4835                            hash: String::new(),
4836                            previous_blockhash: String::new(),
4837                            parent_slot: 0,
4838                            block_time: 0,
4839                            block_height: 0,
4840                            signatures: signatures.clone(),
4841                        },
4842                    )
4843                    .unwrap();
4844            }
4845        });
4846    }
4847
4848    fn fetch_signature_strings(
4849        locker: &SurfnetSvmLocker,
4850        pubkey: &Pubkey,
4851        config: Option<&RpcSignaturesForAddressConfig>,
4852    ) -> Vec<String> {
4853        locker
4854            .get_signatures_for_address_local(pubkey, config)
4855            .inner
4856            .iter()
4857            .map(|s| s.signature.clone())
4858            .collect()
4859    }
4860
4861    #[tokio::test(flavor = "multi_thread")]
4862    async fn test_get_signatures_for_address_ordering_within_block() {
4863        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
4864        let locker = SurfnetSvmLocker::new(svm);
4865
4866        let pubkey = Pubkey::new_unique();
4867        let sig_a = Signature::new_unique();
4868        let sig_b = Signature::new_unique();
4869        let sig_c = Signature::new_unique();
4870        let slot = 5;
4871
4872        seed_signature_history(&locker, &pubkey, &[(slot, vec![sig_a, sig_b, sig_c])]);
4873        let sigs = fetch_signature_strings(&locker, &pubkey, None);
4874
4875        // Last executed (C) should appear first, then B, then A
4876        assert_eq!(sigs.len(), 3);
4877        assert_eq!(sigs[0], sig_c.to_string());
4878        assert_eq!(sigs[1], sig_b.to_string());
4879        assert_eq!(sigs[2], sig_a.to_string());
4880    }
4881
4882    #[tokio::test(flavor = "multi_thread")]
4883    async fn test_get_signatures_for_address_until_excludes_boundary_signature() {
4884        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
4885        let locker = SurfnetSvmLocker::new(svm);
4886
4887        let pubkey = Pubkey::new_unique();
4888        let sig_a = Signature::new_unique();
4889        let sig_b = Signature::new_unique();
4890        let sig_c = Signature::new_unique();
4891        let slot = 5;
4892
4893        seed_signature_history(&locker, &pubkey, &[(slot, vec![sig_a, sig_b, sig_c])]);
4894        let sigs = fetch_signature_strings(
4895            &locker,
4896            &pubkey,
4897            Some(&RpcSignaturesForAddressConfig {
4898                until: Some(sig_b.to_string()),
4899                ..RpcSignaturesForAddressConfig::default()
4900            }),
4901        );
4902
4903        assert_eq!(sigs, vec![sig_c.to_string()]);
4904    }
4905
4906    #[tokio::test(flavor = "multi_thread")]
4907    async fn test_get_signatures_for_address_before_excludes_boundary_signature() {
4908        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
4909        let locker = SurfnetSvmLocker::new(svm);
4910
4911        let pubkey = Pubkey::new_unique();
4912        let sig_a = Signature::new_unique();
4913        let sig_b = Signature::new_unique();
4914        let sig_c = Signature::new_unique();
4915        let slot = 5;
4916
4917        seed_signature_history(&locker, &pubkey, &[(slot, vec![sig_a, sig_b, sig_c])]);
4918        let sigs = fetch_signature_strings(
4919            &locker,
4920            &pubkey,
4921            Some(&RpcSignaturesForAddressConfig {
4922                before: Some(sig_b.to_string()),
4923                ..RpcSignaturesForAddressConfig::default()
4924            }),
4925        );
4926
4927        assert_eq!(sigs, vec![sig_a.to_string()]);
4928    }
4929
4930    #[tokio::test(flavor = "multi_thread")]
4931    async fn test_get_signatures_for_address_before_and_until_form_window() {
4932        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
4933        let locker = SurfnetSvmLocker::new(svm);
4934
4935        let pubkey = Pubkey::new_unique();
4936        let sig_a = Signature::new_unique();
4937        let sig_b = Signature::new_unique();
4938        let sig_c = Signature::new_unique();
4939        let sig_d = Signature::new_unique();
4940        let slot = 5;
4941
4942        seed_signature_history(
4943            &locker,
4944            &pubkey,
4945            &[(slot, vec![sig_a, sig_b, sig_c, sig_d])],
4946        );
4947        let sigs = fetch_signature_strings(
4948            &locker,
4949            &pubkey,
4950            Some(&RpcSignaturesForAddressConfig {
4951                before: Some(sig_d.to_string()),
4952                until: Some(sig_b.to_string()),
4953                ..RpcSignaturesForAddressConfig::default()
4954            }),
4955        );
4956
4957        assert_eq!(sigs, vec![sig_c.to_string()]);
4958    }
4959
4960    #[tokio::test(flavor = "multi_thread")]
4961    async fn test_get_signatures_for_address_before_missing_returns_empty() {
4962        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
4963        let locker = SurfnetSvmLocker::new(svm);
4964
4965        let pubkey = Pubkey::new_unique();
4966        let sig_a = Signature::new_unique();
4967        let sig_b = Signature::new_unique();
4968        let missing_sig = Signature::new_unique();
4969        let slot = 5;
4970
4971        seed_signature_history(&locker, &pubkey, &[(slot, vec![sig_a, sig_b])]);
4972        let sigs = fetch_signature_strings(
4973            &locker,
4974            &pubkey,
4975            Some(&RpcSignaturesForAddressConfig {
4976                before: Some(missing_sig.to_string()),
4977                ..RpcSignaturesForAddressConfig::default()
4978            }),
4979        );
4980
4981        assert!(sigs.is_empty());
4982    }
4983
4984    #[tokio::test(flavor = "multi_thread")]
4985    async fn test_get_signatures_for_address_until_missing_returns_all_results() {
4986        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
4987        let locker = SurfnetSvmLocker::new(svm);
4988
4989        let pubkey = Pubkey::new_unique();
4990        let sig_a = Signature::new_unique();
4991        let sig_b = Signature::new_unique();
4992        let missing_sig = Signature::new_unique();
4993        let slot = 5;
4994
4995        seed_signature_history(&locker, &pubkey, &[(slot, vec![sig_a, sig_b])]);
4996        let sigs = fetch_signature_strings(
4997            &locker,
4998            &pubkey,
4999            Some(&RpcSignaturesForAddressConfig {
5000                until: Some(missing_sig.to_string()),
5001                ..RpcSignaturesForAddressConfig::default()
5002            }),
5003        );
5004
5005        assert_eq!(sigs, vec![sig_b.to_string(), sig_a.to_string()]);
5006    }
5007
5008    #[tokio::test(flavor = "multi_thread")]
5009    async fn test_get_signatures_for_address_limit_applies_after_windowing() {
5010        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
5011        let locker = SurfnetSvmLocker::new(svm);
5012
5013        let pubkey = Pubkey::new_unique();
5014        let sig_a = Signature::new_unique();
5015        let sig_b = Signature::new_unique();
5016        let sig_c = Signature::new_unique();
5017        let sig_d = Signature::new_unique();
5018        let sig_e = Signature::new_unique();
5019        let slot = 5;
5020
5021        seed_signature_history(
5022            &locker,
5023            &pubkey,
5024            &[(slot, vec![sig_a, sig_b, sig_c, sig_d, sig_e])],
5025        );
5026        let sigs = fetch_signature_strings(
5027            &locker,
5028            &pubkey,
5029            Some(&RpcSignaturesForAddressConfig {
5030                before: Some(sig_e.to_string()),
5031                until: Some(sig_a.to_string()),
5032                limit: Some(2),
5033                ..RpcSignaturesForAddressConfig::default()
5034            }),
5035        );
5036
5037        assert_eq!(sigs, vec![sig_d.to_string(), sig_c.to_string()]);
5038    }
5039
5040    #[tokio::test(flavor = "multi_thread")]
5041    async fn test_get_signatures_for_address_until_excludes_boundary_across_slots() {
5042        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
5043        let locker = SurfnetSvmLocker::new(svm);
5044
5045        let pubkey = Pubkey::new_unique();
5046        let sig_s5 = Signature::new_unique();
5047        let sig_s10_a = Signature::new_unique();
5048        let sig_s10_b = Signature::new_unique();
5049        let sig_s15 = Signature::new_unique();
5050
5051        seed_signature_history(
5052            &locker,
5053            &pubkey,
5054            &[
5055                (5, vec![sig_s5]),
5056                (10, vec![sig_s10_a, sig_s10_b]),
5057                (15, vec![sig_s15]),
5058            ],
5059        );
5060        let sigs = fetch_signature_strings(
5061            &locker,
5062            &pubkey,
5063            Some(&RpcSignaturesForAddressConfig {
5064                until: Some(sig_s10_b.to_string()),
5065                ..RpcSignaturesForAddressConfig::default()
5066            }),
5067        );
5068
5069        assert_eq!(sigs, vec![sig_s15.to_string()]);
5070    }
5071
5072    #[tokio::test(flavor = "multi_thread")]
5073    async fn test_get_signatures_for_address_ordering_across_slots() {
5074        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
5075        let locker = SurfnetSvmLocker::new(svm);
5076
5077        let pubkey = Pubkey::new_unique();
5078        let sig_s5_a = Signature::new_unique();
5079        let sig_s5_b = Signature::new_unique();
5080        let sig_s10_a = Signature::new_unique();
5081        let sig_s10_b = Signature::new_unique();
5082
5083        seed_signature_history(
5084            &locker,
5085            &pubkey,
5086            &[
5087                (5, vec![sig_s5_a, sig_s5_b]),
5088                (10, vec![sig_s10_a, sig_s10_b]),
5089            ],
5090        );
5091        let sigs = fetch_signature_strings(&locker, &pubkey, None);
5092
5093        // Slot 10 txs first (descending), then slot 5 txs
5094        // Within each slot: last executed first
5095        assert_eq!(sigs.len(), 4);
5096        assert_eq!(sigs[0], sig_s10_b.to_string());
5097        assert_eq!(sigs[1], sig_s10_a.to_string());
5098        assert_eq!(sigs[2], sig_s5_b.to_string());
5099        assert_eq!(sigs[3], sig_s5_a.to_string());
5100    }
5101
5102    #[tokio::test(flavor = "multi_thread")]
5103    async fn test_get_signatures_for_address_ordering_missing_block_header() {
5104        let (svm, _events_rx, _geyser_rx) = SurfnetSvm::default();
5105        let locker = SurfnetSvmLocker::new(svm);
5106
5107        let pubkey = Pubkey::new_unique();
5108        let sig_a = Signature::new_unique();
5109        let sig_b = Signature::new_unique();
5110        let slot = 5;
5111
5112        locker.with_svm_writer(|svm| {
5113            store_test_tx(svm, sig_a, &pubkey, slot);
5114            store_test_tx(svm, sig_b, &pubkey, slot);
5115            // No block header stored — should not panic
5116        });
5117
5118        let result = locker.get_signatures_for_address_local(&pubkey, None);
5119
5120        // Both transactions should be returned regardless
5121        assert_eq!(result.inner.len(), 2);
5122
5123        // Verify both signatures are present (order not guaranteed without block header)
5124        let sigs: HashSet<String> = result.inner.iter().map(|s| s.signature.clone()).collect();
5125        assert!(sigs.contains(&sig_a.to_string()));
5126        assert!(sigs.contains(&sig_b.to_string()));
5127    }
5128
5129    #[tokio::test(flavor = "multi_thread")]
5130    async fn initializes_epoch_schedule_without_warmup_when_offline() {
5131        let (surfnet_svm, _simnet_events_rx, _geyser_events_rx) = SurfnetSvm::default();
5132        let svm_locker = SurfnetSvmLocker::new(surfnet_svm);
5133
5134        svm_locker
5135            .initialize(&None)
5136            .await
5137            .expect("initialize should succeed");
5138
5139        let epoch_schedule =
5140            svm_locker.with_svm_reader(|svm_reader| svm_reader.inner.get_sysvar::<EpochSchedule>());
5141
5142        assert!(
5143            !epoch_schedule.warmup,
5144            "offline initialization should disable warmup to match mainnet"
5145        );
5146        assert_eq!(
5147            epoch_schedule.get_first_slot_in_epoch(886),
5148            886_u64 * 432_000,
5149            "first slot should align with mainnet epoch boundaries when warmup is disabled"
5150        );
5151    }
5152}