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    // Uses Arc so that we could return a clone of the process-global singleton.
44    static SIM_ADDRESS_MANAGER: std::sync::OnceLock<Arc<SimAddressManager>> =
45        std::sync::OnceLock::new();
46    SIM_ADDRESS_MANAGER
47        .get_or_init(|| Arc::new(SimAddressManager::new()))
48        .clone()
49}
50
51/// In simtest, we generate a new unique IP each time this function is called.
52#[cfg(msim)]
53pub fn get_new_ip() -> String {
54    get_sim_address_manager().get_next_ip()
55}
56
57/// In non-simtest, we always only have one IP address which is localhost.
58#[cfg(not(msim))]
59pub fn get_new_ip() -> String {
60    localhost_for_testing()
61}
62
63/// Returns localhost, which is always 127.0.0.1.
64pub fn localhost_for_testing() -> String {
65    "127.0.0.1".to_string()
66}
67
68/// Returns an available port for the given host in simtest.
69/// We don't care about host because it's all managed by simulator. Just obtain a unique port.
70#[cfg(msim)]
71pub fn get_available_port(_host: &str) -> u16 {
72    get_sim_address_manager().get_next_available_port()
73}
74
75/// Return an ephemeral, available port. On unix systems, the port returned will be in the
76/// TIME_WAIT state ensuring that the OS won't hand out this port for some grace period.
77/// Callers should be able to bind to this port given they use SO_REUSEADDR.
78#[cfg(not(msim))]
79pub fn get_available_port(host: &str) -> u16 {
80    const MAX_PORT_RETRIES: u32 = 1000;
81
82    for _ in 0..MAX_PORT_RETRIES {
83        if let Ok(port) = get_ephemeral_port(host) {
84            return port;
85        }
86    }
87
88    panic!(
89        "Error: could not find an available port on {}: {:?}",
90        host,
91        get_ephemeral_port(host)
92    );
93}
94
95#[cfg(not(msim))]
96fn get_ephemeral_port(host: &str) -> std::io::Result<u16> {
97    use std::net::{TcpListener, TcpStream};
98
99    // Request a random available port from the OS
100    let listener = TcpListener::bind((host, 0))?;
101    let addr = listener.local_addr()?;
102
103    // Create and accept a connection (which we'll promptly drop) in order to force the port
104    // into the TIME_WAIT state, ensuring that the port will be reserved from some limited
105    // amount of time (roughly 60s on some Linux systems)
106    let _sender = TcpStream::connect(addr)?;
107    let _incoming = listener.accept()?;
108
109    Ok(addr.port())
110}
111
112/// Returns a new unique TCP address for the given host, by finding a new available port.
113pub fn new_tcp_address_for_testing(host: &str) -> Multiaddr {
114    format!("/ip4/{}/tcp/{}/https", host, get_available_port(host))
115        .parse()
116        .unwrap()
117}
118
119/// Returns a new unique UDP address for the given host, by finding a new available port.
120pub fn new_udp_address_for_testing(host: &str) -> Multiaddr {
121    format!("/ip4/{}/udp/{}", host, get_available_port(host))
122        .parse()
123        .unwrap()
124}
125
126/// Returns a new unique TCP address in String format for localhost, by finding a new available port on localhost.
127pub fn new_local_tcp_socket_for_testing_string() -> String {
128    format!(
129        "{}:{}",
130        localhost_for_testing(),
131        get_available_port(&localhost_for_testing())
132    )
133}
134
135/// Returns a new unique TCP address (SocketAddr) for localhost, by finding a new available port on localhost.
136pub fn new_local_tcp_socket_for_testing() -> SocketAddr {
137    new_local_tcp_socket_for_testing_string().parse().unwrap()
138}
139
140/// Returns a new unique TCP address (Multiaddr) for localhost, by finding a new available port on localhost.
141pub fn new_local_tcp_address_for_testing() -> Multiaddr {
142    new_tcp_address_for_testing(&localhost_for_testing())
143}
144
145/// Returns a new unique UDP address for localhost, by finding a new available port.
146pub fn new_local_udp_address_for_testing() -> Multiaddr {
147    new_udp_address_for_testing(&localhost_for_testing())
148}
149
150pub fn new_deterministic_tcp_address_for_testing(host: &str, port: u16) -> Multiaddr {
151    format!("/ip4/{host}/tcp/{port}/https").parse().unwrap()
152}
153
154pub fn new_deterministic_udp_address_for_testing(host: &str, port: u16) -> Multiaddr {
155    format!("/ip4/{host}/udp/{port}/https").parse().unwrap()
156}
sui_config/local_ip_utils.rs

sui_config/
local_ip_utils.rs
