fix

binaries/geph5-exit/src/dns.rs

@@ -1,6 +1,6 @@
 use std::{
     fmt::Debug,
-    net::SocketAddr,
+    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
     str::FromStr,
     sync::{Arc, LazyLock},
     time::{Duration, Instant},
@@ -12,7 +12,7 @@ use bytes::Bytes;
 use globset::{Glob, GlobSet};
 use moka::future::Cache;
 use serde::{Deserialize, Serialize};
-use simple_dns::Packet;
+use simple_dns::{rdata::RData, Name, Packet, PacketFlag, Question, CLASS, QCLASS, QTYPE, TYPE};
 
 #[derive(Serialize, Deserialize, Clone, Debug, Copy, Default)]
 pub struct FilterOptions {
@@ -88,28 +88,6 @@ async fn parse_oisd(url: &str) -> anyhow::Result<GlobSet> {
     Ok(builder.build()?)
 }
 
-/// DNS resolve a name
-pub async fn dns_resolve(name: &str, filter: FilterOptions) -> anyhow::Result<Vec<SocketAddr>> {
-    static CACHE: LazyLock<Cache<String, Vec<SocketAddr>>> = LazyLock::new(|| {
-        Cache::builder()
-            .time_to_live(Duration::from_secs(240))
-            .build()
-    });
-    let (host, port) = name
-        .rsplit_once(":")
-        .context("could not split into host and port")?;
-    let port: u16 = port.parse()?;
-    filter.check_host(host).await?;
-    let addr = CACHE
-        .try_get_with(name.to_string(), async {
-            let choices = smol::net::resolve(name).await?;
-            anyhow::Ok(choices)
-        })
-        .await
-        .map_err(|e| anyhow::anyhow!(e))?;
-    Ok(addr)
-}
-
 pub async fn raw_dns_respond(req: Bytes, filter: FilterOptions) -> anyhow::Result<Bytes> {
     if let Ok(packet) = Packet::parse(&req) {
         for q in packet.questions.iter() {
@@ -135,6 +113,92 @@ pub async fn raw_dns_respond(req: Bytes, filter: FilterOptions) -> anyhow::Resul
     Ok(resp.bytes().await?)
 }
 
+pub async fn dns_resolve(name: &str, filter: FilterOptions) -> anyhow::Result<Vec<SocketAddr>> {
+    static CACHE: LazyLock<Cache<String, Vec<SocketAddr>>> = LazyLock::new(|| {
+        Cache::builder()
+            .time_to_live(Duration::from_secs(240))
+            .build()
+    });
+
+    // Split into "host:port"
+    let (host, port_str) = name
+        .rsplit_once(":")
+        .context("could not split into host:port")?;
+    let port: u16 = port_str.parse()?;
+
+    // Check filters
+    filter.check_host(host).await?;
+
+    // Use the cache to avoid repetitive lookups
+    let addrs = CACHE
+        .try_get_with(name.to_string(), async move {
+            // Build a DNS query for A and AAAA
+            let query_data = build_dns_query(host)?;
+
+            // Dispatch the query via DoH
+            let resp_bytes = raw_dns_respond(query_data, filter).await?;
+
+            // Parse out IP addresses
+            let ips = parse_dns_response(&resp_bytes, port)?;
+            anyhow::Ok(ips)
+        })
+        .await
+        .map_err(|e| anyhow::anyhow!(e))?;
+
+    Ok(addrs)
+}
+
+/// Build a simple DNS query packet for both A and AAAA:
+fn build_dns_query(host: &str) -> anyhow::Result<Bytes> {
+    let mut packet = Packet::new_query(rand::random());
+    packet.set_flags(PacketFlag::RECURSION_DESIRED);
+
+    // Ask for A
+    packet.questions.push(Question::new(
+        Name::new_unchecked(host),
+        QTYPE::TYPE(TYPE::A),
+        QCLASS::CLASS(CLASS::IN),
+        false,
+    ));
+
+    // Ask for AAAA
+    packet.questions.push(Question::new(
+        Name::new_unchecked(host),
+        QTYPE::TYPE(TYPE::AAAA),
+        QCLASS::CLASS(CLASS::IN),
+        false,
+    ));
+
+    let bytes = packet.build_bytes_vec()?;
+    Ok(bytes.into())
+}
+
+/// Parse a raw DNS response to gather all A/AAAA records.
+fn parse_dns_response(packet_data: &[u8], port: u16) -> anyhow::Result<Vec<SocketAddr>> {
+    let packet = Packet::parse(packet_data)?;
+    let mut addrs = Vec::new();
+
+    for answer in packet.answers {
+        match answer.rdata {
+            RData::A(ipv4) => addrs.push(SocketAddr::new(
+                IpAddr::V4(Ipv4Addr::from_bits(ipv4.address)),
+                port,
+            )),
+            RData::AAAA(ipv6) => addrs.push(SocketAddr::new(
+                IpAddr::V6(Ipv6Addr::from_bits(ipv6.address)),
+                port,
+            )),
+            _ => {}
+        }
+    }
+
+    if addrs.is_empty() {
+        anyhow::bail!("No A or AAAA records found in DNS response");
+    }
+
+    Ok(addrs)
+}
+
 // /// A udp-socket-efficient DNS responder.
 // pub async fn raw_dns_respond(req: Bytes, filter: FilterOptions) -> anyhow::Result<Bytes> {
 //     if let Ok(packet) = Packet::parse(&req) {
@@ -196,3 +260,24 @@ pub async fn raw_dns_respond(req: Bytes, filter: FilterOptions) -> anyhow::Resul
 //     };
 //     upload.race(download).await
 // }
+
+#[cfg(test)]
+mod tests {
+    use super::dns_resolve;
+
+    #[test]
+    fn resolve_google() {
+        smolscale::block_on(async move {
+            let res = dns_resolve(
+                "google.com:443",
+                super::FilterOptions {
+                    nsfw: false,
+                    ads: false,
+                },
+            )
+            .await
+            .unwrap();
+            eprintln!("{:?}", res);
+        });
+    }
+}