Skip to main content

surfpool_sdk/
surfnet.rs

1use std::{
2    net::TcpListener,
3    thread::sleep,
4    time::{Duration, Instant},
5};
6
7use crossbeam_channel::{Receiver, Sender};
8use solana_commitment_config::CommitmentConfig;
9use solana_keypair::Keypair;
10use solana_pubkey::Pubkey;
11use solana_rpc_client::rpc_client::RpcClient;
12use solana_signer::Signer;
13use surfpool_core::surfnet::{
14    locker::SurfnetSvmLocker,
15    svm::{SurfnetSvm, SurfnetSvmConfig},
16};
17use surfpool_types::{
18    BlockProductionMode, RpcConfig, SimnetCommand, SimnetConfig, SimnetEvent, SurfpoolConfig,
19    SvmFeatureConfig,
20};
21
22use crate::{
23    Cheatcodes,
24    error::{SurfnetError, SurfnetResult},
25};
26
27/// Builder for configuring a [`Surfnet`] instance before starting it.
28///
29/// ```rust
30/// use surfpool_sdk::{Surfnet, BlockProductionMode};
31///
32/// # async fn example() {
33/// let surfnet = Surfnet::builder()
34///     .offline(true)
35///     .block_production_mode(BlockProductionMode::Transaction)
36///     .skip_blockhash_check(true)
37///     .airdrop_sol(10_000_000_000)
38///     .start()
39///     .await
40///     .unwrap();
41/// # }
42/// ```
43pub struct SurfnetBuilder {
44    offline_mode: bool,
45    remote_rpc_url: Option<String>,
46    block_production_mode: BlockProductionMode,
47    slot_time_ms: u64,
48    airdrop_addresses: Vec<Pubkey>,
49    airdrop_lamports: u64,
50    skip_blockhash_check: bool,
51    payer: Option<Keypair>,
52    feature_config: SvmFeatureConfig,
53}
54
55impl Default for SurfnetBuilder {
56    fn default() -> Self {
57        Self {
58            offline_mode: true,
59            remote_rpc_url: None,
60            block_production_mode: BlockProductionMode::Transaction,
61            slot_time_ms: 1,
62            airdrop_addresses: vec![],
63            airdrop_lamports: 10_000_000_000, // 10 SOL
64            skip_blockhash_check: false,
65            payer: None,
66            feature_config: SvmFeatureConfig::default(),
67        }
68    }
69}
70
71impl SurfnetBuilder {
72    /// Run in offline mode (no mainnet RPC fallback). Default: `true`.
73    pub fn offline(mut self, offline: bool) -> Self {
74        self.offline_mode = offline;
75        self
76    }
77
78    /// Set a remote RPC URL for account fallback (implies `offline(false)`).
79    pub fn remote_rpc_url(mut self, url: impl Into<String>) -> Self {
80        self.remote_rpc_url = Some(url.into());
81        self.offline_mode = false;
82        self
83    }
84
85    /// How blocks are produced. Default: `Transaction` (advance on each tx).
86    pub fn block_production_mode(mut self, mode: BlockProductionMode) -> Self {
87        self.block_production_mode = mode;
88        self
89    }
90
91    /// Slot time in milliseconds. Default: `1` (fast for tests).
92    pub fn slot_time_ms(mut self, ms: u64) -> Self {
93        self.slot_time_ms = ms;
94        self
95    }
96
97    /// Additional addresses to airdrop SOL to at startup.
98    pub fn airdrop_addresses(mut self, addresses: Vec<Pubkey>) -> Self {
99        self.airdrop_addresses = addresses;
100        self
101    }
102
103    /// Amount of lamports to airdrop to the payer (and additional addresses) at startup.
104    /// Default: 10 SOL.
105    pub fn airdrop_sol(mut self, lamports: u64) -> Self {
106        self.airdrop_lamports = lamports;
107        self
108    }
109
110    /// Skip blockhash validation for all transactions in this surfnet instance.
111    pub fn skip_blockhash_check(mut self, skip: bool) -> Self {
112        self.skip_blockhash_check = skip;
113        self
114    }
115
116    /// Use a specific keypair as the payer. If not set, a random one is generated.
117    pub fn payer(mut self, keypair: Keypair) -> Self {
118        self.payer = Some(keypair);
119        self
120    }
121
122    /// Enable an SVM feature gate at startup. Mirrors the CLI's `--feature` flag.
123    pub fn enable_feature(mut self, feature: Pubkey) -> Self {
124        self.feature_config = self.feature_config.enable(feature);
125        self
126    }
127
128    /// Disable an SVM feature gate at startup. Mirrors the CLI's `--disable-feature` flag.
129    pub fn disable_feature(mut self, feature: Pubkey) -> Self {
130        self.feature_config = self.feature_config.disable(feature);
131        self
132    }
133
134    /// Replace the feature config wholesale. Default: [`SvmFeatureConfig::default_mainnet_features`].
135    pub fn feature_config(mut self, config: SvmFeatureConfig) -> Self {
136        self.feature_config = config;
137        self
138    }
139
140    /// Start the surfnet with the configured options.
141    pub async fn start(self) -> SurfnetResult<Surfnet> {
142        let SurfnetBuilder {
143            offline_mode,
144            remote_rpc_url,
145            block_production_mode,
146            slot_time_ms,
147            airdrop_addresses,
148            airdrop_lamports,
149            skip_blockhash_check,
150            payer,
151            feature_config,
152        } = self;
153        let payer = payer.unwrap_or_else(Keypair::new);
154
155        let bind_port = get_free_port()?;
156        let ws_port = get_free_port()?;
157        let bind_host = "127.0.0.1".to_string();
158
159        let mut startup_airdrop_addresses = vec![payer.pubkey()];
160        startup_airdrop_addresses.extend(airdrop_addresses);
161        let startup_airdrop_addresses_for_rpc = startup_airdrop_addresses.clone();
162
163        let surfpool_config = SurfpoolConfig {
164            simnets: vec![SimnetConfig {
165                offline_mode,
166                remote_rpc_url,
167                slot_time: slot_time_ms,
168                block_production_mode,
169                airdrop_addresses: startup_airdrop_addresses,
170                airdrop_token_amount: airdrop_lamports,
171                skip_blockhash_check,
172                ..Default::default()
173            }],
174            rpc: RpcConfig {
175                bind_host: bind_host.clone(),
176                bind_port,
177                ws_port,
178                ..Default::default()
179            },
180            ..Default::default()
181        };
182
183        let rpc_url = format!("http://{bind_host}:{bind_port}");
184        let ws_url = format!("ws://{bind_host}:{ws_port}");
185
186        let svm_config = SurfnetSvmConfig {
187            surfnet_id: surfpool_config.simnets[0].surfnet_id.clone(),
188            slot_time: surfpool_config.simnets[0].slot_time,
189            instruction_profiling_enabled: surfpool_config.simnets[0].instruction_profiling_enabled,
190            max_profiles: surfpool_config.simnets[0].max_profiles,
191            log_bytes_limit: surfpool_config.simnets[0].log_bytes_limit,
192            feature_config,
193            skip_blockhash_check,
194        };
195        let (surfnet_svm, simnet_events_rx, geyser_events_rx) = SurfnetSvm::new(svm_config)
196            .map_err(|e| SurfnetError::Runtime(format!("failed to initialize Surfnet SVM: {e}")))?;
197        let (simnet_commands_tx, simnet_commands_rx) = crossbeam_channel::unbounded();
198
199        let svm_locker = SurfnetSvmLocker::new(surfnet_svm);
200        let svm_locker_clone = svm_locker.clone();
201        let simnet_commands_tx_clone = simnet_commands_tx.clone();
202
203        let _handle = std::thread::Builder::new()
204            .name("surfnet-sdk".into())
205            .spawn(move || {
206                let future = surfpool_core::runloops::start_local_surfnet_runloop(
207                    svm_locker_clone,
208                    surfpool_config,
209                    simnet_commands_tx_clone,
210                    simnet_commands_rx,
211                    geyser_events_rx,
212                );
213                if let Err(e) = hiro_system_kit::nestable_block_on(future) {
214                    log::error!("Surfnet exited with error: {e}");
215                }
216            })
217            .map_err(|e| SurfnetError::Runtime(e.to_string()))?;
218
219        // Wait for the runtime to signal ready
220        wait_for_ready(&simnet_events_rx)?;
221        wait_for_startup_airdrops(
222            &rpc_url,
223            &startup_airdrop_addresses_for_rpc,
224            airdrop_lamports,
225        )?;
226
227        Ok(Surfnet {
228            rpc_url,
229            ws_url,
230            payer,
231            simnet_commands_tx,
232            simnet_events_rx,
233            svm_locker,
234            instance_id: uuid::Uuid::new_v4().to_string(),
235            stopped: false,
236        })
237    }
238}
239
240/// A running Surfpool instance with RPC/WS endpoints on dynamic ports.
241///
242/// Provides:
243/// - Pre-funded payer keypair
244/// - [`RpcClient`] connected to the local instance
245/// - [`Cheatcodes`] for direct state manipulation (fund accounts, set token balances, etc.)
246///
247/// The instance is shut down when dropped.
248pub struct Surfnet {
249    rpc_url: String,
250    ws_url: String,
251    payer: Keypair,
252    simnet_commands_tx: Sender<SimnetCommand>,
253    simnet_events_rx: Receiver<SimnetEvent>,
254    #[allow(dead_code)] // retained for future direct profiling access
255    svm_locker: SurfnetSvmLocker,
256    instance_id: String,
257    stopped: bool,
258}
259
260impl Surfnet {
261    /// Start a surfnet with default settings (offline, transaction-mode blocks, 10 SOL payer).
262    pub async fn start() -> SurfnetResult<Self> {
263        SurfnetBuilder::default().start().await
264    }
265
266    /// Create a builder for custom configuration.
267    pub fn builder() -> SurfnetBuilder {
268        SurfnetBuilder::default()
269    }
270
271    /// The HTTP RPC URL (e.g. `http://127.0.0.1:12345`).
272    pub fn rpc_url(&self) -> &str {
273        &self.rpc_url
274    }
275
276    /// The WebSocket URL (e.g. `ws://127.0.0.1:12346`).
277    pub fn ws_url(&self) -> &str {
278        &self.ws_url
279    }
280
281    /// Create a new [`RpcClient`] connected to this surfnet.
282    pub fn rpc_client(&self) -> RpcClient {
283        RpcClient::new(&self.rpc_url)
284    }
285
286    /// The pre-funded payer keypair.
287    pub fn payer(&self) -> &Keypair {
288        &self.payer
289    }
290
291    /// Access cheatcode helpers for direct state manipulation.
292    pub fn cheatcodes(&self) -> Cheatcodes<'_> {
293        Cheatcodes::new(&self.rpc_url)
294    }
295
296    /// Get a reference to the simnet events receiver for observing runtime events.
297    pub fn events(&self) -> &Receiver<SimnetEvent> {
298        &self.simnet_events_rx
299    }
300
301    /// Send a command to the simnet runtime.
302    pub fn send_command(&self, command: SimnetCommand) -> SurfnetResult<()> {
303        self.simnet_commands_tx
304            .send(command)
305            .map_err(|e| SurfnetError::Runtime(format!("failed to send command: {e}")))
306    }
307
308    /// The unique instance ID for this surfnet.
309    pub fn instance_id(&self) -> &str {
310        &self.instance_id
311    }
312
313    /// Gracefully shut down the surfnet, closing the HTTP + WebSocket RPC
314    /// servers and freeing their ports.
315    ///
316    /// Blocks until both RPC servers acknowledge shutdown or the timeout
317    /// elapses. Returns an error if shutdown is not confirmed within the
318    /// timeout (the port may still be bound) — callers that need a
319    /// guaranteed-free port should treat this as fatal. On success, the
320    /// instance is marked stopped and subsequent calls are a no-op.
321    ///
322    /// Note: this drains the simnet events channel while waiting. Don't call
323    /// `events()` / `drain_events()` concurrently from another thread or
324    /// shutdown acknowledgements may be lost, causing this call to time out.
325    pub fn stop(&mut self) -> SurfnetResult<()> {
326        if self.stopped {
327            return Ok(());
328        }
329
330        self.simnet_commands_tx
331            .send(SimnetCommand::Terminate(None))
332            .map_err(|e| SurfnetError::Runtime(format!("failed to send terminate command: {e}")))?;
333
334        let timeout = Duration::from_secs(5);
335        let deadline = Instant::now() + timeout;
336        let mut shutdowns_seen = 0;
337        while shutdowns_seen < 2 {
338            let remaining = match deadline.checked_duration_since(Instant::now()) {
339                Some(d) => d,
340                None => break,
341            };
342            match self.simnet_events_rx.recv_timeout(remaining) {
343                Ok(SimnetEvent::Shutdown) => shutdowns_seen += 1,
344                Ok(_) => continue,
345                Err(_) => break,
346            }
347        }
348
349        if shutdowns_seen < 2 {
350            return Err(SurfnetError::Runtime(format!(
351                "surfnet shutdown not confirmed within {timeout:?}: {shutdowns_seen} of 2 RPC servers acknowledged. The port may still be bound."
352            )));
353        }
354
355        self.stopped = true;
356        Ok(())
357    }
358}
359
360impl Drop for Surfnet {
361    fn drop(&mut self) {
362        if !self.stopped {
363            let _ = self.simnet_commands_tx.send(SimnetCommand::Terminate(None));
364        }
365    }
366}
367
368fn get_free_port() -> SurfnetResult<u16> {
369    let listener = TcpListener::bind("127.0.0.1:0")
370        .map_err(|e| SurfnetError::PortAllocation(e.to_string()))?;
371    let port = listener
372        .local_addr()
373        .map_err(|e| SurfnetError::PortAllocation(e.to_string()))?
374        .port();
375    drop(listener);
376    Ok(port)
377}
378
379fn wait_for_ready(events_rx: &Receiver<SimnetEvent>) -> SurfnetResult<()> {
380    loop {
381        match events_rx.recv() {
382            Ok(SimnetEvent::Ready(_)) => return Ok(()),
383            Ok(SimnetEvent::Aborted(err)) => return Err(SurfnetError::Aborted(err)),
384            Ok(SimnetEvent::Shutdown) => {
385                return Err(SurfnetError::Aborted(
386                    "surfnet shut down during startup".into(),
387                ));
388            }
389            Ok(_) => continue,
390            Err(e) => {
391                return Err(SurfnetError::Startup(format!(
392                    "events channel closed unexpectedly: {e}"
393                )));
394            }
395        }
396    }
397}
398
399fn wait_for_startup_airdrops(
400    rpc_url: &str,
401    addresses: &[Pubkey],
402    expected_lamports: u64,
403) -> SurfnetResult<()> {
404    let rpc_client = RpcClient::new(rpc_url.to_string());
405    let deadline = Instant::now() + Duration::from_secs(5);
406    let mut last_error = None;
407    let mut last_balances = vec![];
408
409    while Instant::now() < deadline {
410        last_balances.clear();
411        let mut all_match = true;
412
413        for address in addresses {
414            match rpc_client.get_balance_with_commitment(address, CommitmentConfig::processed()) {
415                Ok(response) => {
416                    last_balances.push((address.to_string(), response.value));
417                    if response.value != expected_lamports {
418                        all_match = false;
419                    }
420                }
421                Err(err) => {
422                    last_error = Some(err.to_string());
423                    all_match = false;
424                    break;
425                }
426            }
427        }
428
429        if all_match {
430            return Ok(());
431        }
432
433        sleep(Duration::from_millis(25));
434    }
435
436    let balance_summary = if last_balances.is_empty() {
437        "no balances observed".to_string()
438    } else {
439        last_balances
440            .iter()
441            .map(|(address, balance)| format!("{address}={balance}"))
442            .collect::<Vec<_>>()
443            .join(", ")
444    };
445
446    Err(SurfnetError::Startup(format!(
447        "startup balances not visible over RPC within timeout (expected {expected_lamports}); last balances: {balance_summary}; last error: {}",
448        last_error.unwrap_or_else(|| "none".to_string())
449    )))
450}
451
452#[cfg(test)]
453mod tests {
454    use super::*;
455
456    #[test]
457    fn surfnet_builder_skip_blockhash_check_defaults_to_false() {
458        let builder = SurfnetBuilder::default();
459        assert!(!builder.skip_blockhash_check);
460    }
461
462    #[test]
463    fn surfnet_builder_skip_blockhash_check_setter_updates_builder() {
464        let builder = SurfnetBuilder::default().skip_blockhash_check(true);
465        assert!(builder.skip_blockhash_check);
466    }
467}