sui_node/
handle.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! SuiNodeHandle wraps SuiNode in a way suitable for access by test code.
5//!
6//! When starting a SuiNode directly, in a test (as opposed to using Swarm), the node may be
7//! running inside of a simulator node. It is therefore a mistake to do something like:
8//!
9//! ```ignore
10//!     use test_utils::authority::{start_node, spawn_checkpoint_processes};
11//!
12//!     let node = start_node(config, registry).await;
13//!     spawn_checkpoint_processes(config, &[node]).await;
14//! ```
15//!
16//! Because this would cause the checkpointing processes to be running inside the current
17//! simulator node rather than the node in which the SuiNode is running.
18//!
19//! SuiNodeHandle provides an easy way to do the right thing here:
20//!
21//! ```ignore
22//!     let node_handle = start_node(config, registry).await;
23//!     node_handle.with_async(|sui_node| async move {
24//!         spawn_checkpoint_processes(config, &[sui_node]).await;
25//!     });
26//! ```
27//!
28//! Code executed inside of with or with_async will run in the context of the simulator node.
29//! This allows tests to break the simulator abstraction and magically mutate or inspect state that
30//! is conceptually running on a different "machine", but without producing extremely confusing
31//! behavior that might result otherwise. (For instance, any network connection that is initiated
32//! from a task spawned from within a with or with_async will appear to originate from the correct
33//! simulator node.
34//!
35//! It is possible to exfiltrate state:
36//!
37//! ```ignore
38//!    let state = node_handle.with(|sui_node| sui_node.state);
39//!    // DO NOT DO THIS!
40//!    do_stuff_with_state(state)
41//! ```
42//!
43//! We can't prevent this completely, but we can at least make the right way the easy way.
44
45use super::SuiNode;
46use std::future::Future;
47use std::sync::Arc;
48use sui_core::authority::AuthorityState;
49
50/// Wrap SuiNode to allow correct access to SuiNode in simulator tests.
51pub struct SuiNodeHandle {
52    node: Option<Arc<SuiNode>>,
53    shutdown_on_drop: bool,
54}
55
56impl SuiNodeHandle {
57    pub fn new(node: Arc<SuiNode>) -> Self {
58        Self {
59            node: Some(node),
60            shutdown_on_drop: false,
61        }
62    }
63
64    pub fn inner(&self) -> &Arc<SuiNode> {
65        self.node.as_ref().unwrap()
66    }
67
68    pub fn with<T>(&self, cb: impl FnOnce(&SuiNode) -> T) -> T {
69        let _guard = self.guard();
70        cb(self.inner())
71    }
72
73    pub fn state(&self) -> Arc<AuthorityState> {
74        self.with(|sui_node| sui_node.state())
75    }
76
77    pub fn shutdown_on_drop(&mut self) {
78        self.shutdown_on_drop = true;
79    }
80}
81
82impl Clone for SuiNodeHandle {
83    fn clone(&self) -> Self {
84        Self {
85            node: self.node.clone(),
86            shutdown_on_drop: false,
87        }
88    }
89}
90
91#[cfg(not(msim))]
92impl SuiNodeHandle {
93    // Must return something to silence lints above at `let _guard = ...`
94    fn guard(&self) -> u32 {
95        0
96    }
97
98    pub async fn with_async<'a, F, R, T>(&'a self, cb: F) -> T
99    where
100        F: FnOnce(&'a SuiNode) -> R,
101        R: Future<Output = T>,
102    {
103        cb(self.inner()).await
104    }
105}
106
107#[cfg(msim)]
108impl SuiNodeHandle {
109    fn guard(&self) -> sui_simulator::runtime::NodeEnterGuard {
110        self.inner().sim_state.sim_node.enter_node()
111    }
112
113    pub async fn with_async<'a, F, R, T>(&'a self, cb: F) -> T
114    where
115        F: FnOnce(&'a SuiNode) -> R,
116        R: Future<Output = T>,
117    {
118        let fut = cb(self.node.as_ref().unwrap());
119        self.inner()
120            .sim_state
121            .sim_node
122            .await_future_in_node(fut)
123            .await
124    }
125}
126
127#[cfg(msim)]
128impl Drop for SuiNodeHandle {
129    fn drop(&mut self) {
130        if self.shutdown_on_drop {
131            let node_id = self.inner().sim_state.sim_node.id();
132            sui_simulator::runtime::Handle::try_current().map(|h| h.delete_node(node_id));
133        }
134    }
135}
136
137impl From<Arc<SuiNode>> for SuiNodeHandle {
138    fn from(node: Arc<SuiNode>) -> Self {
139        SuiNodeHandle::new(node)
140    }
141}
sui_node/handle.rs

sui_node/
handle.rs
