Add: socks5 support. It may have problems with DoT, will see.

crates/wtfnet-dns/Cargo.toml

@@ -7,9 +7,14 @@ edition = "2024"
 hickory-resolver = { version = "0.24", features = ["dns-over-tls", "dns-over-https", "dns-over-https-rustls", "dns-over-rustls", "native-certs"] }
 hickory-proto = "0.24"
 reqwest = { version = "0.11", features = ["rustls-tls", "socks"] }
+rustls = "0.21"
+rustls-native-certs = "0.6"
 serde = { version = "1", features = ["derive"] }
 thiserror = "2"
-tokio = { version = "1", features = ["time"] }
+tokio = { version = "1", features = ["io-util", "time"] }
+tokio-rustls = "0.24"
+tokio-socks = "0.5"
+url = "2"
 pnet = { version = "0.34", optional = true }
 
 [features]

crates/wtfnet-dns/src/lib.rs

@@ -8,12 +8,18 @@ use hickory_resolver::system_conf::read_system_conf;
 use hickory_proto::op::{Message, MessageType, Query};
 use hickory_proto::rr::Name;
 use reqwest::Proxy;
+use rustls::{Certificate, ClientConfig, RootCertStore, ServerName};
 use serde::{Deserialize, Serialize};
 use std::collections::BTreeSet;
 use std::net::{IpAddr, SocketAddr};
 use std::str::FromStr;
+use std::sync::Arc;
 use std::time::{Duration, Instant};
 use thiserror::Error;
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio_rustls::TlsConnector;
+use tokio_socks::tcp::Socks5Stream;
+use url::Url;
 
 #[cfg(feature = "pcap")]
 use pnet::datalink::{self, Channel, Config as DatalinkConfig, NetworkInterface};
@@ -164,11 +170,16 @@ pub async fn query(
 ) -> Result<DnsQueryReport, DnsError> {
     let record_type = parse_record_type(record_type)?;
     if let Some(proxy) = proxy {
-        if transport != DnsTransport::Doh {
-            return Err(DnsError::ProxyUnsupported(transport.to_string()));
-        }
         let server = server.ok_or_else(|| DnsError::MissingServer(transport.to_string()))?;
-        return doh_query_via_proxy(domain, record_type, server, timeout_ms, proxy).await;
+        return match transport {
+            DnsTransport::Doh => {
+                doh_query_via_proxy(domain, record_type, server, timeout_ms, proxy).await
+            }
+            DnsTransport::Dot => {
+                dot_query_via_proxy(domain, record_type, server, timeout_ms, proxy).await
+            }
+            _ => Err(DnsError::ProxyUnsupported(transport.to_string())),
+        };
     }
     let resolver = build_resolver(server.clone(), transport, timeout_ms)?;
     let start = Instant::now();
@@ -512,6 +523,151 @@ async fn doh_query_via_proxy(
     })
 }
 
+async fn dot_query_via_proxy(
+    domain: &str,
+    record_type: RecordType,
+    server: DnsServerTarget,
+    timeout_ms: u64,
+    proxy: String,
+) -> Result<DnsQueryReport, DnsError> {
+    let tls_name = server
+        .name
+        .clone()
+        .ok_or_else(|| DnsError::MissingTlsName("dot".to_string()))?;
+
+    let name = Name::from_ascii(domain)
+        .map_err(|err| DnsError::Resolver(format!("invalid domain: {err}")))?;
+    let mut message = Message::new();
+    message
+        .set_id(0)
+        .set_message_type(MessageType::Query)
+        .set_recursion_desired(true)
+        .add_query(Query::query(name, record_type));
+    let body = message
+        .to_vec()
+        .map_err(|err| DnsError::Resolver(err.to_string()))?;
+    if body.len() > u16::MAX as usize {
+        return Err(DnsError::Resolver("dns message too large".to_string()));
+    }
+
+    let connector = build_tls_connector()?;
+    let proxy_config = parse_socks5_proxy(&proxy)?;
+    let target = if proxy_config.remote_dns {
+        (tls_name.clone(), server.addr.port())
+    } else {
+        (server.addr.ip().to_string(), server.addr.port())
+    };
+    let timeout = Duration::from_millis(timeout_ms);
+    let tcp = tokio::time::timeout(
+        timeout,
+        Socks5Stream::connect(proxy_config.addr.as_str(), target),
+    )
+    .await
+    .map_err(|_| DnsError::Resolver("timeout".to_string()))?
+    .map_err(|err| DnsError::Proxy(err.to_string()))?
+    .into_inner();
+    let server_name = ServerName::try_from(tls_name.as_str())
+        .map_err(|_| DnsError::MissingTlsName(tls_name.clone()))?;
+    let mut stream = tokio::time::timeout(timeout, connector.connect(server_name, tcp))
+        .await
+        .map_err(|_| DnsError::Resolver("timeout".to_string()))?
+        .map_err(|err| DnsError::Resolver(err.to_string()))?;
+
+    let start = Instant::now();
+    let response_bytes = tokio::time::timeout(timeout, async {
+        let length = (body.len() as u16).to_be_bytes();
+        stream.write_all(&length).await?;
+        stream.write_all(&body).await?;
+        stream.flush().await?;
+        let mut len_buf = [0u8; 2];
+        stream.read_exact(&mut len_buf).await?;
+        let response_len = u16::from_be_bytes(len_buf) as usize;
+        let mut response = vec![0u8; response_len];
+        stream.read_exact(&mut response).await?;
+        Ok::<Vec<u8>, std::io::Error>(response)
+    })
+    .await
+    .map_err(|_| DnsError::Resolver("timeout".to_string()))?
+    .map_err(|err| DnsError::Resolver(err.to_string()))?;
+
+    let response =
+        Message::from_vec(&response_bytes).map_err(|err| DnsError::Resolver(err.to_string()))?;
+    let duration_ms = start.elapsed().as_millis();
+
+    let mut answers = Vec::new();
+    for record in response.answers() {
+        let ttl = record.ttl();
+        let name = record.name().to_string();
+        let record_type = record.record_type().to_string();
+        if let Some(data) = record.data() {
+            if let Some(data) = format_rdata(data) {
+                answers.push(DnsAnswer {
+                    name,
+                    record_type,
+                    ttl,
+                    data,
+                });
+            }
+        }
+    }
+
+    Ok(DnsQueryReport {
+        domain: domain.to_string(),
+        record_type: record_type.to_string(),
+        transport: DnsTransport::Dot.to_string(),
+        server: Some(server.addr.to_string()),
+        server_name: Some(tls_name),
+        proxy: Some(proxy),
+        rcode: response.response_code().to_string(),
+        answers,
+        duration_ms,
+    })
+}
+
+fn build_tls_connector() -> Result<TlsConnector, DnsError> {
+    let mut roots = RootCertStore::empty();
+    let store = rustls_native_certs::load_native_certs()
+        .map_err(|err| DnsError::Io(err.to_string()))?;
+    for cert in store {
+        roots
+            .add(&Certificate(cert.0))
+            .map_err(|err| DnsError::Resolver(err.to_string()))?;
+    }
+    let config = ClientConfig::builder()
+        .with_safe_defaults()
+        .with_root_certificates(roots)
+        .with_no_client_auth();
+    Ok(TlsConnector::from(Arc::new(config)))
+}
+
+struct Socks5Proxy {
+    addr: String,
+    remote_dns: bool,
+}
+
+fn parse_socks5_proxy(value: &str) -> Result<Socks5Proxy, DnsError> {
+    let url = Url::parse(value).map_err(|_| DnsError::Proxy(value.to_string()))?;
+    let scheme = url.scheme();
+    let remote_dns = match scheme {
+        "socks5" => false,
+        "socks5h" => true,
+        _ => return Err(DnsError::ProxyUnsupported(scheme.to_string())),
+    };
+    if !url.username().is_empty() || url.password().is_some() {
+        return Err(DnsError::Proxy("proxy auth not supported".to_string()));
+    }
+    let host = url
+        .host_str()
+        .ok_or_else(|| DnsError::Proxy(value.to_string()))?;
+    let port = url
+        .port_or_known_default()
+        .ok_or_else(|| DnsError::Proxy(value.to_string()))?;
+    Ok(Socks5Proxy {
+        addr: format!("{host}:{port}"),
+        remote_dns,
+    })
+}
+
 #[cfg(feature = "pcap")]
 fn select_interface(name: Option<&str>) -> Option<NetworkInterface> {
     let interfaces = datalink::interfaces();