sui_config/
local_ip_utils.rs

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

sui_config/
local_ip_utils.rs
