feat: Add DNS timings and resolution overrides (#360)

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

crates/yaak-http/src/client.rs

@@ -2,6 +2,8 @@ use crate::dns::LocalhostResolver;
 use crate::error::Result;
 use log::{debug, info, warn};
 use reqwest::{Client, Proxy, redirect};
+use std::sync::Arc;
+use yaak_models::models::DnsOverride;
 use yaak_tls::{ClientCertificateConfig, get_tls_config};
 
 #[derive(Clone)]
@@ -28,10 +30,14 @@ pub struct HttpConnectionOptions {
     pub validate_certificates: bool,
     pub proxy: HttpConnectionProxySetting,
     pub client_certificate: Option<ClientCertificateConfig>,
+    pub dns_overrides: Vec<DnsOverride>,
 }
 
 impl HttpConnectionOptions {
-    pub(crate) fn build_client(&self) -> Result<Client> {
+    /// Build a reqwest Client and return it along with the DNS resolver.
+    /// The resolver is returned separately so it can be configured per-request
+    /// to emit DNS timing events to the appropriate channel.
+    pub(crate) fn build_client(&self) -> Result<(Client, Arc<LocalhostResolver>)> {
         let mut client = Client::builder()
             .connection_verbose(true)
             .redirect(redirect::Policy::none())
@@ -40,15 +46,19 @@ impl HttpConnectionOptions {
             .no_brotli()
             .no_deflate()
             .referer(false)
-            .tls_info(true);
+            .tls_info(true)
+            // Disable connection pooling to ensure DNS resolution happens on each request
+            // This is needed so we can emit DNS timing events for each request
+            .pool_max_idle_per_host(0);
 
         // Configure TLS with optional client certificate
         let config =
             get_tls_config(self.validate_certificates, true, self.client_certificate.clone())?;
         client = client.use_preconfigured_tls(config);
 
-        // Configure DNS resolver
-        client = client.dns_resolver(LocalhostResolver::new());
+        // Configure DNS resolver - keep a reference to configure per-request
+        let resolver = LocalhostResolver::new(self.dns_overrides.clone());
+        client = client.dns_resolver(resolver.clone());
 
         // Configure proxy
         match self.proxy.clone() {
@@ -69,7 +79,7 @@ impl HttpConnectionOptions {
             self.client_certificate.is_some()
         );
 
-        Ok(client.build()?)
+        Ok((client.build()?, resolver))
     }
 }
 

crates/yaak-http/src/dns.rs

@@ -1,53 +1,185 @@
+use crate::sender::HttpResponseEvent;
 use hyper_util::client::legacy::connect::dns::{
     GaiResolver as HyperGaiResolver, Name as HyperName,
 };
+use log::info;
 use reqwest::dns::{Addrs, Name, Resolve, Resolving};
+use std::collections::HashMap;
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
 use std::str::FromStr;
 use std::sync::Arc;
+use std::time::Instant;
+use tokio::sync::{RwLock, mpsc};
 use tower_service::Service;
+use yaak_models::models::DnsOverride;
+
+/// Stores resolved addresses for a hostname override
+#[derive(Clone)]
+pub struct ResolvedOverride {
+    pub ipv4: Vec<Ipv4Addr>,
+    pub ipv6: Vec<Ipv6Addr>,
+}
 
 #[derive(Clone)]
 pub struct LocalhostResolver {
     fallback: HyperGaiResolver,
+    event_tx: Arc<RwLock<Option<mpsc::Sender<HttpResponseEvent>>>>,
+    overrides: Arc<HashMap<String, ResolvedOverride>>,
 }
 
 impl LocalhostResolver {
-    pub fn new() -> Arc<Self> {
+    pub fn new(dns_overrides: Vec<DnsOverride>) -> Arc<Self> {
         let resolver = HyperGaiResolver::new();
-        Arc::new(Self { fallback: resolver })
+
+        // Pre-parse DNS overrides into a lookup map
+        let mut overrides = HashMap::new();
+        for o in dns_overrides {
+            if !o.enabled {
+                continue;
+            }
+            let hostname = o.hostname.to_lowercase();
+
+            let ipv4: Vec<Ipv4Addr> =
+                o.ipv4.iter().filter_map(|s| s.parse::<Ipv4Addr>().ok()).collect();
+
+            let ipv6: Vec<Ipv6Addr> =
+                o.ipv6.iter().filter_map(|s| s.parse::<Ipv6Addr>().ok()).collect();
+
+            // Only add if at least one address is valid
+            if !ipv4.is_empty() || !ipv6.is_empty() {
+                overrides.insert(hostname, ResolvedOverride { ipv4, ipv6 });
+            }
+        }
+
+        Arc::new(Self {
+            fallback: resolver,
+            event_tx: Arc::new(RwLock::new(None)),
+            overrides: Arc::new(overrides),
+        })
+    }
+
+    /// Set the event sender for the current request.
+    /// This should be called before each request to direct DNS events
+    /// to the appropriate channel.
+    pub async fn set_event_sender(&self, tx: Option<mpsc::Sender<HttpResponseEvent>>) {
+        let mut guard = self.event_tx.write().await;
+        *guard = tx;
     }
 }
 
 impl Resolve for LocalhostResolver {
     fn resolve(&self, name: Name) -> Resolving {
         let host = name.as_str().to_lowercase();
+        let event_tx = self.event_tx.clone();
+        let overrides = self.overrides.clone();
 
+        info!("DNS resolve called for: {}", host);
+
+        // Check for DNS override first
+        if let Some(resolved) = overrides.get(&host) {
+            log::debug!("DNS override found for: {}", host);
+            let hostname = host.clone();
+            let mut addrs: Vec<SocketAddr> = Vec::new();
+
+            // Add IPv4 addresses
+            for ip in &resolved.ipv4 {
+                addrs.push(SocketAddr::new(IpAddr::V4(*ip), 0));
+            }
+
+            // Add IPv6 addresses
+            for ip in &resolved.ipv6 {
+                addrs.push(SocketAddr::new(IpAddr::V6(*ip), 0));
+            }
+
+            let addresses: Vec<String> = addrs.iter().map(|a| a.ip().to_string()).collect();
+
+            return Box::pin(async move {
+                // Emit DNS event for override
+                let guard = event_tx.read().await;
+                if let Some(tx) = guard.as_ref() {
+                    let _ = tx
+                        .send(HttpResponseEvent::DnsResolved {
+                            hostname,
+                            addresses,
+                            duration: 0,
+                            overridden: true,
+                        })
+                        .await;
+                }
+
+                Ok::<Addrs, Box<dyn std::error::Error + Send + Sync>>(Box::new(addrs.into_iter()))
+            });
+        }
+
+        // Check for .localhost suffix
         let is_localhost = host.ends_with(".localhost");
         if is_localhost {
+            let hostname = host.clone();
             // Port 0 is fine; reqwest replaces it with the URL's explicit
-            // port or the scheme’s default (80/443, etc.).
-            // (See docs note below.)
+            // port or the scheme's default (80/443, etc.).
             let addrs: Vec<SocketAddr> = vec![
                 SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
                 SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 0),
             ];
 
+            let addresses: Vec<String> = addrs.iter().map(|a| a.ip().to_string()).collect();
+
             return Box::pin(async move {
+                // Emit DNS event for localhost resolution
+                let guard = event_tx.read().await;
+                if let Some(tx) = guard.as_ref() {
+                    let _ = tx
+                        .send(HttpResponseEvent::DnsResolved {
+                            hostname,
+                            addresses,
+                            duration: 0,
+                            overridden: false,
+                        })
+                        .await;
+                }
+
                 Ok::<Addrs, Box<dyn std::error::Error + Send + Sync>>(Box::new(addrs.into_iter()))
             });
         }
 
+        // Fall back to system DNS
         let mut fallback = self.fallback.clone();
         let name_str = name.as_str().to_string();
+        let hostname = host.clone();
+
         Box::pin(async move {
-            match HyperName::from_str(&name_str) {
-                Ok(n) => fallback
-                    .call(n)
-                    .await
-                    .map(|addrs| Box::new(addrs) as Addrs)
-                    .map_err(|err| Box::new(err) as Box<dyn std::error::Error + Send + Sync>),
-                Err(e) => Err(Box::new(e) as Box<dyn std::error::Error + Send + Sync>),
+            let start = Instant::now();
+
+            let result = match HyperName::from_str(&name_str) {
+                Ok(n) => fallback.call(n).await,
+                Err(e) => return Err(Box::new(e) as Box<dyn std::error::Error + Send + Sync>),
+            };
+
+            let duration = start.elapsed().as_millis() as u64;
+
+            match result {
+                Ok(addrs) => {
+                    // Collect addresses for event emission
+                    let addr_vec: Vec<SocketAddr> = addrs.collect();
+                    let addresses: Vec<String> =
+                        addr_vec.iter().map(|a| a.ip().to_string()).collect();
+
+                    // Emit DNS event
+                    let guard = event_tx.read().await;
+                    if let Some(tx) = guard.as_ref() {
+                        let _ = tx
+                            .send(HttpResponseEvent::DnsResolved {
+                                hostname,
+                                addresses,
+                                duration,
+                                overridden: false,
+                            })
+                            .await;
+                    }
+
+                    Ok(Box::new(addr_vec.into_iter()) as Addrs)
+                }
+                Err(err) => Err(Box::new(err) as Box<dyn std::error::Error + Send + Sync>),
             }
         })
     }

crates/yaak-http/src/manager.rs

@@ -1,4 +1,5 @@
 use crate::client::HttpConnectionOptions;
+use crate::dns::LocalhostResolver;
 use crate::error::Result;
 use log::info;
 use reqwest::Client;
@@ -7,8 +8,15 @@ use std::sync::Arc;
 use std::time::{Duration, Instant};
 use tokio::sync::RwLock;
 
+/// A cached HTTP client along with its DNS resolver.
+/// The resolver is needed to set the event sender per-request.
+pub struct CachedClient {
+    pub client: Client,
+    pub resolver: Arc<LocalhostResolver>,
+}
+
 pub struct HttpConnectionManager {
-    connections: Arc<RwLock<BTreeMap<String, (Client, Instant)>>>,
+    connections: Arc<RwLock<BTreeMap<String, (CachedClient, Instant)>>>,
     ttl: Duration,
 }
 
@@ -20,21 +28,26 @@ impl HttpConnectionManager {
         }
     }
 
-    pub async fn get_client(&self, opt: &HttpConnectionOptions) -> Result<Client> {
+    pub async fn get_client(&self, opt: &HttpConnectionOptions) -> Result<CachedClient> {
         let mut connections = self.connections.write().await;
         let id = opt.id.clone();
 
         // Clean old connections
         connections.retain(|_, (_, last_used)| last_used.elapsed() <= self.ttl);
 
-        if let Some((c, last_used)) = connections.get_mut(&id) {
+        if let Some((cached, last_used)) = connections.get_mut(&id) {
             info!("Re-using HTTP client {id}");
             *last_used = Instant::now();
-            return Ok(c.clone());
+            return Ok(CachedClient {
+                client: cached.client.clone(),
+                resolver: cached.resolver.clone(),
+            });
         }
 
-        let c = opt.build_client()?;
-        connections.insert(id.into(), (c.clone(), Instant::now()));
-        Ok(c)
+        let (client, resolver) = opt.build_client()?;
+        let cached = CachedClient { client: client.clone(), resolver: resolver.clone() };
+        connections.insert(id.into(), (cached, Instant::now()));
+
+        Ok(CachedClient { client, resolver })
     }
 }

crates/yaak-http/src/sender.rs

@@ -45,6 +45,12 @@ pub enum HttpResponseEvent {
     ChunkReceived {
         bytes: usize,
     },
+    DnsResolved {
+        hostname: String,
+        addresses: Vec<String>,
+        duration: u64,
+        overridden: bool,
+    },
 }
 
 impl Display for HttpResponseEvent {
@@ -67,6 +73,19 @@ impl Display for HttpResponseEvent {
             HttpResponseEvent::HeaderDown(name, value) => write!(f, "< {}: {}", name, value),
             HttpResponseEvent::ChunkSent { bytes } => write!(f, "> [{} bytes sent]", bytes),
             HttpResponseEvent::ChunkReceived { bytes } => write!(f, "< [{} bytes received]", bytes),
+            HttpResponseEvent::DnsResolved { hostname, addresses, duration, overridden } => {
+                if *overridden {
+                    write!(f, "* DNS override {} -> {}", hostname, addresses.join(", "))
+                } else {
+                    write!(
+                        f,
+                        "* DNS resolved {} to {} ({}ms)",
+                        hostname,
+                        addresses.join(", "),
+                        duration
+                    )
+                }
+            }
         }
     }
 }
@@ -93,6 +112,9 @@ impl From<HttpResponseEvent> for yaak_models::models::HttpResponseEventData {
             HttpResponseEvent::HeaderDown(name, value) => D::HeaderDown { name, value },
             HttpResponseEvent::ChunkSent { bytes } => D::ChunkSent { bytes },
             HttpResponseEvent::ChunkReceived { bytes } => D::ChunkReceived { bytes },
+            HttpResponseEvent::DnsResolved { hostname, addresses, duration, overridden } => {
+                D::DnsResolved { hostname, addresses, duration, overridden }
+            }
         }
     }
 }