iota_config/
local_ip_utils.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::net::SocketAddr;
6#[cfg(msim)]
7use std::sync::{Arc, atomic::AtomicI16};
8
9use iota_types::multiaddr::Multiaddr;
10
11/// A singleton struct to manage IP addresses and ports for simtest.
12/// This allows us to generate unique IP addresses and ports for each node in
13/// simtest.
14#[cfg(msim)]
15pub struct SimAddressManager {
16    next_ip_offset: AtomicI16,
17    next_port: AtomicI16,
18}
19
20#[cfg(msim)]
21impl Default for SimAddressManager {
22    fn default() -> Self {
23        Self {
24            next_ip_offset: AtomicI16::new(1),
25            next_port: AtomicI16::new(9000),
26        }
27    }
28}
29
30#[cfg(msim)]
31impl SimAddressManager {
32    pub fn new() -> Self {
33        Self::default()
34    }
35
36    pub fn get_next_ip(&self) -> String {
37        let offset = self
38            .next_ip_offset
39            .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
40        // If offset ever goes beyond 255, we could use more bytes in the IP.
41        assert!(offset <= 255);
42        format!("10.10.0.{}", offset)
43    }
44
45    pub fn get_next_available_port(&self) -> u16 {
46        self.next_port
47            .fetch_add(1, std::sync::atomic::Ordering::SeqCst) as u16
48    }
49}
50
51#[cfg(msim)]
52fn get_sim_address_manager() -> Arc<SimAddressManager> {
53    thread_local! {
54        // Uses Arc so that we could return a clone of the thread local singleton.
55        static SIM_ADDRESS_MANAGER: Arc<SimAddressManager> = Arc::new(SimAddressManager::new());
56    }
57    SIM_ADDRESS_MANAGER.with(|s| s.clone())
58}
59
60/// In simtest, we generate a new unique IP each time this function is called.
61#[cfg(msim)]
62pub fn get_new_ip() -> String {
63    get_sim_address_manager().get_next_ip()
64}
65
66/// In non-simtest, we always only have one IP address which is localhost.
67#[cfg(not(msim))]
68pub fn get_new_ip() -> String {
69    localhost_for_testing()
70}
71
72/// Returns localhost, which is always 127.0.0.1.
73pub fn localhost_for_testing() -> String {
74    "127.0.0.1".to_string()
75}
76
77/// Returns an available port for the given host in simtest.
78/// We don't care about host because it's all managed by simulator. Just obtain
79/// a unique port.
80#[cfg(msim)]
81pub fn get_available_port(_host: &str) -> u16 {
82    get_sim_address_manager().get_next_available_port()
83}
84
85/// Return an ephemeral, available port. On unix systems, the port returned will
86/// be in the TIME_WAIT state ensuring that the OS won't hand out this port for
87/// some grace period. Callers should be able to bind to this port given they
88/// use SO_REUSEADDR.
89#[cfg(not(msim))]
90pub fn get_available_port(host: &str) -> u16 {
91    const MAX_PORT_RETRIES: u32 = 1000;
92
93    for _ in 0..MAX_PORT_RETRIES {
94        if let Ok(port) = get_ephemeral_port(host) {
95            return port;
96        }
97    }
98
99    panic!(
100        "error: could not find an available port on {}: {:?}",
101        host,
102        get_ephemeral_port(host)
103    );
104}
105
106#[cfg(not(msim))]
107fn get_ephemeral_port(host: &str) -> std::io::Result<u16> {
108    use std::net::{TcpListener, TcpStream};
109
110    // Request a random available port from the OS
111    let listener = TcpListener::bind((host, 0))?;
112    let addr = listener.local_addr()?;
113
114    // Create and accept a connection (which we'll promptly drop) in order to force
115    // the port into the TIME_WAIT state, ensuring that the port will be
116    // reserved from some limited amount of time (roughly 60s on some Linux
117    // systems)
118    let _sender = TcpStream::connect(addr)?;
119    let _incoming = listener.accept()?;
120
121    Ok(addr.port())
122}
123
124/// Returns a new unique TCP address for the given host, by finding a new
125/// available port.
126pub fn new_tcp_address_for_testing(host: &str) -> Multiaddr {
127    format!("/ip4/{}/tcp/{}/http", host, get_available_port(host))
128        .parse()
129        .unwrap()
130}
131
132/// Returns a new unique UDP address for the given host, by finding a new
133/// available port.
134pub fn new_udp_address_for_testing(host: &str) -> Multiaddr {
135    format!("/ip4/{}/udp/{}", host, get_available_port(host))
136        .parse()
137        .unwrap()
138}
139
140/// Returns a new unique TCP address (SocketAddr) for localhost, by finding a
141/// new available port on localhost.
142pub fn new_local_tcp_socket_for_testing() -> SocketAddr {
143    format!(
144        "{}:{}",
145        localhost_for_testing(),
146        get_available_port(&localhost_for_testing())
147    )
148    .parse()
149    .unwrap()
150}
151
152/// Returns a new unique TCP address (Multiaddr) for localhost, by finding a new
153/// available port on localhost.
154pub fn new_local_tcp_address_for_testing() -> Multiaddr {
155    new_tcp_address_for_testing(&localhost_for_testing())
156}
157
158/// Returns a new unique UDP address for localhost, by finding a new available
159/// port.
160pub fn new_local_udp_address_for_testing() -> Multiaddr {
161    new_udp_address_for_testing(&localhost_for_testing())
162}
163
164pub fn new_deterministic_tcp_address_for_testing(host: &str, port: u16) -> Multiaddr {
165    format!("/ip4/{host}/tcp/{port}/http").parse().unwrap()
166}
167
168pub fn new_deterministic_udp_address_for_testing(host: &str, port: u16) -> Multiaddr {
169    format!("/ip4/{host}/udp/{port}/http").parse().unwrap()
170}
iota_config/local_ip_utils.rs

iota_config/
local_ip_utils.rs
