Fix socks proxy local DNS resolution not respected (#30619)

Jonathan LEI created 6 months ago

Closes #30618

Release Notes:

- Fixed SOCKS proxy incorrectly always uses remote DNS resolution.

Change summary

crates/client/src/proxy/socks_proxy.rs | 127 ++++++++++++++++++++++++---
1 file changed, 111 insertions(+), 16 deletions(-)

Detailed changes

crates/client/src/proxy/socks_proxy.rs

@@ -3,7 +3,10 @@
 use anyhow::{Context as _, Result};
 use http_client::Url;
 use tokio::net::TcpStream;
-use tokio_socks::tcp::{Socks4Stream, Socks5Stream};
+use tokio_socks::{
+    IntoTargetAddr, TargetAddr,
+    tcp::{Socks4Stream, Socks5Stream},
+};
 
 use super::AsyncReadWrite;
 
@@ -23,8 +26,14 @@ pub(super) struct Socks5Authorization<'a> {
 /// V4 allows idenfication using a user_id
 /// V5 allows authorization using a username and password
 pub(super) enum SocksVersion<'a> {
-    V4(Option<Socks4Identification<'a>>),
-    V5(Option<Socks5Authorization<'a>>),
+    V4 {
+        local_dns: bool,
+        identification: Option<Socks4Identification<'a>>,
+    },
+    V5 {
+        local_dns: bool,
+        authorization: Option<Socks5Authorization<'a>>,
+    },
 }
 
 pub(super) fn parse_socks_proxy<'t>(scheme: &str, proxy: &'t Url) -> SocksVersion<'t> {
@@ -33,13 +42,19 @@ pub(super) fn parse_socks_proxy<'t>(scheme: &str, proxy: &'t Url) -> SocksVersio
             "" => None,
             username => Some(Socks4Identification { user_id: username }),
         };
-        SocksVersion::V4(identification)
+        SocksVersion::V4 {
+            local_dns: scheme != "socks4a",
+            identification,
+        }
     } else {
         let authorization = proxy.password().map(|password| Socks5Authorization {
             username: proxy.username(),
             password,
         });
-        SocksVersion::V5(authorization)
+        SocksVersion::V5 {
+            local_dns: scheme != "socks5h",
+            authorization,
+        }
     }
 }
 
@@ -48,26 +63,58 @@ pub(super) async fn connect_socks_proxy_stream(
     socks_version: SocksVersion<'_>,
     rpc_host: (&str, u16),
 ) -> Result<Box<dyn AsyncReadWrite>> {
+    let rpc_host = rpc_host
+        .into_target_addr()
+        .context("Failed to parse target addr")?;
+
+    let local_dns = match &socks_version {
+        SocksVersion::V4 { local_dns, .. } => local_dns,
+        SocksVersion::V5 { local_dns, .. } => local_dns,
+    };
+    let rpc_host = match (rpc_host, local_dns) {
+        (TargetAddr::Domain(domain, port), true) => {
+            let ip_addr = tokio::net::lookup_host((domain.as_ref(), port))
+                .await
+                .with_context(|| format!("Failed to lookup domain {}", domain))?
+                .next()
+                .ok_or_else(|| anyhow::anyhow!("Failed to lookup domain {}", domain))?;
+            TargetAddr::Ip(ip_addr)
+        }
+        (rpc_host, _) => rpc_host,
+    };
+
     match socks_version {
-        SocksVersion::V4(None) => {
+        SocksVersion::V4 {
+            identification: None,
+            ..
+        } => {
             let socks = Socks4Stream::connect_with_socket(stream, rpc_host)
                 .await
                 .context("error connecting to socks")?;
             Ok(Box::new(socks))
         }
-        SocksVersion::V4(Some(Socks4Identification { user_id })) => {
+        SocksVersion::V4 {
+            identification: Some(Socks4Identification { user_id }),
+            ..
+        } => {
             let socks = Socks4Stream::connect_with_userid_and_socket(stream, rpc_host, user_id)
                 .await
                 .context("error connecting to socks")?;
             Ok(Box::new(socks))
         }
-        SocksVersion::V5(None) => {
+        SocksVersion::V5 {
+            authorization: None,
+            ..
+        } => {
             let socks = Socks5Stream::connect_with_socket(stream, rpc_host)
                 .await
                 .context("error connecting to socks")?;
             Ok(Box::new(socks))
         }
-        SocksVersion::V5(Some(Socks5Authorization { username, password })) => {
+        SocksVersion::V5 {
+            authorization: Some(Socks5Authorization { username, password }),
+            ..
+        } => {
             let socks = Socks5Stream::connect_with_password_and_socket(
                 stream, rpc_host, username, password,
             )
@@ -90,7 +137,13 @@ mod tests {
         let scheme = proxy.scheme();
 
         let version = parse_socks_proxy(scheme, &proxy);
-        assert!(matches!(version, SocksVersion::V4(None)))
+        assert!(matches!(
+            version,
+            SocksVersion::V4 {
+                local_dns: true,
+                identification: None
+            }
+        ))
     }
 
     #[test]
@@ -101,7 +154,25 @@ mod tests {
         let version = parse_socks_proxy(scheme, &proxy);
         assert!(matches!(
             version,
-            SocksVersion::V4(Some(Socks4Identification { user_id: "userid" }))
+            SocksVersion::V4 {
+                local_dns: true,
+                identification: Some(Socks4Identification { user_id: "userid" })
+            }
+        ))
+    }
+
+    #[test]
+    fn parse_socks4_with_remote_dns() {
+        let proxy = Url::parse("socks4a://proxy.example.com:1080").unwrap();
+        let scheme = proxy.scheme();
+
+        let version = parse_socks_proxy(scheme, &proxy);
+        assert!(matches!(
+            version,
+            SocksVersion::V4 {
+                local_dns: false,
+                identification: None
+            }
         ))
     }
 
@@ -111,7 +182,13 @@ mod tests {
         let scheme = proxy.scheme();
 
         let version = parse_socks_proxy(scheme, &proxy);
-        assert!(matches!(version, SocksVersion::V5(None)))
+        assert!(matches!(
+            version,
+            SocksVersion::V5 {
+                local_dns: true,
+                authorization: None
+            }
+        ))
     }
 
     #[test]
@@ -122,10 +199,28 @@ mod tests {
         let version = parse_socks_proxy(scheme, &proxy);
         assert!(matches!(
             version,
-            SocksVersion::V5(Some(Socks5Authorization {
-                username: "username",
-                password: "password"
-            }))
+            SocksVersion::V5 {
+                local_dns: true,
+                authorization: Some(Socks5Authorization {
+                    username: "username",
+                    password: "password"
+                })
+            }
+        ))
+    }
+
+    #[test]
+    fn parse_socks5_with_remote_dns() {
+        let proxy = Url::parse("socks5h://proxy.example.com:1080").unwrap();
+        let scheme = proxy.scheme();
+
+        let version = parse_socks_proxy(scheme, &proxy);
+        assert!(matches!(
+            version,
+            SocksVersion::V5 {
+                local_dns: false,
+                authorization: None
+            }
         ))
     }
 }