iota_node/
handle.rs

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

iota_node/
handle.rs
