use std::collections::HashMap;
use std::net::{IpAddr, SocketAddr};
use std::time::{Duration, Instant};
use hickory_resolver::TokioAsyncResolver;
use hickory_resolver::config::{ResolverConfig, ResolverOpts};
#[derive(Debug, Clone)]
pub struct DnsEntry {
pub hostname: String,
pub addresses: Vec<SocketAddr>,
pub resolved_at: Instant,
pub ttl: Duration,
pub ipv4_preferred: bool,
}
impl DnsEntry {
pub fn is_expired(&self) -> bool {
self.resolved_at.elapsed() > self.ttl
}
pub fn best_address(&self) -> Option<SocketAddr> {
if self.addresses.is_empty() {
return None;
}
if self.ipv4_preferred {
self.addresses
.iter()
.find(|a| matches!(a.ip(), IpAddr::V4(_)))
.copied()
.or_else(|| self.addresses.first().copied())
} else {
Some(self.addresses[0])
}
}
pub fn all_addresses(&self) -> Vec<SocketAddr> {
self.addresses.clone()
}
}
pub struct DnsCache {
cache: HashMap<String, DnsEntry>,
default_ttl: Duration,
negative_ttl: Duration,
negative_entries: HashMap<String, Instant>,
ipv4_preference: bool,
resolver: Option<TokioAsyncResolver>,
}
impl DnsCache {
pub fn new() -> Self {
Self {
cache: HashMap::new(),
default_ttl: Duration::from_secs(300), negative_ttl: Duration::from_secs(60), negative_entries: HashMap::new(),
ipv4_preference: true, resolver: build_default_resolver(),
}
}
pub fn with_ttl(default_ttl_secs: u64, negative_ttl_secs: u64) -> Self {
Self {
default_ttl: Duration::from_secs(default_ttl_secs),
negative_ttl: Duration::from_secs(negative_ttl_secs),
..Self::new()
}
}
pub fn with_resolver(mut self, resolver: TokioAsyncResolver) -> Self {
self.resolver = Some(resolver);
self
}
pub async fn resolve(&mut self, hostname: &str, port: u16) -> Result<Vec<SocketAddr>, String> {
if let Some(entry) = self.cache.get(hostname)
&& !entry.is_expired()
{
return Ok(entry.all_addresses());
}
if let Some(failed_at) = self.negative_entries.get(hostname)
&& failed_at.elapsed() < self.negative_ttl
{
return Err(format!(
"DNS lookup recently failed for {} (retry after {:?})",
hostname,
self.negative_ttl.saturating_sub(failed_at.elapsed())
));
}
if let Some(resolver) = self.resolver.as_ref() {
match resolver.lookup_ip(hostname).await {
Ok(lookup) => {
let mut addrs: Vec<SocketAddr> =
lookup.iter().map(|ip| SocketAddr::new(ip, port)).collect();
if !addrs.is_empty() {
let record_ttl = lookup
.valid_until()
.saturating_duration_since(Instant::now());
let actual_ttl = record_ttl.min(self.default_ttl);
if self.ipv4_preference {
addrs.sort_by_key(|a| match a.ip() {
IpAddr::V4(_) => 0u8,
IpAddr::V6(_) => 1u8,
});
}
let hostname_owned = hostname.to_string();
let entry = DnsEntry {
hostname: hostname_owned.clone(),
addresses: addrs.clone(),
resolved_at: Instant::now(),
ttl: actual_ttl,
ipv4_preferred: self.ipv4_preference,
};
self.cache.insert(hostname_owned, entry);
self.negative_entries.remove(hostname);
tracing::trace!(
hostname = hostname,
addr_count = addrs.len(),
ttl_secs = actual_ttl.as_secs(),
"DNS resolved via hickory async resolver"
);
return Ok(addrs);
}
tracing::debug!(
hostname = hostname,
"hickory returned no addresses, falling back to tokio::net::lookup_host"
);
}
Err(e) => {
tracing::debug!(
hostname = hostname,
error = %e,
"hickory DNS lookup failed, falling back to tokio::net::lookup_host"
);
}
}
}
let addr_str = format!("{}:{}", hostname, port);
match tokio::net::lookup_host(&addr_str).await {
Ok(addrs) => {
let mut sorted: Vec<SocketAddr> = addrs.collect();
if self.ipv4_preference {
sorted.sort_by_key(|a| match a.ip() {
IpAddr::V4(_) => 0u8,
IpAddr::V6(_) => 1u8,
});
}
let hostname_owned = hostname.to_string();
let entry = DnsEntry {
hostname: hostname_owned.clone(),
addresses: sorted.clone(),
resolved_at: Instant::now(),
ttl: self.default_ttl,
ipv4_preferred: self.ipv4_preference,
};
self.cache.insert(hostname_owned, entry);
self.negative_entries.remove(hostname);
tracing::trace!(
hostname = hostname,
addr_count = sorted.len(),
"DNS resolved via tokio::net::lookup_host fallback"
);
Ok(sorted)
}
Err(e) => {
self.negative_entries
.insert(hostname.to_string(), Instant::now());
tracing::debug!(
hostname = hostname,
error = %e,
"DNS resolution failed via fallback"
);
Err(e.to_string())
}
}
}
pub async fn force_refresh(
&mut self,
hostname: &str,
port: u16,
) -> Result<Vec<SocketAddr>, String> {
self.cache.remove(hostname);
self.resolve(hostname, port).await
}
pub fn clear(&mut self) {
self.cache.clear();
self.negative_entries.clear();
}
pub fn purge_expired(&mut self) -> usize {
let before = self.cache.len();
self.cache.retain(|_, v| !v.is_expired());
self.negative_entries
.retain(|_, t| t.elapsed() < self.negative_ttl);
before - self.cache.len()
}
pub fn set_ipv4_preference(&mut self, prefer_ipv4: bool) {
self.ipv4_preference = prefer_ipv4;
}
pub fn len(&self) -> usize {
self.cache.len()
}
pub fn is_empty(&self) -> bool {
self.cache.is_empty()
}
pub fn default_ttl(&self) -> Duration {
self.default_ttl
}
pub fn negative_ttl(&self) -> Duration {
self.negative_ttl
}
pub fn record_failure(&mut self, hostname: &str) {
self.negative_entries
.insert(hostname.to_string(), Instant::now());
}
pub fn resolve_no_network(
&mut self,
hostname: &str,
port: u16,
) -> Result<Vec<SocketAddr>, String> {
if let Some(entry) = self.cache.get(hostname)
&& !entry.is_expired()
{
return Ok(entry.all_addresses());
}
if let Some(failed_at) = self.negative_entries.get(hostname)
&& failed_at.elapsed() < self.negative_ttl
{
return Err(format!(
"DNS lookup recently failed for {} (retry after {:?})",
hostname,
self.negative_ttl.saturating_sub(failed_at.elapsed())
));
}
Err(format!("No cached entry for {}:{}", hostname, port))
}
}
fn build_default_resolver() -> Option<TokioAsyncResolver> {
let mut opts = ResolverOpts::default();
opts.use_hosts_file = true;
Some(TokioAsyncResolver::tokio(ResolverConfig::default(), opts))
}
impl Default for DnsCache {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn create_test_cache() -> DnsCache {
DnsCache::with_ttl(10, 1) }
#[test]
fn test_dns_entry_is_expired() {
let entry = DnsEntry {
hostname: "test.com".to_string(),
addresses: vec![],
resolved_at: Instant::now(),
ttl: Duration::from_secs(60),
ipv4_preferred: true,
};
assert!(!entry.is_expired());
let expired_entry = DnsEntry {
hostname: "old.com".to_string(),
addresses: vec![],
resolved_at: Instant::now() - Duration::from_secs(61),
ttl: Duration::from_secs(60),
ipv4_preferred: false,
};
assert!(expired_entry.is_expired());
}
#[test]
fn test_dns_entry_best_address_ipv4_preferred() {
let ipv6_addr: SocketAddr = "[::1]:8080".parse().unwrap();
let ipv4_addr: SocketAddr = "127.0.0.1:8080".parse().unwrap();
let entry = DnsEntry {
hostname: "mixed.com".to_string(),
addresses: vec![ipv6_addr, ipv4_addr], resolved_at: Instant::now(),
ttl: Duration::from_secs(60),
ipv4_preferred: true,
};
let best = entry.best_address().unwrap();
assert_eq!(best, ipv4_addr);
}
#[test]
fn test_dns_entry_best_address_no_ipv4() {
let ipv6_addr: SocketAddr = "[::1]:8080".parse().unwrap();
let entry = DnsEntry {
hostname: "ipv6only.com".to_string(),
addresses: vec![ipv6_addr],
resolved_at: Instant::now(),
ttl: Duration::from_secs(60),
ipv4_preferred: true,
};
let best = entry.best_address().unwrap();
assert_eq!(best, ipv6_addr);
}
#[test]
fn test_dns_entry_best_address_empty() {
let entry = DnsEntry {
hostname: "empty.com".to_string(),
addresses: vec![],
resolved_at: Instant::now(),
ttl: Duration::from_secs(60),
ipv4_preferred: true,
};
assert!(entry.best_address().is_none());
}
#[test]
fn test_dns_cache_creation() {
let cache = DnsCache::new();
assert!(cache.is_empty());
assert_eq!(cache.default_ttl(), Duration::from_secs(300));
assert_eq!(cache.negative_ttl(), Duration::from_secs(60));
}
#[test]
fn test_dns_cache_with_custom_ttl() {
let cache = DnsCache::with_ttl(600, 30);
assert_eq!(cache.default_ttl(), Duration::from_secs(600));
assert_eq!(cache.negative_ttl(), Duration::from_secs(30));
}
#[test]
fn test_dns_cache_clear() {
let mut cache = create_test_cache();
let entry = DnsEntry {
hostname: "example.com".to_string(),
addresses: vec!["127.0.0.1:80".parse().unwrap()],
resolved_at: Instant::now(),
ttl: Duration::from_secs(60),
ipv4_preferred: true,
};
cache.cache.insert("example.com".to_string(), entry);
cache
.negative_entries
.insert("failed.com".to_string(), Instant::now());
assert_eq!(cache.len(), 1);
cache.clear();
assert!(cache.is_empty());
assert!(cache.negative_entries.is_empty());
}
#[tokio::test]
async fn test_resolve_caches_result() {
let mut cache = create_test_cache();
let result1 = cache.resolve("localhost", 80).await;
assert!(
result1.is_ok(),
"First resolve of localhost should succeed: {:?}",
result1.err()
);
let addrs1 = result1.unwrap();
assert!(
!addrs1.is_empty(),
"localhost should resolve to at least one address"
);
let result2 = cache.resolve("localhost", 80).await;
assert!(result2.is_ok(), "Second resolve should succeed from cache");
let addrs2 = result2.unwrap();
assert_eq!(
addrs1, addrs2,
"Cached result should match original resolution"
);
assert_eq!(cache.len(), 1);
}
#[test]
fn test_negative_cache_blocks_retry() {
let mut cache = DnsCache::with_ttl(300, 2);
cache.record_failure("test-host.invalid");
let result = cache.resolve_no_network("test-host.invalid", 80);
assert!(
result.is_err(),
"Lookup should be blocked by negative cache"
);
let err_msg = result.unwrap_err();
assert!(
err_msg.contains("recently failed"),
"Error should mention recent failure: {}",
err_msg
);
}
#[test]
fn test_purge_expired_removes_old() {
let mut cache = DnsCache::with_ttl(1, 60);
let expired_entry = DnsEntry {
hostname: "expired.example.com".to_string(),
addresses: vec!["10.0.0.1:80".parse().unwrap()],
resolved_at: Instant::now() - Duration::from_secs(5), ttl: Duration::from_secs(1),
ipv4_preferred: true,
};
cache
.cache
.insert("expired.example.com".to_string(), expired_entry);
let fresh_entry = DnsEntry {
hostname: "fresh.example.com".to_string(),
addresses: vec!["10.0.0.2:80".parse().unwrap()],
resolved_at: Instant::now(),
ttl: Duration::from_secs(3600), ipv4_preferred: true,
};
cache
.cache
.insert("fresh.example.com".to_string(), fresh_entry);
assert_eq!(cache.len(), 2, "Should have 2 entries before purge");
let removed = cache.purge_expired();
assert_eq!(removed, 1, "Should remove exactly 1 expired entry");
assert_eq!(cache.len(), 1, "Should have 1 entry remaining");
assert!(
cache.cache.contains_key("fresh.example.com"),
"Fresh entry should still exist"
);
assert!(
!cache.cache.contains_key("expired.example.com"),
"Expired entry should be removed"
);
}
#[tokio::test]
async fn test_ipv4_preferred_sorting() {
let mut cache = create_test_cache();
let result = cache.resolve("localhost", 8080).await;
assert!(
result.is_ok(),
"localhost resolution should succeed: {:?}",
result.err()
);
let addrs = result.unwrap();
let has_ipv4 = addrs.iter().any(|a| matches!(a.ip(), IpAddr::V4(_)));
let has_ipv6 = addrs.iter().any(|a| matches!(a.ip(), IpAddr::V6(_)));
if has_ipv4 && has_ipv6 {
let first_ipv4_pos = addrs
.iter()
.position(|a| matches!(a.ip(), IpAddr::V4(_)))
.unwrap();
let first_ipv6_pos = addrs
.iter()
.position(|a| matches!(a.ip(), IpAddr::V6(_)))
.unwrap();
assert!(
first_ipv4_pos < first_ipv6_pos,
"IPv4 addresses should come before IPv6 when preferred. Got order: {:?}",
addrs
);
}
cache.set_ipv4_preference(false);
let result2 = cache.force_refresh("localhost", 8080).await;
assert!(result2.is_ok(), "Force refresh should succeed");
}
#[test]
fn test_force_refresh_clears_cache_entry() {
let mut cache = create_test_cache();
let entry = DnsEntry {
hostname: "preloaded.com".to_string(),
addresses: vec!["192.168.1.1:443".parse().unwrap()],
resolved_at: Instant::now(),
ttl: Duration::from_secs(3600),
ipv4_preferred: true,
};
cache.cache.insert("preloaded.com".to_string(), entry);
assert_eq!(cache.len(), 1);
}
#[test]
fn test_default_impl() {
let cache = DnsCache::default();
assert!(cache.is_empty());
assert_eq!(cache.default_ttl(), Duration::from_secs(300));
}
#[tokio::test]
async fn test_async_dns_resolves_localhost() {
let mut cache = DnsCache::with_ttl(300, 60);
let result = cache.resolve("localhost", 80).await;
assert!(
result.is_ok(),
"localhost should resolve: {:?}",
result.err()
);
let addrs = result.unwrap();
assert!(!addrs.is_empty(), "should have at least one address");
assert_eq!(cache.len(), 1, "localhost should be cached after resolve");
}
#[tokio::test]
async fn test_dns_ttl_capped_at_default() {
let mut cache = DnsCache::with_ttl(60, 10); let _ = cache.resolve("localhost", 80).await;
if let Some(entry) = cache.cache.get("localhost") {
assert!(
entry.ttl <= Duration::from_secs(60),
"ttl should be capped at default, got {:?}",
entry.ttl
);
}
}
}