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

crates/wtfnet-probe/Cargo.toml

@@ -12,3 +12,5 @@ tokio = { version = "1", features = ["net", "time"] }
 surge-ping = "0.8"
 wtfnet-geoip = { path = "../wtfnet-geoip" }
 libc = "0.2"
+tokio-socks = "0.5"
+url = "2"

crates/wtfnet-probe/src/lib.rs

@@ -20,6 +20,8 @@ use std::time::{Duration, Instant};
 use thiserror::Error;
 use tokio::net::{TcpStream, lookup_host};
 use tokio::time::timeout;
+use tokio_socks::tcp::Socks5Stream;
+use url::Url;
 use wtfnet_geoip::GeoIpRecord;
 
 #[derive(Debug, Error)]
@@ -28,6 +30,10 @@ pub enum ProbeError {
     Resolve(String),
     #[error("io error: {0}")]
     Io(String),
+    #[error("invalid proxy: {0}")]
+    InvalidProxy(String),
+    #[error("proxy error: {0}")]
+    Proxy(String),
     #[error("timeout")]
     Timeout,
     #[error("ping error: {0}")]
@@ -184,9 +190,30 @@ pub async fn tcp_ping(
     port: u16,
     count: u32,
     timeout_ms: u64,
+    proxy: Option<&str>,
+    prefer_ipv4: bool,
 ) -> Result<TcpPingReport, ProbeError> {
-    let addr = resolve_one(target).await?;
-    let socket_addr = SocketAddr::new(addr, port);
+    let (report_ip, target_host, proxy_addr) = if let Some(proxy) = proxy {
+        let proxy = parse_socks5_proxy(proxy)?;
+        if proxy.remote_dns {
+            (None, target.to_string(), proxy.addr)
+        } else {
+            let addr = if prefer_ipv4 {
+                resolve_one_prefer_ipv4(target).await?
+            } else {
+                resolve_one(target).await?
+            };
+            (Some(addr), addr.to_string(), proxy.addr)
+        }
+    } else {
+        let addr = if prefer_ipv4 {
+            resolve_one_prefer_ipv4(target).await?
+        } else {
+            resolve_one(target).await?
+        };
+        (Some(addr), addr.to_string(), String::new())
+    };
+    let socket_addr = report_ip.map(|addr| SocketAddr::new(addr, port));
     let timeout_dur = Duration::from_millis(timeout_ms);
     let mut results = Vec::new();
     let mut received = 0u32;
@@ -197,9 +224,27 @@ pub async fn tcp_ping(
     for seq in 0..count {
         let seq = seq as u16;
         let start = Instant::now();
-        let attempt = timeout(timeout_dur, TcpStream::connect(socket_addr)).await;
+        let attempt: Result<TcpStream, ProbeError> = if proxy.is_some() {
+            let target = (target_host.as_str(), port);
+            let stream = timeout(
+                timeout_dur,
+                Socks5Stream::connect(proxy_addr.as_str(), target),
+            )
+            .await
+            .map_err(|_| ProbeError::Timeout)?
+            .map_err(|err| ProbeError::Proxy(err.to_string()))?;
+            Ok(stream.into_inner())
+        } else {
+            timeout(
+                timeout_dur,
+                TcpStream::connect(socket_addr.expect("missing socket addr")),
+            )
+            .await
+            .map_err(|_| ProbeError::Timeout)?
+            .map_err(|err| ProbeError::Io(err.to_string()))
+        };
         match attempt {
-            Ok(Ok(_stream)) => {
+            Ok(_stream) => {
                 let rtt = start.elapsed().as_millis();
                 received += 1;
                 min = Some(min.map_or(rtt, |value: u128| value.min(rtt)));
@@ -211,27 +256,20 @@ pub async fn tcp_ping(
                     error: None,
                 });
             }
-            Ok(Err(err)) => {
+            Err(err) => {
                 results.push(TcpPingResult {
                     seq,
                     rtt_ms: None,
                     error: Some(err.to_string()),
                 });
             }
-            Err(_) => {
-                results.push(TcpPingResult {
-                    seq,
-                    rtt_ms: None,
-                    error: Some("timeout".to_string()),
-                });
-            }
         }
     }
 
     let summary = build_summary(count, received, min, max, sum);
     Ok(TcpPingReport {
         target: target.to_string(),
-        ip: Some(addr.to_string()),
+        ip: report_ip.map(|addr| addr.to_string()),
         geoip: None,
         port,
         timeout_ms,
@@ -389,6 +427,50 @@ async fn resolve_one(target: &str) -> Result<IpAddr, ProbeError> {
         .ok_or_else(|| ProbeError::Resolve("no address found".to_string()))
 }
 
+async fn resolve_one_prefer_ipv4(target: &str) -> Result<IpAddr, ProbeError> {
+    let mut iter = lookup_host((target, 0))
+        .await
+        .map_err(|err| ProbeError::Resolve(err.to_string()))?;
+    let mut fallback = None;
+    for addr in iter.by_ref() {
+        if addr.ip().is_ipv4() {
+            return Ok(addr.ip());
+        }
+        if fallback.is_none() {
+            fallback = Some(addr.ip());
+        }
+    }
+    fallback.ok_or_else(|| ProbeError::Resolve("no address found".to_string()))
+}
+
+struct Socks5Proxy {
+    addr: String,
+    remote_dns: bool,
+}
+
+fn parse_socks5_proxy(value: &str) -> Result<Socks5Proxy, ProbeError> {
+    let url = Url::parse(value).map_err(|_| ProbeError::InvalidProxy(value.to_string()))?;
+    let scheme = url.scheme();
+    let remote_dns = match scheme {
+        "socks5" => false,
+        "socks5h" => true,
+        _ => return Err(ProbeError::InvalidProxy(value.to_string())),
+    };
+    if !url.username().is_empty() || url.password().is_some() {
+        return Err(ProbeError::Proxy("proxy auth not supported".to_string()));
+    }
+    let host = url
+        .host_str()
+        .ok_or_else(|| ProbeError::InvalidProxy(value.to_string()))?;
+    let port = url
+        .port_or_known_default()
+        .ok_or_else(|| ProbeError::InvalidProxy(value.to_string()))?;
+    Ok(Socks5Proxy {
+        addr: format!("{host}:{port}"),
+        remote_dns,
+    })
+}
+
 fn tcp_connect_with_ttl(addr: SocketAddr, ttl: u8, timeout: Duration) -> Result<(), ProbeError> {
     let domain = match addr.ip() {
         IpAddr::V4(_) => Domain::IPV4,