Skip to main content

surfpool_core/surfnet/
locker.rs

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