From fa3b39a42e5b95dc06f41fd9a0b12bcc10a0eda1 Mon Sep 17 00:00:00 2001 From: dvloplerz Date: Wed, 30 Jul 2025 00:47:19 +0700 Subject: [PATCH 1/2] feat(dns): add exponentail backoff retry for reverse lookup. --- Cargo.toml | 14 ++--- src/network/dns/resolver.rs | 105 +++++++++++++++++++++++++++++++----- src/tests/cases/mod.rs | 1 + src/tests/cases/network.rs | 63 ++++++++++++++++++++++ 4 files changed, 163 insertions(+), 20 deletions(-) create mode 100644 src/tests/cases/network.rs diff --git a/Cargo.toml b/Cargo.toml index 29c2a630..73b7599a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,12 +2,12 @@ name = "bandwhich" version = "0.23.1" authors = [ - "Aram Drevekenin ", - "Eduardo Toledo ", - "Eduardo Broto ", - "Kelvin Zhang ", - "Brooks Rady ", - "cyqsimon <28627918+cyqsimon@users.noreply.github.com>", + "Aram Drevekenin ", + "Eduardo Toledo ", + "Eduardo Broto ", + "Kelvin Zhang ", + "Brooks Rady ", + "cyqsimon <28627918+cyqsimon@users.noreply.github.com>", ] categories = ["network-programming", "command-line-utilities"] edition = "2021" @@ -42,7 +42,7 @@ ratatui = "0.29.0" resolv-conf = "0.7.4" simplelog = "0.12.2" thiserror = "2.0.12" -tokio = { version = "1.46", features = ["rt", "sync"] } +tokio = { version = "1.46", features = ["rt", "sync", "macros"] } trust-dns-resolver = "0.23.2" unicode-width = "0.2.0" strum = { version = "0.27.1", features = ["derive"] } diff --git a/src/network/dns/resolver.rs b/src/network/dns/resolver.rs index 46995452..c74c5bf4 100644 --- a/src/network/dns/resolver.rs +++ b/src/network/dns/resolver.rs @@ -1,9 +1,13 @@ -use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4}; +use std::{ + net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4}, + time::Duration, +}; use async_trait::async_trait; +use log::warn; +use tokio::time::sleep; use trust_dns_resolver::{ config::{NameServerConfig, Protocol, ResolverConfig, ResolverOpts}, - error::ResolveErrorKind, TokioAsyncResolver, }; @@ -40,18 +44,93 @@ impl Resolver { #[async_trait] impl Lookup for Resolver { async fn lookup(&self, ip: IpAddr) -> Option { - let lookup_future = self.0.reverse_lookup(ip); - match lookup_future.await { - Ok(names) => { - // Take the first result and convert it to a string - names.into_iter().next().map(|name| name.to_string()) - } - Err(e) => match e.kind() { - // If the IP is not associated with a hostname, store the IP - // so that we don't retry indefinitely - ResolveErrorKind::NoRecordsFound { .. } => Some(ip.to_string()), - _ => None, + let mut retries = 0_u8; + // loop { + // match self.0.reverse_lookup(ip).await { + // Ok(names) => { + // // Take the first result and convert it to a string + // return names.into_iter().next().map(|name| name.to_string()); + // } + // Err(e) => match e.kind() { + // // If the IP is not associated with a hostname, store the IP + // // so that we don't retry indefinitely + // ResolveErrorKind::NoRecordsFound { .. } => return Some(ip.to_string()), + // ResolveErrorKind::Timeout if retries < 5 => { + // sleep(Duration::from_millis(1000)).await; + // retries += 1; + // continue; + // } + // ResolveErrorKind::Timeout => { + // return Some(String::from("DNS lookup timeout.")); + // } + // _ => break None, + // }, + // }; + // } + let retry_config = RetryPolicy { + max_retries: 3, + ..Default::default() + }; + + retry_with_backoff( + || { + let resolver = &self.0; + async move { + resolver + .reverse_lookup(ip) + .await + .ok() + .and_then(|names| names.iter().next().map(|n| n.to_string())) + .or_else(|| Some(ip.to_string())) + } }, + retry_config.max_retries, + retry_config.base_delay, + ) + .await + .or_else(|| Some("DNS lookup timeout.".into())) + } +} + +struct RetryPolicy { + max_retries: u8, + base_delay: tokio::time::Duration, +} + +impl Default for RetryPolicy { + fn default() -> Self { + Self { + max_retries: 2, + base_delay: Duration::from_millis(1000), + } + } +} + +pub async fn retry_with_backoff( + mut operation: F, + max_reties: u8, + inittial_delay: tokio::time::Duration, +) -> Option +where + F: FnMut() -> Fut, + Fut: std::future::Future>, +{ + let mut delay = inittial_delay; + for attemp in 0..=max_reties { + match operation().await { + Some(value) => return Some(value), + None if attemp < max_reties => { + warn!( + "Retrying.. attemp: {}/{} (waiting {:?})", + attemp + 1, + max_reties, + delay + ); + sleep(delay).await; + delay *= 2; + } + None => return None, } } + None } diff --git a/src/tests/cases/mod.rs b/src/tests/cases/mod.rs index d0236bf6..5abb7bec 100644 --- a/src/tests/cases/mod.rs +++ b/src/tests/cases/mod.rs @@ -1,3 +1,4 @@ +pub mod network; pub mod raw_mode; pub mod test_utils; #[cfg(feature = "ui_test")] diff --git a/src/tests/cases/network.rs b/src/tests/cases/network.rs new file mode 100644 index 00000000..a70e8db3 --- /dev/null +++ b/src/tests/cases/network.rs @@ -0,0 +1,63 @@ +#[cfg(test)] +mod tests { + use super::*; + use crate::network::dns; + use std::sync::{ + atomic::{AtomicU8, Ordering}, + Arc, + }; + use tokio::time::Instant; + + #[tokio::test] + async fn retry_should_succeed_after_3_attempts() { + let counter = Arc::new(AtomicU8::new(0)); + let counter_clone = counter.clone(); + + let start = Instant::now(); + + let result = dns::retry_with_backoff( + move || { + let counter = counter_clone.clone(); + async move { + let attempt = counter.fetch_add(1, Ordering::SeqCst); + if attempt >= 2 { + Some("Success".to_string()) + } else { + None + } + } + }, + 5, + std::time::Duration::from_millis(50), + ) + .await; + + let duration = start.elapsed(); + + assert_eq!(result, Some("Success".to_string())); + assert!(duration >= std::time::Duration::from_millis(50 + 100)); // 2 delays + assert!(counter.load(Ordering::SeqCst) == 3); // called 3 times + } + + #[tokio::test] + async fn retry_should_fail_after_max_retries() { + let counter = Arc::new(AtomicU8::new(0)); + let counter_clone = counter.clone(); + + let result: Option<()> = dns::retry_with_backoff( + move || { + let counter = counter_clone.clone(); + async move { + counter.fetch_add(1, Ordering::SeqCst); + None + } + }, + 3, + std::time::Duration::from_millis(10), + ) + .await; + + assert_eq!(result, None); + assert_eq!(counter.load(Ordering::SeqCst), 4); // initial try + 3 retries + } +} From 475299d2ac1b7eb794651c62176cfa8302becd4d Mon Sep 17 00:00:00 2001 From: dvloplerz Date: Wed, 30 Jul 2025 00:59:34 +0700 Subject: [PATCH 2/2] Clean up previous dirty things! --- src/network/dns/resolver.rs | 23 ----------------------- 1 file changed, 23 deletions(-) diff --git a/src/network/dns/resolver.rs b/src/network/dns/resolver.rs index c74c5bf4..e7324090 100644 --- a/src/network/dns/resolver.rs +++ b/src/network/dns/resolver.rs @@ -44,29 +44,6 @@ impl Resolver { #[async_trait] impl Lookup for Resolver { async fn lookup(&self, ip: IpAddr) -> Option { - let mut retries = 0_u8; - // loop { - // match self.0.reverse_lookup(ip).await { - // Ok(names) => { - // // Take the first result and convert it to a string - // return names.into_iter().next().map(|name| name.to_string()); - // } - // Err(e) => match e.kind() { - // // If the IP is not associated with a hostname, store the IP - // // so that we don't retry indefinitely - // ResolveErrorKind::NoRecordsFound { .. } => return Some(ip.to_string()), - // ResolveErrorKind::Timeout if retries < 5 => { - // sleep(Duration::from_millis(1000)).await; - // retries += 1; - // continue; - // } - // ResolveErrorKind::Timeout => { - // return Some(String::from("DNS lookup timeout.")); - // } - // _ => break None, - // }, - // }; - // } let retry_config = RetryPolicy { max_retries: 3, ..Default::default()