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
35pub(crate) const TIMEOUT: Duration = Duration::from_secs(5);
37
38pub(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
52pub(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
63async 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 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 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 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
141pub(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 let bind_address = listeners[0]
158 .local_addr()
159 .expect("Sockets should be bound")
160 .ip();
161
162 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 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 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 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 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 let _ = sender.send(());
206 }
207 None => warn!("tcp incoming failed"),
208 }
209 });
210
211 let receiver = tokio::time::timeout(timeout, receiver);
213 checkers.push((listening_addr, thread_handle, receiver));
214 }
215 }
216
217 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 TcpStream::connect_timeout(&listening_addr, timeout).unwrap();
236 ok = false;
238 }
239 }
240 thread_handle.await.expect("Thread should exit cleanly");
241 }
242
243 ok
244}
245
246pub(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 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 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 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; }
354 }
355
356 error!("Maximum retry count reached. Some ports for IP {bind_ip} unreachable.");
357 return false;
358 }
359 }
360 true
361}