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