Add API to resolve hostnames and get OS's DNS servers

Cargo.toml

@@ -46,6 +46,7 @@ futures-util = { version = "0.3.30", features = ["sink"] }
 lru_time_cache = "0.11.11"
 internet-packet = { version = "0.2.0", features = ["smoltcp"] }
 data-encoding = "2.4.0"
+hickory-resolver = "0.24.1"
 
 [patch.crates-io]
 # tokio = { path = "../tokio/tokio" }
@@ -80,6 +81,7 @@ sysinfo = "0.29.10"
 env_logger = "0.11"
 rand = "0.8"
 criterion = "0.5.1"
+hickory-server = "0.24.1"
 
 
 [[bench]]

benches/dns.py

@@ -0,0 +1,26 @@
+import mitmproxy_rs
+import asyncio
+import socket
+
+async def main():
+    builder = mitmproxy_rs.DnsResolverBuilder()
+    builder.use_hosts_file(False)
+    builder.use_nameserver(["8.8.8.8"])
+    resolver = builder.build()
+
+    async def lookup(host: str):
+        try:
+            r = await resolver.lookup_ip(host)
+        except socket.gaierror as e:
+            print(f"{host=} {e=}")
+        else:
+            print(f"{host=} {r=}")
+
+    await lookup("example.com.")
+    await lookup("nxdomain.mitmproxy.org.")
+    await lookup("no-a-records.mitmproxy.org.")
+
+    print(f"{mitmproxy_rs.get_system_dns_servers()=}")
+
+
+asyncio.run(main())

mitmproxy-rs/mitmproxy_rs/__init__.pyi

@@ -124,3 +124,11 @@ class Process:
     @property
     def is_system(self) -> bool: ...
 
+# DNS resolution
+
+@final
+class DnsResolver:
+    def __init__(self, *, name_servers: list[str] | None = None, use_hosts_file: bool = True) -> None: ...
+    async def lookup_ip(self, host: str) -> list[str]: ...
+
+def get_system_dns_servers(): list[str]

mitmproxy-rs/src/dns_resolver.rs

@@ -0,0 +1,71 @@
+use mitmproxy::dns::{ResolveErrorKind, ResponseCode, DNS_SERVERS};
+use pyo3::exceptions::socket::gaierror;
+use pyo3::prelude::*;
+use pyo3::types::PyAny;
+use std::{net::IpAddr, net::SocketAddr, sync::Arc};
+
+/// A DNS resolver backed by [hickory-dns](https://github.com/hickory-dns/hickory-dns).
+/// This can serve as a replacement for `getaddrinfo` with configurable resolution behavior.
+///
+/// By default, the resolver will use the name servers configured by the operating system.
+/// It can optionally be configured to use custom name servers or ignore the hosts file.
+#[pyclass]
+pub struct DnsResolver(Arc<mitmproxy::dns::DnsResolver>);
+
+#[pymethods]
+impl DnsResolver {
+    #[new]
+    #[pyo3(signature = (*, name_servers, use_hosts_file=true))]
+    fn new(name_servers: Option<Vec<IpAddr>>, use_hosts_file: bool) -> PyResult<Self> {
+        let name_servers = name_servers.map(|ns| {
+            ns.into_iter()
+                .map(|ip| SocketAddr::from((ip, 53)))
+                .collect()
+        });
+        let resolver =
+            mitmproxy::dns::DnsResolver::new(name_servers, use_hosts_file).map_err(|e| {
+                pyo3::exceptions::PyRuntimeError::new_err(format!(
+                    "failed to create dns resolver: {}",
+                    e
+                ))
+            })?;
+        Ok(Self(Arc::new(resolver)))
+    }
+
+    /// Lookup the IPv4 and IPv6 addresses for a hostname.
+    ///
+    /// Raises `socket.gaierror` if the domain does not exist, has no records, or there is a general connectivity failure.
+    pub fn lookup_ip<'py>(&self, py: Python<'py>, host: String) -> PyResult<Bound<'py, PyAny>> {
+        let resolver = self.0.clone();
+        pyo3_asyncio_0_21::tokio::future_into_py(py, async move {
+            match resolver.lookup_ip(host).await {
+                Ok(resp) => Ok(resp
+                    .into_iter()
+                    .map(|ip| ip.to_string())
+                    .collect::<Vec<String>>()),
+                Err(e) => match *e.kind() {
+                    ResolveErrorKind::NoRecordsFound {
+                        response_code: ResponseCode::NXDomain,
+                        ..
+                    } => Err(gaierror::new_err("NXDOMAIN")),
+                    ResolveErrorKind::NoRecordsFound {
+                        response_code: ResponseCode::NoError,
+                        ..
+                    } => Err(gaierror::new_err("NOERROR")),
+                    _ => Err(gaierror::new_err(e.to_string())),
+                },
+            }
+        })
+    }
+}
+
+/// Returns the operating system's DNS servers as IP addresses.
+/// Raises a RuntimeError on unsupported platforms.
+///
+/// *Availability: Windows, Unix*
+#[pyfunction]
+pub fn get_system_dns_servers() -> PyResult<Vec<String>> {
+    DNS_SERVERS.clone().map_err(|e| {
+        pyo3::exceptions::PyRuntimeError::new_err(format!("failed to get dns servers: {}", e))
+    })
+}

mitmproxy-rs/src/lib.rs

@@ -5,6 +5,7 @@ use std::sync::RwLock;
 use once_cell::sync::Lazy;
 use pyo3::{exceptions::PyException, prelude::*};
 
+mod dns_resolver;
 mod process_info;
 mod server;
 mod stream;
@@ -59,6 +60,9 @@ pub fn mitmproxy_rs(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<process_info::Process>()?;
     m.add_function(wrap_pyfunction!(process_info::executable_icon, m)?)?;
 
+    m.add_class::<dns_resolver::DnsResolver>()?;
+    m.add_function(wrap_pyfunction!(dns_resolver::get_system_dns_servers, m)?)?;
+
     m.add_class::<stream::Stream>()?;
 
     // Import platform-specific modules here so that missing dependencies are raising immediately.

mitmproxy-rs/stubtest-allowlist.txt

@@ -2,4 +2,5 @@ mitmproxy_rs.mitmproxy_rs
 mitmproxy_rs._pyinstaller.hook-mitmproxy_rs
 mitmproxy_rs._pyinstaller.hook-mitmproxy_windows
 mitmproxy_rs._pyinstaller.hook-mitmproxy_macos
-mitmproxy_rs.T
+mitmproxy_rs.T
+mitmproxy_rs.DnsResolver.__init__

src/dns.rs

@@ -0,0 +1,145 @@
+use hickory_resolver::config::LookupIpStrategy;
+use hickory_resolver::error::ResolveResult;
+use hickory_resolver::lookup_ip::LookupIp;
+use hickory_resolver::system_conf::read_system_conf;
+use hickory_resolver::TokioAsyncResolver;
+use once_cell::sync::Lazy;
+use std::net::IpAddr;
+use std::net::SocketAddr;
+
+use hickory_resolver::config::NameServerConfig;
+use hickory_resolver::config::Protocol;
+use hickory_resolver::config::ResolverConfig;
+pub use hickory_resolver::error::ResolveErrorKind;
+pub use hickory_resolver::proto::op::ResponseCode;
+
+pub static DNS_SERVERS: Lazy<ResolveResult<Vec<String>>> = Lazy::new(|| {
+    let (config, _opts) = read_system_conf()?;
+    Ok(config
+        .name_servers()
+        .iter()
+        .filter(|ns| ns.protocol == Protocol::Udp)
+        .map(|ns| ns.socket_addr.ip().to_string())
+        .collect::<Vec<String>>())
+});
+
+pub struct DnsResolver(TokioAsyncResolver);
+
+impl DnsResolver {
+    pub fn new(name_servers: Option<Vec<SocketAddr>>, use_hosts_file: bool) -> ResolveResult<Self> {
+        let (config, mut opts) = if let Some(ns) = name_servers {
+            // Try to get opts from system, but fall back gracefully if that is unavailable.
+            let opts = read_system_conf().map(|r| r.1).unwrap_or_default();
+
+            let mut conf = ResolverConfig::new();
+            for addr in ns.into_iter() {
+                conf.add_name_server(NameServerConfig::new(addr, Protocol::Udp));
+                conf.add_name_server(NameServerConfig::new(addr, Protocol::Tcp));
+            }
+            (conf, opts)
+        } else {
+            read_system_conf()?
+        };
+        opts.use_hosts_file = use_hosts_file;
+        opts.ip_strategy = LookupIpStrategy::Ipv4AndIpv6;
+        Ok(Self(TokioAsyncResolver::tokio(config, opts)))
+    }
+
+    pub async fn lookup_ip(&self, host: String) -> ResolveResult<Vec<IpAddr>> {
+        self.0.lookup_ip(host).await.map(_interleave_addrinfos)
+    }
+}
+
+fn _interleave_addrinfos(lookup_ip: LookupIp) -> Vec<IpAddr> {
+    let (mut ipv4_addrs, mut ipv6_addrs): (Vec<IpAddr>, Vec<IpAddr>) =
+        lookup_ip.into_iter().partition(|addr| addr.is_ipv4());
+
+    let mut interleaved: Vec<IpAddr> = Vec::with_capacity(ipv4_addrs.len() + ipv6_addrs.len());
+
+    while let Some(ipv4) = ipv4_addrs.pop() {
+        interleaved.push(ipv4);
+        if let Some(ipv6) = ipv6_addrs.pop() {
+            interleaved.push(ipv6);
+        }
+    }
+    interleaved.append(&mut ipv6_addrs);
+    interleaved
+}
+
+#[cfg(test)]
+mod tests {
+
+    use hickory_resolver::config::NameServerConfig;
+
+    use hickory_server::proto::rr::rdata::{A, AAAA};
+    use hickory_server::proto::rr::{DNSClass, Name, RData, Record};
+    use std::net::SocketAddr;
+    use std::str::FromStr;
+    use std::sync::Arc;
+
+    use hickory_server::authority::ZoneType;
+    use hickory_server::store::in_memory::InMemoryAuthority;
+
+    use super::*;
+
+    #[test]
+    fn dns_servers() {
+        assert!(DNS_SERVERS.as_ref().is_ok_and(|s| !s.is_empty()))
+    }
+
+    #[tokio::test]
+    async fn resolver() -> anyhow::Result<()> {
+        let listen_addr = test_server().await?;
+
+        let mut config = ResolverConfig::new();
+        config.add_name_server(NameServerConfig::new(listen_addr, Protocol::Udp));
+        let results = DnsResolver::new(Some(vec![listen_addr]), false)?
+            .lookup_ip("example.com.".to_string())
+            .await?;
+
+        assert_eq!(
+            results,
+            vec![
+                IpAddr::from_str("93.184.215.14")?,
+                IpAddr::from_str("2606:2800:21f:cb07:6820:80da:af6b:8b2c")?,
+            ]
+        );
+
+        Ok(())
+    }
+
+    async fn test_server() -> anyhow::Result<SocketAddr> {
+        let sock = tokio::net::UdpSocket::bind("127.0.0.1:0").await?;
+        let listen_addr = sock.local_addr()?;
+
+        let origin: Name = Name::parse("example.com.", None).unwrap();
+        let mut records = InMemoryAuthority::empty(origin.clone(), ZoneType::Primary, false);
+        records.upsert_mut(
+            Record::from_rdata(origin.clone(), 86400, RData::A(A::new(93, 184, 215, 14)))
+                .set_dns_class(DNSClass::IN)
+                .clone(),
+            0,
+        );
+        records.upsert_mut(
+            Record::from_rdata(
+                origin,
+                86400,
+                RData::AAAA(AAAA::new(
+                    0x2606, 0x2800, 0x21f, 0xcb07, 0x6820, 0x80da, 0xaf6b, 0x8b2c,
+                )),
+            )
+            .set_dns_class(DNSClass::IN)
+            .clone(),
+            0,
+        );
+
+        let mut catalog = hickory_server::authority::Catalog::new();
+        catalog.upsert(Name::root().into(), Box::new(Arc::new(records)));
+
+        let mut server = hickory_server::ServerFuture::new(catalog);
+        server.register_socket(sock);
+
+        tokio::spawn(async move { server.block_until_done().await });
+        Ok(listen_addr)
+    }
+}

src/lib.rs

@@ -1,6 +1,7 @@
 pub use network::MAX_PACKET_SIZE;
 
 pub mod certificates;
+pub mod dns;
 pub mod intercept_conf;
 pub mod ipc;
 pub mod messages;