1#![cfg(feature = "agave-unstable-api")]
2#![warn(if_let_rescope)]
6#![warn(keyword_idents_2024)]
7#![warn(rust_2024_incompatible_pat)]
8#![warn(tail_expr_drop_order)]
9#![warn(unsafe_attr_outside_unsafe)]
10#![warn(unsafe_op_in_unsafe_fn)]
11
12pub mod banlist;
13mod ip_echo_client;
14mod ip_echo_server;
15pub mod multihomed_sockets;
16pub mod socket_addr_space;
17pub mod sockets;
18#[cfg(any(target_os = "android", target_os = "windows"))]
19#[path = "test_port_allocator_legacy.rs"]
20pub(crate) mod test_port_allocator;
21#[cfg(not(any(target_os = "android", target_os = "windows")))]
22pub(crate) mod test_port_allocator;
23pub mod token_bucket;
24
25#[cfg(feature = "dev-context-only-utils")]
26pub mod tooling_for_tests;
27
28pub use {
29 ip_echo_client::IpEchoClientError,
30 ip_echo_server::{
31 DEFAULT_IP_ECHO_SERVER_THREADS, IpEchoServer, MAX_PORT_COUNT_PER_MESSAGE, ip_echo_server,
32 },
33 socket_addr_space::SocketAddrSpace,
34};
35use {
36 ip_echo_client::{ip_echo_server_request, ip_echo_server_request_with_binding},
37 ip_echo_server::IpEchoServerMessage,
38 rand::{Rng, rng},
39 std::{
40 io::{self},
41 net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener, ToSocketAddrs, UdpSocket},
42 },
43 url::Url,
44};
45
46pub struct UdpSocketPair {
48 pub addr: SocketAddr, pub receiver: UdpSocket, pub sender: UdpSocket, }
52
53pub type PortRange = (u16, u16);
54
55#[cfg(not(debug_assertions))]
56pub const VALIDATOR_PORT_RANGE: PortRange = (8000, 10_000);
58
59#[cfg(debug_assertions)]
62pub const VALIDATOR_PORT_RANGE: PortRange = (
63 crate::sockets::UNIQUE_ALLOC_BASE_PORT - 512,
64 crate::sockets::UNIQUE_ALLOC_BASE_PORT,
65);
66
67pub const MINIMUM_VALIDATOR_PORT_RANGE_WIDTH: u16 = 26; pub(crate) const HEADER_LENGTH: usize = 4;
70pub(crate) const IP_ECHO_SERVER_RESPONSE_LENGTH: usize = HEADER_LENGTH + 23;
71
72pub fn get_public_ip_addr_with_binding(
75 ip_echo_server_addr: &SocketAddr,
76 bind_address: IpAddr,
77) -> Result<IpAddr, IpEchoClientError> {
78 let fut = ip_echo_server_request_with_binding(
79 *ip_echo_server_addr,
80 IpEchoServerMessage::default(),
81 bind_address,
82 );
83 let rt = tokio::runtime::Builder::new_current_thread()
84 .enable_all()
85 .build()?;
86 let resp = rt.block_on(fut)?;
87 Ok(resp.address)
88}
89
90pub fn get_cluster_shred_version(ip_echo_server_addr: &SocketAddr) -> Result<u16, String> {
92 let fut = ip_echo_server_request(*ip_echo_server_addr, IpEchoServerMessage::default());
93 let rt = tokio::runtime::Builder::new_current_thread()
94 .enable_all()
95 .build()
96 .map_err(|e| e.to_string())?;
97 let resp = rt.block_on(fut).map_err(|e| e.to_string())?;
98 resp.shred_version
99 .ok_or_else(|| "IP echo server does not return a shred-version".to_owned())
100}
101
102pub fn get_cluster_shred_version_with_binding(
105 ip_echo_server_addr: &SocketAddr,
106 bind_address: IpAddr,
107) -> Result<u16, IpEchoClientError> {
108 let fut = ip_echo_server_request_with_binding(
109 *ip_echo_server_addr,
110 IpEchoServerMessage::default(),
111 bind_address,
112 );
113 let rt = tokio::runtime::Builder::new_current_thread()
114 .enable_all()
115 .build()?;
116 let resp = rt.block_on(fut)?;
117 resp.shred_version.ok_or_else(|| {
118 IpEchoClientError::InvalidResponse(
119 "IP echo server does not return a shred-version".to_owned(),
120 )
121 })
122}
123
124const MAX_PORT_VERIFY_THREADS: usize = 64;
127
128pub fn verify_all_reachable_udp(
133 ip_echo_server_addr: &SocketAddr,
134 udp_sockets: &[&UdpSocket],
135) -> bool {
136 let rt = tokio::runtime::Builder::new_current_thread()
137 .enable_all()
138 .max_blocking_threads(MAX_PORT_VERIFY_THREADS)
139 .build()
140 .expect("Tokio builder should be able to reliably create a current thread runtime");
141 let fut = ip_echo_client::verify_all_reachable_udp(
142 *ip_echo_server_addr,
143 udp_sockets,
144 ip_echo_client::TIMEOUT,
145 ip_echo_client::DEFAULT_RETRY_COUNT,
146 );
147 rt.block_on(fut)
148}
149
150pub fn verify_all_reachable_tcp(
155 ip_echo_server_addr: &SocketAddr,
156 tcp_listeners: Vec<TcpListener>,
157) -> bool {
158 let rt = tokio::runtime::Builder::new_current_thread()
159 .enable_all()
160 .max_blocking_threads(MAX_PORT_VERIFY_THREADS)
161 .build()
162 .expect("Tokio builder should be able to reliably create a current thread runtime");
163 let fut = ip_echo_client::verify_all_reachable_tcp(
164 *ip_echo_server_addr,
165 tcp_listeners,
166 ip_echo_client::TIMEOUT,
167 );
168 rt.block_on(fut)
169}
170
171pub fn parse_port_or_addr(optstr: Option<&str>, default_addr: SocketAddr) -> SocketAddr {
172 if let Some(addrstr) = optstr {
173 if let Ok(port) = addrstr.parse() {
174 let mut addr = default_addr;
175 addr.set_port(port);
176 addr
177 } else if let Ok(addr) = addrstr.parse() {
178 addr
179 } else {
180 default_addr
181 }
182 } else {
183 default_addr
184 }
185}
186
187pub fn parse_port_range(port_range: &str) -> Option<PortRange> {
188 let ports: Vec<&str> = port_range.split('-').collect();
189 if ports.len() != 2 {
190 return None;
191 }
192
193 let start_port = ports[0].parse();
194 let end_port = ports[1].parse();
195
196 if start_port.is_err() || end_port.is_err() {
197 return None;
198 }
199 let start_port = start_port.unwrap();
200 let end_port = end_port.unwrap();
201 if end_port < start_port {
202 return None;
203 }
204 Some((start_port, end_port))
205}
206
207fn select_ipv4(host: &str, mut ips: impl Iterator<Item = IpAddr>) -> Result<IpAddr, String> {
208 let Some(first_ip) = ips.next() else {
209 return Err(format!("Unable to resolve host: {host}"));
210 };
211
212 if first_ip.is_ipv4() {
213 return Ok(first_ip);
214 }
215
216 ips.find(IpAddr::is_ipv4)
217 .ok_or_else(|| format!("IPv6 addresses are not supported: {host}"))
218}
219
220pub fn parse_host(host: &str) -> Result<IpAddr, String> {
221 if let Ok(IpAddr::V6(_)) = host.parse::<IpAddr>() {
222 return Err(format!("IPv6 addresses are not supported: {host}"));
223 }
224
225 let parsed_url = Url::parse(&format!("http://{host}")).map_err(|e| e.to_string())?;
228 if parsed_url.port().is_some() {
229 return Err(format!("Expected port in URL: {host}"));
230 }
231
232 let mut ips = (host, 0)
234 .to_socket_addrs()
235 .map_err(|err| err.to_string())?
236 .map(|socket_address| socket_address.ip());
237
238 select_ipv4(host, &mut ips)
239}
240
241pub fn is_host(string: String) -> Result<(), String> {
242 parse_host(&string).map(|_| ())
243}
244
245pub fn parse_host_port(host_port: &str) -> Result<SocketAddr, String> {
246 let addrs: Vec<_> = host_port
247 .to_socket_addrs()
248 .map_err(|err| format!("Unable to resolve host {host_port}: {err}"))?
249 .collect();
250 if addrs.is_empty() {
251 Err(format!("Unable to resolve host: {host_port}"))
252 } else {
253 Ok(addrs[0])
254 }
255}
256
257pub fn is_host_port(string: String) -> Result<(), String> {
258 parse_host_port(&string).map(|_| ())
259}
260
261pub fn bind_in_range(ip_addr: IpAddr, range: PortRange) -> io::Result<(u16, UdpSocket)> {
262 let config = sockets::SocketConfiguration::default();
263 sockets::bind_in_range_with_config(ip_addr, range, config)
264}
265
266pub fn bind_to_unspecified() -> io::Result<UdpSocket> {
267 let config = sockets::SocketConfiguration::default();
268 sockets::bind_to_with_config(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0, config)
269}
270
271pub fn find_available_port_in_range(ip_addr: IpAddr, range: PortRange) -> io::Result<u16> {
278 let [port] = find_available_ports_in_range(ip_addr, range)?;
279 Ok(port)
280}
281
282pub fn find_available_ports_in_range<const N: usize>(
287 ip_addr: IpAddr,
288 range: PortRange,
289) -> io::Result<[u16; N]> {
290 let mut result = [0u16; N];
291 let range = range.0..range.1;
292 let mut next_port_to_try = range
293 .clone()
294 .cycle() .skip(rng().random_range(range.clone()) as usize) .take(range.len()) .peekable();
298 let mut num = 0;
299 let config = sockets::SocketConfiguration::default();
300 while num < N {
301 let port_to_try = next_port_to_try.next().unwrap(); let bind = sockets::bind_common_with_config(ip_addr, port_to_try, config);
303 match bind {
304 Ok(_) => {
305 result[num] = port_to_try;
306 num = num.saturating_add(1);
307 }
308 Err(err) => {
309 if next_port_to_try.peek().is_none() {
310 return Err(err);
311 }
312 }
313 }
314 }
315 Ok(result)
316}
317
318#[cfg(test)]
319mod tests {
320 use {
321 super::*, ip_echo_server::IpEchoServerResponse, itertools::Itertools, std::net::Ipv4Addr,
322 };
323
324 #[test]
325 fn test_response_length() {
326 let resp = IpEchoServerResponse {
327 address: IpAddr::from([u16::MAX; 8]), shred_version: Some(u16::MAX),
329 };
330 let resp_size = bincode::serialized_size(&resp).unwrap();
331 assert_eq!(
332 IP_ECHO_SERVER_RESPONSE_LENGTH,
333 HEADER_LENGTH + resp_size as usize
334 );
335 }
336
337 #[test]
339 fn test_backward_compat() {
340 let address = IpAddr::from([
341 525u16, 524u16, 523u16, 522u16, 521u16, 520u16, 519u16, 518u16,
342 ]);
343 let response = IpEchoServerResponse {
344 address,
345 shred_version: Some(42),
346 };
347 let mut data = vec![0u8; IP_ECHO_SERVER_RESPONSE_LENGTH];
348 bincode::serialize_into(&mut data[HEADER_LENGTH..], &response).unwrap();
349 data.truncate(HEADER_LENGTH + 20);
350 assert_eq!(
351 bincode::deserialize::<IpAddr>(&data[HEADER_LENGTH..]).unwrap(),
352 address
353 );
354 }
355
356 #[test]
358 fn test_forward_compat() {
359 let address = IpAddr::from([
360 525u16, 524u16, 523u16, 522u16, 521u16, 520u16, 519u16, 518u16,
361 ]);
362 let mut data = [0u8; IP_ECHO_SERVER_RESPONSE_LENGTH];
363 bincode::serialize_into(&mut data[HEADER_LENGTH..], &address).unwrap();
364 let response: Result<IpEchoServerResponse, _> =
365 bincode::deserialize(&data[HEADER_LENGTH..]);
366 assert_eq!(
367 response.unwrap(),
368 IpEchoServerResponse {
369 address,
370 shred_version: None,
371 }
372 );
373 }
374
375 #[test]
376 fn test_parse_port_or_addr() {
377 let p1 = parse_port_or_addr(Some("9000"), SocketAddr::from(([1, 2, 3, 4], 1)));
378 assert_eq!(p1.port(), 9000);
379 let p2 = parse_port_or_addr(Some("127.0.0.1:7000"), SocketAddr::from(([1, 2, 3, 4], 1)));
380 assert_eq!(p2.port(), 7000);
381 let p2 = parse_port_or_addr(Some("hi there"), SocketAddr::from(([1, 2, 3, 4], 1)));
382 assert_eq!(p2.port(), 1);
383 let p3 = parse_port_or_addr(None, SocketAddr::from(([1, 2, 3, 4], 1)));
384 assert_eq!(p3.port(), 1);
385 }
386
387 #[test]
388 fn test_parse_port_range() {
389 assert_eq!(parse_port_range("garbage"), None);
390 assert_eq!(parse_port_range("1-"), None);
391 assert_eq!(parse_port_range("1-2"), Some((1, 2)));
392 assert_eq!(parse_port_range("1-2-3"), None);
393 assert_eq!(parse_port_range("2-1"), None);
394 }
395
396 #[test]
397 fn test_parse_host() {
398 parse_host("localhost:1234").unwrap_err();
399 parse_host("localhost").unwrap();
400 parse_host("127.0.0.0:1234").unwrap_err();
401 parse_host("127.0.0.0").unwrap();
402 parse_host("2001:db8:abcd:42::dead:beef").unwrap_err();
403
404 assert_eq!(
405 select_ipv4(
406 "ipv6-only.test",
407 [IpAddr::V6(std::net::Ipv6Addr::LOCALHOST)].into_iter()
408 )
409 .unwrap_err(),
410 "IPv6 addresses are not supported: ipv6-only.test",
411 );
412 assert_eq!(
413 select_ipv4(
414 "dual-stack.test",
415 [
416 IpAddr::V6(std::net::Ipv6Addr::LOCALHOST),
417 IpAddr::V4(Ipv4Addr::LOCALHOST),
418 ]
419 .into_iter(),
420 )
421 .unwrap(),
422 IpAddr::V4(Ipv4Addr::LOCALHOST),
423 );
424 }
425
426 #[test]
427 fn test_parse_host_port() {
428 parse_host_port("localhost:1234").unwrap();
429 parse_host_port("localhost").unwrap_err();
430 parse_host_port("127.0.0.0:1234").unwrap();
431 parse_host_port("127.0.0.0").unwrap_err();
432 }
433
434 #[test]
435 fn test_is_host_port() {
436 assert!(is_host_port("localhost:1234".to_string()).is_ok());
437 assert!(is_host_port("localhost".to_string()).is_err());
438 }
439
440 #[test]
441 fn test_find_available_port_in_range() {
442 let ip_addr = IpAddr::V4(Ipv4Addr::LOCALHOST);
443 let range = sockets::unique_port_range_for_tests(4);
444 let (pr_s, pr_e) = (range.start, range.end);
445 assert_eq!(
446 find_available_port_in_range(ip_addr, (pr_s, pr_s + 1)).unwrap(),
447 pr_s
448 );
449 let port = find_available_port_in_range(ip_addr, (pr_s, pr_e)).unwrap();
450 assert!((pr_s..pr_e).contains(&port));
451
452 let _socket = sockets::bind_to(ip_addr, port).unwrap();
453 find_available_port_in_range(ip_addr, (port, port + 1)).unwrap_err();
454 }
455
456 #[test]
457 fn test_find_available_ports_in_range() {
458 let ip_addr = IpAddr::V4(Ipv4Addr::LOCALHOST);
459 let port_range = sockets::localhost_port_range_for_tests();
460 assert!(port_range.1 - port_range.0 > 16);
461 let sock = sockets::bind_to_with_config(
463 ip_addr,
464 port_range.0 + 2,
465 sockets::SocketConfiguration::default(),
466 )
467 .unwrap();
468 let ports: [u16; 15] = find_available_ports_in_range(ip_addr, port_range).unwrap();
469 let mut ports_vec = Vec::from(ports);
470 ports_vec.push(sock.local_addr().unwrap().port());
471 let res: Vec<_> = ports_vec.into_iter().unique().collect();
472 assert_eq!(res.len(), 16, "Should reserve 16 unique ports");
473 }
474}