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