use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::net::TcpStream;
use tokio::sync::Mutex;
use crate::net::error::NetError;
const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10);
const HAPPY_EYEBALLS_DELAY: Duration = Duration::from_millis(250);
const DNS_TTL: Duration = Duration::from_secs(300);
#[derive(Clone, Default)]
pub struct DnsCache {
inner: Arc<Mutex<HashMap<String, DnsEntry>>>,
}
struct DnsEntry {
addrs: Vec<SocketAddr>,
resolved_at: Instant,
}
impl DnsCache {
pub fn new() -> Self {
Self::default()
}
pub async fn resolve(&self, host: &str, port: u16) -> Result<Vec<SocketAddr>, NetError> {
{
let cache = self.inner.lock().await;
if let Some(entry) = cache.get(host) {
if entry.resolved_at.elapsed() < DNS_TTL {
return Ok(entry
.addrs
.iter()
.map(|a| SocketAddr::new(a.ip(), port))
.collect());
}
}
}
let addr_str = format!("{host}:{port}");
let addrs: Vec<SocketAddr> = tokio::net::lookup_host(&addr_str)
.await
.map_err(|e| NetError::Tcp(format!("DNS lookup failed for {host}: {e}")))?
.collect();
{
let mut cache = self.inner.lock().await;
cache.insert(
host.to_string(),
DnsEntry {
addrs: addrs.clone(),
resolved_at: Instant::now(),
},
);
}
Ok(addrs)
}
}
pub async fn connect(host: &str, port: u16, timeout: Duration) -> Result<TcpStream, NetError> {
connect_with_cache(host, port, timeout, None).await
}
pub async fn connect_via_proxy(
host: &str,
port: u16,
timeout: Duration,
dns_cache: Option<&DnsCache>,
proxy: Option<&crate::net::proxy::ProxyConfig>,
) -> Result<TcpStream, NetError> {
if let Some(proxy) = proxy {
return tokio::time::timeout(
timeout,
crate::net::proxy::connect(host, port, timeout, dns_cache, proxy),
)
.await
.map_err(|_| {
NetError::Tcp(format!(
"proxy connect to {host}:{port} timed out (timeout={}s)",
timeout.as_secs()
))
})?;
}
connect_with_cache(host, port, timeout, dns_cache).await
}
pub async fn connect_with_cache(
host: &str,
port: u16,
timeout: Duration,
dns_cache: Option<&DnsCache>,
) -> Result<TcpStream, NetError> {
let addrs: Vec<SocketAddr> = if let Some(cache) = dns_cache {
cache.resolve(host, port).await?
} else {
let addr_str = format!("{host}:{port}");
let resolved: Vec<SocketAddr> = tokio::net::lookup_host(&addr_str)
.await
.map_err(|e| NetError::Tcp(format!("DNS lookup failed for {host}: {e}")))?
.collect();
resolved
};
if addrs.is_empty() {
return Err(NetError::Tcp(format!("no addresses found for {host}")));
}
let mut ipv6: Vec<SocketAddr> = Vec::new();
let mut ipv4: Vec<SocketAddr> = Vec::new();
for addr in &addrs {
if addr.is_ipv6() {
ipv6.push(*addr);
} else {
ipv4.push(*addr);
}
}
let stream = tokio::time::timeout(timeout, happy_eyeballs(&ipv6, &ipv4))
.await
.map_err(|_| NetError::Tcp(format!("connection to {host}:{port} timed out")))??;
stream
.set_nodelay(true)
.map_err(|e| NetError::Tcp(format!("failed to set TCP_NODELAY: {e}")))?;
Ok(stream)
}
pub async fn connect_default(host: &str, port: u16) -> Result<TcpStream, NetError> {
connect(host, port, DEFAULT_TIMEOUT).await
}
async fn happy_eyeballs(ipv6: &[SocketAddr], ipv4: &[SocketAddr]) -> Result<TcpStream, NetError> {
if ipv6.is_empty() {
return try_addrs(ipv4).await;
}
if ipv4.is_empty() {
return try_addrs(ipv6).await;
}
tokio::select! {
result = try_addrs(ipv6) => {
match result {
Ok(stream) => Ok(stream),
Err(_) => try_addrs(ipv4).await,
}
}
_ = tokio::time::sleep(HAPPY_EYEBALLS_DELAY) => {
tokio::select! {
result = try_addrs(ipv6) => {
match result {
Ok(stream) => Ok(stream),
Err(_) => try_addrs(ipv4).await,
}
}
result = try_addrs(ipv4) => {
result
}
}
}
}
}
async fn try_addrs(addrs: &[SocketAddr]) -> Result<TcpStream, NetError> {
let mut last_err = None;
for addr in addrs {
match TcpStream::connect(addr).await {
Ok(stream) => return Ok(stream),
Err(e) => last_err = Some(e),
}
}
Err(NetError::Tcp(format!(
"all addresses failed: {}",
last_err.map_or_else(|| "no addresses".to_string(), |e| e.to_string())
)))
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
#[ignore] async fn connect_ipv4_httpbin() {
let stream = connect("httpbin.org", 443, Duration::from_secs(10)).await;
assert!(stream.is_ok(), "Failed to connect: {:?}", stream.err());
}
#[tokio::test]
#[ignore] async fn connect_ipv6_example_com() {
let stream = connect("example.com", 443, Duration::from_secs(10)).await;
assert!(stream.is_ok(), "Failed to connect: {:?}", stream.err());
}
#[tokio::test]
async fn connect_invalid_host() {
let result = connect(
"this.host.does.not.exist.example",
443,
Duration::from_secs(2),
)
.await;
assert!(result.is_err());
}
}