1use std::collections::HashSet;
14use std::fmt;
15use std::future::Future;
16use std::net::{IpAddr, SocketAddr};
17use std::time::Duration;
18
19use futures::StreamExt;
20use futures::stream::FuturesUnordered;
21
22pub(crate) const DEFAULT_DELAY: Duration = Duration::from_millis(250);
25
26#[derive(Debug, Clone, PartialEq, Eq)]
32pub struct Failure<E> {
33 pub addr: SocketAddr,
35
36 pub error: E,
38}
39
40impl<E: fmt::Display> fmt::Display for Failure<E> {
41 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
42 write!(f, "{}: {}", self.addr, self.error)
43 }
44}
45
46impl<E: std::error::Error + 'static> std::error::Error for Failure<E> {
47 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
48 Some(&self.error)
49 }
50}
51
52pub(crate) trait Aggregate: Sized {
57 fn aggregate(failures: Vec<Failure<Self>>) -> Self;
62}
63
64pub(crate) fn interleave(addrs: impl IntoIterator<Item = SocketAddr>) -> Vec<SocketAddr> {
70 let (mut a, mut b): (Vec<SocketAddr>, Vec<SocketAddr>) = (Vec::new(), Vec::new());
71 for addr in addrs {
72 if a.is_empty() || a[0].is_ipv4() == addr.is_ipv4() {
73 a.push(addr);
74 } else {
75 b.push(addr);
76 }
77 }
78
79 let mut out = Vec::with_capacity(a.len() + b.len());
80 let (mut a, mut b) = (a.into_iter(), b.into_iter());
81 loop {
82 match (a.next(), b.next()) {
83 (Some(x), Some(y)) => {
84 out.push(x);
85 out.push(y);
86 }
87 (Some(x), None) => out.push(x),
88 (None, Some(y)) => out.push(y),
89 (None, None) => break,
90 }
91 }
92 out
93}
94
95pub(crate) fn match_local(
106 addrs: impl IntoIterator<Item = SocketAddr>,
107 local: SocketAddr,
108 dual_stack: bool,
109) -> Vec<SocketAddr> {
110 let mut seen = HashSet::new();
116 let candidates: Vec<SocketAddr> = interleave(addrs)
117 .into_iter()
118 .map(|addr| normalize_family(addr, local))
119 .filter(|addr| seen.insert(*addr))
120 .collect();
121
122 let usable: Vec<SocketAddr> = candidates
123 .iter()
124 .copied()
125 .filter(|addr| addressable(*addr, local, dual_stack))
126 .collect();
127
128 if usable.is_empty() { candidates } else { usable }
129}
130
131fn addressable(dest: SocketAddr, local: SocketAddr, dual_stack: bool) -> bool {
142 let (SocketAddr::V6(dest), SocketAddr::V6(local)) = (dest, local) else {
143 return dest.is_ipv4() == local.is_ipv4();
144 };
145
146 match (dest.ip().to_ipv4_mapped(), local.ip().to_ipv4_mapped()) {
147 (Some(_), None) => dual_stack && local.ip().is_unspecified(),
148 (None, Some(_)) => false,
149 _ => true,
150 }
151}
152
153fn normalize_family(addr: SocketAddr, local: SocketAddr) -> SocketAddr {
156 match (addr, local.is_ipv4()) {
157 (SocketAddr::V6(v6), true) => match v6.ip().to_ipv4_mapped() {
158 Some(v4) => SocketAddr::new(IpAddr::V4(v4), v6.port()),
159 None => addr,
160 },
161 (SocketAddr::V4(v4), false) => SocketAddr::new(IpAddr::V6(v4.ip().to_ipv6_mapped()), v4.port()),
162 _ => addr,
163 }
164}
165
166pub(crate) fn describe<E: fmt::Display>(failures: &[Failure<E>]) -> String {
168 failures.iter().map(|f| f.to_string()).collect::<Vec<_>>().join("; ")
169}
170
171pub(crate) async fn race<C, E, F, Fut>(candidates: Vec<SocketAddr>, delay: Duration, mut dial: F) -> Result<C, E>
190where
191 F: FnMut(SocketAddr) -> Fut,
192 Fut: Future<Output = Result<C, E>>,
193 E: Aggregate + fmt::Display,
194{
195 let mut remaining = candidates.into_iter();
196 let mut attempts = FuturesUnordered::new();
197 let mut failures: Vec<(usize, Failure<E>)> = Vec::new();
198
199 let mut next_index = 0;
200 let mut start = |addr: SocketAddr, attempts: &mut FuturesUnordered<_>| {
201 let index = next_index;
202 next_index += 1;
203 tracing::debug!(%addr, index, "dialing");
204 let attempt = dial(addr);
205 attempts.push(async move { (index, addr, attempt.await) });
206 };
207
208 let first = remaining.next().expect("no candidates to dial");
209 start(first, &mut attempts);
210
211 loop {
212 tokio::select! {
213 biased;
216
217 res = attempts.next() => {
218 let (index, addr, res) = res.expect("attempts can't be empty here");
219 match res {
220 Ok(conn) => {
221 tracing::debug!(%addr, index, "connected");
222 return Ok(conn);
223 }
224 Err(err) => {
225 tracing::debug!(%addr, index, %err, "connection attempt failed");
230 failures.push((index, Failure { addr, error: err }));
231 if let Some(addr) = remaining.next() {
234 start(addr, &mut attempts);
235 } else if attempts.is_empty() {
236 failures.sort_by_key(|(index, _)| *index);
239 return Err(collapse(failures.into_iter().map(|(_, failure)| failure).collect()));
240 }
241 }
242 }
243 }
244
245 _ = tokio::time::sleep(delay), if remaining.len() > 0 => {
248 let addr = remaining.next().expect("guarded by remaining.len()");
249 start(addr, &mut attempts);
250 }
251 }
252 }
253}
254
255fn collapse<E: Aggregate>(mut failures: Vec<Failure<E>>) -> E {
258 match failures.len() {
259 1 => failures.pop().expect("checked len").error,
260 _ => E::aggregate(failures),
261 }
262}
263
264#[cfg(test)]
265mod tests {
266 use super::*;
267 use std::sync::Arc;
268 use std::sync::atomic::{AtomicUsize, Ordering};
269
270 fn v4(s: &str) -> SocketAddr {
271 s.parse().unwrap()
272 }
273
274 #[derive(Debug, PartialEq, Eq)]
277 enum TestError {
278 Dial(&'static str),
279 All(Vec<Failure<TestError>>),
280 }
281
282 impl fmt::Display for TestError {
283 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
284 match self {
285 Self::Dial(err) => write!(f, "{err}"),
286 Self::All(failures) => write!(f, "all {} attempts failed: {}", failures.len(), describe(failures)),
287 }
288 }
289 }
290
291 impl Aggregate for TestError {
292 fn aggregate(failures: Vec<Failure<Self>>) -> Self {
293 Self::All(failures)
294 }
295 }
296
297 fn failed(addr: &str, err: &'static str) -> Failure<TestError> {
298 Failure {
299 addr: v4(addr),
300 error: TestError::Dial(err),
301 }
302 }
303
304 #[test]
305 fn interleave_alternates_families() {
306 let addrs = [
307 v4("[2001:db8::1]:443"),
308 v4("[2001:db8::2]:443"),
309 v4("1.2.3.4:443"),
310 v4("5.6.7.8:443"),
311 ];
312 assert_eq!(
313 interleave(addrs),
314 vec![
315 v4("[2001:db8::1]:443"),
316 v4("1.2.3.4:443"),
317 v4("[2001:db8::2]:443"),
318 v4("5.6.7.8:443"),
319 ]
320 );
321 }
322
323 #[test]
324 fn interleave_keeps_the_resolver_preferred_family_first() {
325 let addrs = [v4("1.2.3.4:443"), v4("[2001:db8::1]:443")];
327 assert_eq!(interleave(addrs), vec![v4("1.2.3.4:443"), v4("[2001:db8::1]:443")]);
328 }
329
330 #[test]
331 fn interleave_single_family_passthrough() {
332 let addrs = [v4("1.2.3.4:443"), v4("5.6.7.8:443")];
333 assert_eq!(interleave(addrs), addrs.to_vec());
334 }
335
336 #[test]
337 fn match_local_prefers_matching_family() {
338 let a4 = v4("127.0.0.1:443");
339 let a6 = v4("[::1]:443");
340
341 assert_eq!(match_local([a6, a4], v4("0.0.0.0:0"), false), vec![a4]);
343 assert_eq!(
345 match_local([a4, a6], v4("[::]:0"), true),
346 vec![v4("[::ffff:127.0.0.1]:443"), a6]
347 );
348 }
349
350 #[test]
351 fn match_local_skips_mapped_ipv4_on_a_v6_only_socket() {
352 let a4 = v4("192.0.2.1:443");
353 let a6 = v4("[2001:db8::1]:443");
354 assert_eq!(match_local([a4, a6], v4("[::]:0"), false), vec![a6]);
355 }
356
357 #[test]
358 fn match_local_skips_ipv4_for_a_concrete_v6_bind() {
359 let a4 = v4("192.0.2.1:443");
360 let a6 = v4("[2001:db8::1]:443");
361 assert_eq!(match_local([a4, a6], v4("[2001:db8::5]:0"), true), vec![a6]);
362 }
363
364 #[test]
365 fn match_local_keeps_normalized_fallback_when_none_are_usable() {
366 let a4 = v4("192.0.2.1:443");
367 assert_eq!(
368 match_local([a4], v4("[::]:0"), false),
369 vec![v4("[::ffff:192.0.2.1]:443")]
370 );
371 }
372
373 #[test]
374 fn match_local_unwraps_v4_mapped_for_v4_socket() {
375 let mapped = v4("[::ffff:127.0.0.1]:443");
376 assert_eq!(match_local([mapped], v4("0.0.0.0:0"), false), vec![v4("127.0.0.1:443")]);
377 }
378
379 #[test]
380 fn match_local_falls_back_for_unmappable_v6() {
381 let a6 = v4("[2001:db8::1]:443");
384 assert_eq!(match_local([a6], v4("0.0.0.0:0"), false), vec![a6]);
385 }
386
387 #[test]
388 fn match_local_empty() {
389 assert!(match_local(std::iter::empty(), v4("0.0.0.0:0"), false).is_empty());
390 }
391
392 #[test]
393 fn match_local_dedups_across_the_interleave() {
394 let a4 = v4("1.2.3.4:443");
397 let a6 = v4("[2001:db8::1]:443");
398 assert_eq!(match_local([a4, a4, a6], v4("0.0.0.0:0"), false), vec![a4]);
399 assert_eq!(
400 match_local([a4, a4, a6], v4("[::]:0"), true),
401 vec![v4("[::ffff:1.2.3.4]:443"), a6]
402 );
403 }
404
405 #[test]
406 fn match_local_dedups_normalized_forms() {
407 let a4 = v4("1.2.3.4:443");
410 let mapped = v4("[::ffff:1.2.3.4]:443");
411 assert_eq!(match_local([a4, mapped], v4("[::]:0"), true), vec![mapped]);
412 assert_eq!(match_local([mapped, a4], v4("0.0.0.0:0"), false), vec![a4]);
413 }
414
415 #[tokio::test(start_paused = true)]
416 async fn first_success_returns_immediately() {
417 let dials = Arc::new(AtomicUsize::new(0));
418 let counter = dials.clone();
419 let res: Result<&str, TestError> = race(vec![v4("1.1.1.1:1"), v4("2.2.2.2:2")], DEFAULT_DELAY, move |_| {
420 counter.fetch_add(1, Ordering::SeqCst);
421 async { Ok("winner") }
422 })
423 .await;
424 assert_eq!(res, Ok("winner"));
425 assert_eq!(dials.load(Ordering::SeqCst), 1, "no second dial after a fast success");
426 }
427
428 #[tokio::test(start_paused = true)]
429 async fn second_wins_when_first_hangs() {
430 let start = tokio::time::Instant::now();
431 let res: Result<&str, TestError> = race(
432 vec![v4("1.1.1.1:1"), v4("2.2.2.2:2")],
433 DEFAULT_DELAY,
434 |addr| async move {
435 if addr == v4("1.1.1.1:1") {
436 std::future::pending().await
437 } else {
438 Ok("second")
439 }
440 },
441 )
442 .await;
443 assert_eq!(res, Ok("second"));
444 assert_eq!(start.elapsed(), DEFAULT_DELAY, "second dial waits out the stagger");
445 }
446
447 #[tokio::test(start_paused = true)]
448 async fn failure_starts_the_next_attempt_immediately() {
449 let start = tokio::time::Instant::now();
450 let res: Result<&str, TestError> = race(
451 vec![v4("1.1.1.1:1"), v4("2.2.2.2:2")],
452 DEFAULT_DELAY,
453 |addr| async move {
454 if addr == v4("1.1.1.1:1") {
455 Err(TestError::Dial("boom"))
456 } else {
457 Ok("second")
458 }
459 },
460 )
461 .await;
462 assert_eq!(res, Ok("second"));
463 assert_eq!(start.elapsed(), Duration::ZERO, "failure must not wait for the timer");
464 }
465
466 #[tokio::test(start_paused = true)]
470 async fn all_failures_are_reported_when_the_preferred_fails_first() {
471 let res: Result<&str, TestError> = race(
472 vec![v4("1.1.1.1:1"), v4("2.2.2.2:2")],
473 Duration::from_millis(10),
474 |addr| async move {
475 if addr == v4("1.1.1.1:1") {
476 Err(TestError::Dial("network unreachable"))
477 } else {
478 tokio::time::sleep(Duration::from_secs(1)).await;
479 Err(TestError::Dial("invalid peer certificate"))
480 }
481 },
482 )
483 .await;
484 assert_eq!(
485 res,
486 Err(TestError::All(vec![
487 failed("1.1.1.1:1", "network unreachable"),
488 failed("2.2.2.2:2", "invalid peer certificate"),
489 ]))
490 );
491 }
492
493 #[tokio::test(start_paused = true)]
498 async fn all_failures_are_reported_when_the_preferred_times_out_last() {
499 let res: Result<&str, TestError> = race(
500 vec![v4("1.1.1.1:1"), v4("2.2.2.2:2")],
501 Duration::from_millis(10),
502 |addr| async move {
503 if addr == v4("1.1.1.1:1") {
504 tokio::time::sleep(Duration::from_secs(30)).await;
505 Err(TestError::Dial("timed out"))
506 } else {
507 Err(TestError::Dial("invalid peer certificate"))
508 }
509 },
510 )
511 .await;
512 assert_eq!(
513 res,
514 Err(TestError::All(vec![
515 failed("1.1.1.1:1", "timed out"),
516 failed("2.2.2.2:2", "invalid peer certificate"),
517 ]))
518 );
519 }
520
521 #[tokio::test(start_paused = true)]
524 async fn a_lone_failure_is_returned_unwrapped() {
525 let res: Result<&str, TestError> = race(vec![v4("1.1.1.1:1")], DEFAULT_DELAY, |_| async {
526 Err(TestError::Dial("invalid peer certificate"))
527 })
528 .await;
529 assert_eq!(res, Err(TestError::Dial("invalid peer certificate")));
530 }
531
532 #[test]
533 fn describe_lists_every_attempt() {
534 let failures = [failed("1.1.1.1:1", "timed out"), failed("2.2.2.2:2", "bad cert")];
535 assert_eq!(describe(&failures), "1.1.1.1:1: timed out; 2.2.2.2:2: bad cert");
536 }
537
538 #[tokio::test(start_paused = true)]
539 async fn losers_are_dropped_on_success() {
540 struct Guard(Arc<AtomicUsize>);
542 impl Drop for Guard {
543 fn drop(&mut self) {
544 self.0.fetch_add(1, Ordering::SeqCst);
545 }
546 }
547
548 let dropped = Arc::new(AtomicUsize::new(0));
549 let count = dropped.clone();
550 let res: Result<&str, TestError> = race(vec![v4("1.1.1.1:1"), v4("2.2.2.2:2")], Duration::ZERO, move |addr| {
551 let guard = Guard(count.clone());
552 async move {
553 if addr == v4("1.1.1.1:1") {
554 let _guard = guard;
555 std::future::pending().await
556 } else {
557 drop(guard);
558 tokio::time::sleep(Duration::from_millis(1)).await;
559 Ok("second")
560 }
561 }
562 })
563 .await;
564 assert_eq!(res, Ok("second"));
565 assert_eq!(dropped.load(Ordering::SeqCst), 2, "the hung attempt was not aborted");
566 }
567
568 #[tokio::test(start_paused = true)]
569 async fn zero_delay_dials_all_at_once() {
570 let start = tokio::time::Instant::now();
571 let res: Result<&str, TestError> = race(
572 vec![v4("1.1.1.1:1"), v4("2.2.2.2:2")],
573 Duration::ZERO,
574 |addr| async move {
575 if addr == v4("1.1.1.1:1") {
576 std::future::pending().await
577 } else {
578 Ok("second")
579 }
580 },
581 )
582 .await;
583 assert_eq!(res, Ok("second"));
584 assert_eq!(start.elapsed(), Duration::ZERO);
585 }
586}