Skip to main content

solana_net_utils/
ip_echo_client.rs

1use {
2    crate::{
3        HEADER_LENGTH, IP_ECHO_SERVER_RESPONSE_LENGTH, MAX_PORT_COUNT_PER_MESSAGE,
4        ip_echo_server::{IpEchoServerMessage, IpEchoServerResponse},
5    },
6    bytes::{BufMut, BytesMut},
7    itertools::Itertools,
8    log::*,
9    std::{
10        collections::{BTreeMap, HashMap, HashSet},
11        net::{IpAddr, SocketAddr, TcpListener, TcpStream, UdpSocket},
12        sync::{Arc, RwLock},
13        time::{Duration, Instant},
14    },
15    tokio::{
16        io::{AsyncReadExt, AsyncWriteExt},
17        net::TcpSocket,
18        sync::oneshot,
19        task::JoinSet,
20    },
21};
22
23#[derive(Debug, thiserror::Error)]
24pub enum IpEchoClientError {
25    #[error(transparent)]
26    Io(#[from] std::io::Error),
27    #[error(transparent)]
28    Timeout(#[from] tokio::time::error::Elapsed),
29    #[error(transparent)]
30    BincodeError(#[from] bincode::Error),
31    #[error("{0}")]
32    InvalidResponse(String),
33}
34
35/// Applies to all operations with the echo server.
36pub(crate) const TIMEOUT: Duration = Duration::from_secs(5);
37
38/// Make a request to the echo server, binding the client socket to the provided IP.
39pub(crate) async fn ip_echo_server_request_with_binding(
40    ip_echo_server_addr: SocketAddr,
41    msg: IpEchoServerMessage,
42    bind_address: IpAddr,
43) -> Result<IpEchoServerResponse, IpEchoClientError> {
44    let socket = TcpSocket::new_v4()?;
45    socket.bind(SocketAddr::new(bind_address, 0))?;
46
47    let response =
48        tokio::time::timeout(TIMEOUT, make_request(socket, ip_echo_server_addr, msg)).await??;
49    parse_response(response, ip_echo_server_addr)
50}
51
52/// Make a request to the echo server, client socket will be bound by the OS.
53pub(crate) async fn ip_echo_server_request(
54    ip_echo_server_addr: SocketAddr,
55    msg: IpEchoServerMessage,
56) -> Result<IpEchoServerResponse, IpEchoClientError> {
57    let socket = TcpSocket::new_v4()?;
58    let response =
59        tokio::time::timeout(TIMEOUT, make_request(socket, ip_echo_server_addr, msg)).await??;
60    parse_response(response, ip_echo_server_addr)
61}
62
63/// Makes the request to the specified server and returns reply as Bytes.
64async fn make_request(
65    socket: TcpSocket,
66    ip_echo_server_addr: SocketAddr,
67    msg: IpEchoServerMessage,
68) -> Result<BytesMut, IpEchoClientError> {
69    let mut stream = socket.connect(ip_echo_server_addr).await?;
70    // Start with HEADER_LENGTH null bytes to avoid looking like an HTTP GET/POST request
71    let mut bytes = BytesMut::with_capacity(IP_ECHO_SERVER_RESPONSE_LENGTH);
72    bytes.extend_from_slice(&[0u8; HEADER_LENGTH]);
73    bytes.extend_from_slice(&bincode::serialize(&msg)?);
74
75    // End with '\n' to make this request look HTTP-ish and tickle an error response back
76    // from an HTTP server
77    bytes.put_u8(b'\n');
78    stream.write_all(&bytes).await?;
79    stream.flush().await?;
80
81    bytes.clear();
82    let _n = stream.read_buf(&mut bytes).await?;
83    stream.shutdown().await?;
84
85    Ok(bytes)
86}
87
88fn parse_response(
89    response: BytesMut,
90    ip_echo_server_addr: SocketAddr,
91) -> Result<IpEchoServerResponse, IpEchoClientError> {
92    // It's common for users to accidentally confuse the validator's gossip port and JSON
93    // RPC port.  Attempt to detect when this occurs by looking for the standard HTTP
94    // response header and provide the user with a helpful error message
95    if response.len() < HEADER_LENGTH {
96        return Err(IpEchoClientError::InvalidResponse(format!(
97            "Response too short, received {} bytes",
98            response.len()
99        )));
100    }
101
102    let (response_header, body) =
103        response
104            .split_first_chunk::<HEADER_LENGTH>()
105            .ok_or_else(|| {
106                IpEchoClientError::InvalidResponse(format!(
107                    "Not enough data in the response from {ip_echo_server_addr}!"
108                ))
109            })?;
110    let payload = match response_header {
111        [0, 0, 0, 0] => bincode::deserialize(body)?,
112        [b'H', b'T', b'T', b'P'] => {
113            let http_response = std::str::from_utf8(body);
114            match http_response {
115                Ok(r) => {
116                    return Err(IpEchoClientError::InvalidResponse(format!(
117                        "Invalid gossip entrypoint. {ip_echo_server_addr} looks to be an HTTP \
118                         port replying with {r}"
119                    )));
120                }
121                Err(_) => {
122                    return Err(IpEchoClientError::InvalidResponse(format!(
123                        "Invalid gossip entrypoint. {ip_echo_server_addr} looks to be an HTTP \
124                         port."
125                    )));
126                }
127            }
128        }
129        _ => {
130            return Err(IpEchoClientError::InvalidResponse(format!(
131                "Invalid gossip entrypoint. {ip_echo_server_addr} provided unexpected header \
132                 bytes {response_header:?} "
133            )));
134        }
135    };
136    Ok(payload)
137}
138
139pub(crate) const DEFAULT_RETRY_COUNT: usize = 5;
140
141/// Checks if all of the provided TCP ports are reachable by the machine at
142/// `ip_echo_server_addr`. Tests must complete within timeout provided.
143/// Tests will run concurrently to avoid head-of-line blocking. Will return false on any error.
144/// This function may panic.
145/// All listening sockets should be bound to the same IP.
146pub(crate) async fn verify_all_reachable_tcp(
147    ip_echo_server_addr: SocketAddr,
148    listeners: Vec<TcpListener>,
149    timeout: Duration,
150) -> bool {
151    if listeners.is_empty() {
152        warn!("No ports provided for verify_all_reachable_tcp to check");
153        return true;
154    }
155
156    // Extract the bind address for requests to remote server
157    let bind_address = listeners[0]
158        .local_addr()
159        .expect("Sockets should be bound")
160        .ip();
161
162    // Verify that all other sockets are bound to the same address
163    for listener in listeners.iter() {
164        let local_binding = listener.local_addr().expect("Sockets should be bound");
165        assert_eq!(
166            local_binding.ip(),
167            bind_address,
168            "All sockets should be bound to the same IP"
169        );
170    }
171    let mut checkers = Vec::new();
172    let mut ok = true;
173
174    // Chunk the port range into slices small enough to fit into one packet
175    for chunk in &listeners.into_iter().chunks(MAX_PORT_COUNT_PER_MESSAGE) {
176        let listeners = chunk.collect_vec();
177        let ports = listeners
178            .iter()
179            .map(|l| l.local_addr().expect("Sockets should be bound").port())
180            .collect_vec();
181        info!("Checking that tcp ports {ports:?} are reachable from {ip_echo_server_addr:?}");
182
183        // make request to the echo server
184        let _ = ip_echo_server_request_with_binding(
185            ip_echo_server_addr,
186            IpEchoServerMessage::new(&ports, &[]),
187            bind_address,
188        )
189        .await
190        .map_err(|err| warn!("ip_echo_server request failed: {err}"));
191
192        // spawn checker to wait for reply
193        // since we do not know if tcp_listeners are nonblocking, we have to run them in native threads.
194        for (port, tcp_listener) in ports.into_iter().zip(listeners) {
195            let listening_addr = tcp_listener.local_addr().unwrap();
196            let (sender, receiver) = oneshot::channel();
197
198            // Use blocking API since we have no idea if sockets given to us are nonblocking or not
199            let thread_handle = tokio::task::spawn_blocking(move || {
200                debug!("Waiting for incoming connection on tcp/{port}");
201                match tcp_listener.incoming().next() {
202                    Some(_) => {
203                        // ignore errors here since this can only happen if a timeout was detected.
204                        // timeout drops the receiver part of the channel resulting in failure to send.
205                        let _ = sender.send(());
206                    }
207                    None => warn!("tcp incoming failed"),
208                }
209            });
210
211            // Set the timeout on the receiver
212            let receiver = tokio::time::timeout(timeout, receiver);
213            checkers.push((listening_addr, thread_handle, receiver));
214        }
215    }
216
217    // now wait for notifications from all the tasks we have spawned.
218    for (listening_addr, thread_handle, receiver) in checkers.drain(..) {
219        match receiver.await {
220            Ok(Ok(_)) => {
221                info!("tcp/{} is reachable", listening_addr.port());
222            }
223            Ok(Err(_v)) => {
224                unreachable!("The receive on oneshot channel should never fail");
225            }
226            Err(_t) => {
227                error!(
228                    "Received no response at tcp/{}, check your port configuration",
229                    listening_addr.port()
230                );
231                // Ugh, std rustc doesn't provide accepting with timeout or restoring original
232                // nonblocking-status of sockets because of lack of getter, only the setter...
233                // So, to close the thread cleanly, just connect from here.
234                // ref: https://github.com/rust-lang/rust/issues/31615
235                TcpStream::connect_timeout(&listening_addr, timeout).unwrap();
236                // Mark that we have found error, but do  not exit yet, as we will have stuck ports otherwise.
237                ok = false;
238            }
239        }
240        thread_handle.await.expect("Thread should exit cleanly");
241    }
242
243    ok
244}
245
246/// Checks if all of the provided UDP ports on all of the provided IPs are
247/// reachable by the machine at `ip_echo_server_addr`.
248/// This function will test a few ports at a time, retrying if necessary.
249/// Tests must complete within timeout provided, so a longer timeout may be
250/// necessary if checking many ports.
251/// A given amount of retries will be made to accommodate packet loss.
252/// This function may panic.
253///
254pub(crate) async fn verify_all_reachable_udp(
255    ip_echo_server_addr: SocketAddr,
256    sockets: &[&UdpSocket],
257    timeout: Duration,
258    retry_count: usize,
259) -> bool {
260    if sockets.is_empty() {
261        warn!("No ports provided for verify_all_reachable_udp to check");
262        return true;
263    }
264    let mut ip_to_ports: HashMap<IpAddr, BTreeMap<u16, Vec<&UdpSocket>>> = HashMap::new();
265    for &socket in sockets.iter() {
266        let local_addr = socket.local_addr().expect("Socket must be bound");
267        ip_to_ports
268            .entry(local_addr.ip())
269            .or_default()
270            .entry(local_addr.port())
271            .or_default()
272            .push(socket);
273    }
274    for (bind_ip, ports_to_socks_map) in ip_to_ports {
275        let ports: Vec<u16> = ports_to_socks_map.keys().copied().collect();
276
277        info!("Checking that udp ports {ports:?} are reachable from bind IP {bind_ip:?}");
278
279        'outer: for chunk_to_check in ports.chunks(MAX_PORT_COUNT_PER_MESSAGE) {
280            let ports_to_check = chunk_to_check.to_vec();
281
282            for attempt in 0..retry_count {
283                if attempt > 0 {
284                    error!("There are some udp ports with no response!! Retrying...");
285                }
286                // clone off the sockets that use ports within our chunk
287                let sockets_to_check: Vec<UdpSocket> = ports_to_check
288                    .iter()
289                    .flat_map(|port| ports_to_socks_map.get(port).unwrap())
290                    .map(|&s| s.try_clone().expect("Unable to clone UDP socket"))
291                    .collect();
292
293                let _ = ip_echo_server_request_with_binding(
294                    ip_echo_server_addr,
295                    IpEchoServerMessage::new(&[], &ports_to_check),
296                    bind_ip,
297                )
298                .await
299                .map_err(|err| warn!("ip_echo_server request failed: {err}"));
300
301                let reachable_ports = Arc::new(RwLock::new(HashSet::new()));
302                // Spawn threads for each socket to check
303                let mut checkers = JoinSet::new();
304                for socket in sockets_to_check {
305                    let port = socket.local_addr().expect("Socket should be bound").port();
306                    let reachable_ports = reachable_ports.clone();
307
308                    checkers.spawn_blocking(move || {
309                        let start = Instant::now();
310
311                        let original_read_timeout = socket.read_timeout().unwrap();
312                        socket
313                            .set_read_timeout(Some(Duration::from_millis(250)))
314                            .unwrap();
315
316                        loop {
317                            if reachable_ports.read().unwrap().contains(&port)
318                                || Instant::now().duration_since(start) >= timeout
319                            {
320                                break;
321                            }
322
323                            let recv_result = socket.recv(&mut [0; 1]);
324                            debug!("Waited for incoming datagram on udp/{port}: {recv_result:?}");
325
326                            if recv_result.is_ok() {
327                                reachable_ports.write().unwrap().insert(port);
328                                break;
329                            }
330                        }
331
332                        socket.set_read_timeout(original_read_timeout).unwrap();
333                    });
334                }
335                loop {
336                    let next = checkers.join_next().await;
337                    let Some(r) = next else {
338                        break;
339                    };
340                    r.expect("Threads should exit cleanly");
341                }
342                // Might have lost a UDP packet, check that all ports were reached
343                let reachable_ports = Arc::into_inner(reachable_ports)
344                    .expect("Single owner expected")
345                    .into_inner()
346                    .expect("No threads should hold the lock");
347                info!(
348                    "checked udp ports: {ports_to_check:?}, reachable udp ports: \
349                     {reachable_ports:?}"
350                );
351                if reachable_ports.len() == ports_to_check.len() {
352                    continue 'outer; // starts checking next chunk of ports, if any
353                }
354            }
355
356            error!("Maximum retry count reached. Some ports for IP {bind_ip} unreachable.");
357            return false;
358        }
359    }
360    true
361}