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 ip_echo_client::IpEchoClientError,
23 ip_echo_server::{
24 DEFAULT_IP_ECHO_SERVER_THREADS, IpEchoServer, MAX_PORT_COUNT_PER_MESSAGE, ip_echo_server,
25 },
26 pinned_xdp_sender::PinnedXdpSender,
27 socket_addr_space::SocketAddrSpace,
28};
29use {
30 ip_echo_client::{ip_echo_server_request, ip_echo_server_request_with_binding},
31 ip_echo_server::IpEchoServerMessage,
32 rand::{Rng, rng},
33 std::{
34 io::{self},
35 net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener, ToSocketAddrs, UdpSocket},
36 },
37 url::Url,
38};
39
40pub struct UdpSocketPair {
42 pub addr: SocketAddr, pub receiver: UdpSocket, pub sender: UdpSocket, }
46
47pub type PortRange = (u16, u16);
48
49#[cfg(not(debug_assertions))]
50pub const VALIDATOR_PORT_RANGE: PortRange = (8000, 10_000);
52
53#[cfg(debug_assertions)]
56pub const VALIDATOR_PORT_RANGE: PortRange = (
57 crate::sockets::UNIQUE_ALLOC_BASE_PORT - 512,
58 crate::sockets::UNIQUE_ALLOC_BASE_PORT,
59);
60
61pub const MINIMUM_VALIDATOR_PORT_RANGE_WIDTH: u16 = 26; pub(crate) const HEADER_LENGTH: usize = 4;
64pub(crate) const IP_ECHO_SERVER_RESPONSE_LENGTH: usize = HEADER_LENGTH + 23;
65
66pub fn get_public_ip_addr_with_binding(
69 ip_echo_server_addr: &SocketAddr,
70 bind_address: IpAddr,
71) -> Result<IpAddr, IpEchoClientError> {
72 let fut = ip_echo_server_request_with_binding(
73 *ip_echo_server_addr,
74 IpEchoServerMessage::default(),
75 bind_address,
76 );
77 let rt = tokio::runtime::Builder::new_current_thread()
78 .enable_all()
79 .build()?;
80 let resp = rt.block_on(fut)?;
81 Ok(resp.address)
82}
83
84pub fn get_cluster_shred_version(ip_echo_server_addr: &SocketAddr) -> Result<u16, String> {
86 let fut = ip_echo_server_request(*ip_echo_server_addr, IpEchoServerMessage::default());
87 let rt = tokio::runtime::Builder::new_current_thread()
88 .enable_all()
89 .build()
90 .map_err(|e| e.to_string())?;
91 let resp = rt.block_on(fut).map_err(|e| e.to_string())?;
92 resp.shred_version
93 .ok_or_else(|| "IP echo server does not return a shred-version".to_owned())
94}
95
96pub fn get_cluster_shred_version_with_binding(
99 ip_echo_server_addr: &SocketAddr,
100 bind_address: IpAddr,
101) -> Result<u16, IpEchoClientError> {
102 let fut = ip_echo_server_request_with_binding(
103 *ip_echo_server_addr,
104 IpEchoServerMessage::default(),
105 bind_address,
106 );
107 let rt = tokio::runtime::Builder::new_current_thread()
108 .enable_all()
109 .build()?;
110 let resp = rt.block_on(fut)?;
111 resp.shred_version.ok_or_else(|| {
112 IpEchoClientError::InvalidResponse(
113 "IP echo server does not return a shred-version".to_owned(),
114 )
115 })
116}
117
118const MAX_PORT_VERIFY_THREADS: usize = 64;
121
122pub fn verify_all_reachable_udp(
127 ip_echo_server_addr: &SocketAddr,
128 udp_sockets: &[&UdpSocket],
129) -> bool {
130 let rt = tokio::runtime::Builder::new_current_thread()
131 .enable_all()
132 .max_blocking_threads(MAX_PORT_VERIFY_THREADS)
133 .build()
134 .expect("Tokio builder should be able to reliably create a current thread runtime");
135 let fut = ip_echo_client::verify_all_reachable_udp(
136 *ip_echo_server_addr,
137 udp_sockets,
138 ip_echo_client::TIMEOUT,
139 ip_echo_client::DEFAULT_RETRY_COUNT,
140 );
141 rt.block_on(fut)
142}
143
144pub fn verify_all_reachable_tcp(
149 ip_echo_server_addr: &SocketAddr,
150 tcp_listeners: Vec<TcpListener>,
151) -> bool {
152 let rt = tokio::runtime::Builder::new_current_thread()
153 .enable_all()
154 .max_blocking_threads(MAX_PORT_VERIFY_THREADS)
155 .build()
156 .expect("Tokio builder should be able to reliably create a current thread runtime");
157 let fut = ip_echo_client::verify_all_reachable_tcp(
158 *ip_echo_server_addr,
159 tcp_listeners,
160 ip_echo_client::TIMEOUT,
161 );
162 rt.block_on(fut)
163}
164
165pub fn parse_port_or_addr(optstr: Option<&str>, default_addr: SocketAddr) -> SocketAddr {
166 if let Some(addrstr) = optstr {
167 if let Ok(port) = addrstr.parse() {
168 let mut addr = default_addr;
169 addr.set_port(port);
170 addr
171 } else if let Ok(addr) = addrstr.parse() {
172 addr
173 } else {
174 default_addr
175 }
176 } else {
177 default_addr
178 }
179}
180
181pub fn parse_port_range(port_range: &str) -> Option<PortRange> {
182 let ports: Vec<&str> = port_range.split('-').collect();
183 if ports.len() != 2 {
184 return None;
185 }
186
187 let start_port = ports[0].parse();
188 let end_port = ports[1].parse();
189
190 if start_port.is_err() || end_port.is_err() {
191 return None;
192 }
193 let start_port = start_port.unwrap();
194 let end_port = end_port.unwrap();
195 if end_port < start_port {
196 return None;
197 }
198 Some((start_port, end_port))
199}
200
201fn select_ipv4<T>(
202 host: &str,
203 mut values: impl Iterator<Item = T>,
204 mut ip_addr: impl FnMut(&T) -> IpAddr,
205) -> Result<T, String> {
206 let Some(first_value) = values.next() else {
207 return Err(format!("Unable to resolve host: {host}"));
208 };
209
210 if ip_addr(&first_value).is_ipv4() {
211 return Ok(first_value);
212 }
213
214 values
215 .find(|value| ip_addr(value).is_ipv4())
216 .ok_or_else(|| format!("IPv6 addresses are not supported: {host}"))
217}
218
219pub fn parse_host(host: &str) -> Result<IpAddr, String> {
220 if let Ok(IpAddr::V6(_)) = host.parse::<IpAddr>() {
221 return Err(format!("IPv6 addresses are not supported: {host}"));
222 }
223
224 let parsed_url = Url::parse(&format!("http://{host}")).map_err(|e| e.to_string())?;
227 if parsed_url.port().is_some() {
228 return Err(format!("Expected port in URL: {host}"));
229 }
230
231 let ips = (host, 0)
233 .to_socket_addrs()
234 .map_err(|err| err.to_string())?
235 .map(|socket_address| socket_address.ip());
236
237 select_ipv4(host, ips, |ip| *ip)
238}
239
240pub fn is_host(string: String) -> Result<(), String> {
241 parse_host(&string).map(|_| ())
242}
243
244pub fn parse_host_port(host_port: &str) -> Result<SocketAddr, String> {
245 let addrs = host_port
246 .to_socket_addrs()
247 .map_err(|err| format!("Unable to resolve host {host_port}: {err}"))?;
248 select_ipv4(host_port, addrs, SocketAddr::ip)
249}
250
251pub fn is_host_port(string: String) -> Result<(), String> {
252 parse_host_port(&string).map(|_| ())
253}
254
255pub fn bind_in_range(ip_addr: IpAddr, range: PortRange) -> io::Result<(u16, UdpSocket)> {
256 let config = sockets::SocketConfiguration::default();
257 sockets::bind_in_range_with_config(ip_addr, range, config)
258}
259
260pub fn bind_to_unspecified() -> io::Result<UdpSocket> {
261 let config = sockets::SocketConfiguration::default();
262 sockets::bind_to_with_config(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0, config)
263}
264
265pub fn find_available_port_in_range(ip_addr: IpAddr, range: PortRange) -> io::Result<u16> {
272 let [port] = find_available_ports_in_range(ip_addr, range)?;
273 Ok(port)
274}
275
276pub fn find_available_ports_in_range<const N: usize>(
281 ip_addr: IpAddr,
282 range: PortRange,
283) -> io::Result<[u16; N]> {
284 let mut result = [0u16; N];
285 let range = range.0..range.1;
286 let mut next_port_to_try = range
287 .clone()
288 .cycle() .skip(rng().random_range(range.clone()) as usize) .take(range.len()) .peekable();
292 let mut num = 0;
293 let config = sockets::SocketConfiguration::default();
294 while num < N {
295 let port_to_try = next_port_to_try.next().unwrap(); let bind = sockets::bind_common_with_config(ip_addr, port_to_try, config);
297 match bind {
298 Ok(_) => {
299 result[num] = port_to_try;
300 num = num.saturating_add(1);
301 }
302 Err(err) => {
303 if next_port_to_try.peek().is_none() {
304 return Err(err);
305 }
306 }
307 }
308 }
309 Ok(result)
310}
311
312#[cfg(test)]
313mod tests {
314 use {
315 super::*, ip_echo_server::IpEchoServerResponse, itertools::Itertools, std::net::Ipv4Addr,
316 };
317
318 #[test]
319 fn test_response_length() {
320 let resp = IpEchoServerResponse {
321 address: IpAddr::from([u16::MAX; 8]), shred_version: Some(u16::MAX),
323 };
324 let resp_size = bincode::serialized_size(&resp).unwrap();
325 assert_eq!(
326 IP_ECHO_SERVER_RESPONSE_LENGTH,
327 HEADER_LENGTH + resp_size as usize
328 );
329 }
330
331 #[test]
333 fn test_backward_compat() {
334 let address = IpAddr::from([
335 525u16, 524u16, 523u16, 522u16, 521u16, 520u16, 519u16, 518u16,
336 ]);
337 let response = IpEchoServerResponse {
338 address,
339 shred_version: Some(42),
340 };
341 let mut data = vec![0u8; IP_ECHO_SERVER_RESPONSE_LENGTH];
342 bincode::serialize_into(&mut data[HEADER_LENGTH..], &response).unwrap();
343 data.truncate(HEADER_LENGTH + 20);
344 assert_eq!(
345 bincode::deserialize::<IpAddr>(&data[HEADER_LENGTH..]).unwrap(),
346 address
347 );
348 }
349
350 #[test]
352 fn test_forward_compat() {
353 let address = IpAddr::from([
354 525u16, 524u16, 523u16, 522u16, 521u16, 520u16, 519u16, 518u16,
355 ]);
356 let mut data = [0u8; IP_ECHO_SERVER_RESPONSE_LENGTH];
357 bincode::serialize_into(&mut data[HEADER_LENGTH..], &address).unwrap();
358 let response: Result<IpEchoServerResponse, _> =
359 bincode::deserialize(&data[HEADER_LENGTH..]);
360 assert_eq!(
361 response.unwrap(),
362 IpEchoServerResponse {
363 address,
364 shred_version: None,
365 }
366 );
367 }
368
369 #[test]
370 fn test_parse_port_or_addr() {
371 let p1 = parse_port_or_addr(Some("9000"), SocketAddr::from(([1, 2, 3, 4], 1)));
372 assert_eq!(p1.port(), 9000);
373 let p2 = parse_port_or_addr(Some("127.0.0.1:7000"), SocketAddr::from(([1, 2, 3, 4], 1)));
374 assert_eq!(p2.port(), 7000);
375 let p2 = parse_port_or_addr(Some("hi there"), SocketAddr::from(([1, 2, 3, 4], 1)));
376 assert_eq!(p2.port(), 1);
377 let p3 = parse_port_or_addr(None, SocketAddr::from(([1, 2, 3, 4], 1)));
378 assert_eq!(p3.port(), 1);
379 }
380
381 #[test]
382 fn test_parse_port_range() {
383 assert_eq!(parse_port_range("garbage"), None);
384 assert_eq!(parse_port_range("1-"), None);
385 assert_eq!(parse_port_range("1-2"), Some((1, 2)));
386 assert_eq!(parse_port_range("1-2-3"), None);
387 assert_eq!(parse_port_range("2-1"), None);
388 }
389
390 #[test]
391 fn test_parse_host() {
392 parse_host("localhost:1234").unwrap_err();
393 parse_host("localhost").unwrap();
394 parse_host("127.0.0.0:1234").unwrap_err();
395 parse_host("127.0.0.0").unwrap();
396 parse_host("2001:db8:abcd:42::dead:beef").unwrap_err();
397
398 assert_eq!(
399 select_ipv4(
400 "ipv6-only.test",
401 [IpAddr::V6(std::net::Ipv6Addr::LOCALHOST)].into_iter(),
402 |ip| *ip,
403 )
404 .unwrap_err(),
405 "IPv6 addresses are not supported: ipv6-only.test",
406 );
407 assert_eq!(
408 select_ipv4(
409 "dual-stack.test",
410 [
411 IpAddr::V6(std::net::Ipv6Addr::LOCALHOST),
412 IpAddr::V4(Ipv4Addr::LOCALHOST),
413 ]
414 .into_iter(),
415 |ip| *ip,
416 )
417 .unwrap(),
418 IpAddr::V4(Ipv4Addr::LOCALHOST),
419 );
420 }
421
422 #[test]
423 fn test_parse_host_port() {
424 parse_host_port("localhost:1234").unwrap();
425 parse_host_port("localhost").unwrap_err();
426 parse_host_port("127.0.0.0:1234").unwrap();
427 parse_host_port("127.0.0.0").unwrap_err();
428 assert_eq!(
429 parse_host_port("[2001:db8:abcd:42::dead:beef]:1234").unwrap_err(),
430 "IPv6 addresses are not supported: [2001:db8:abcd:42::dead:beef]:1234",
431 );
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}