From 2639193e9134a235db42cca16c8cff7f21f61661 Mon Sep 17 00:00:00 2001
From: Nuh <Ar.nazeh@gmail.com>
Date: Tue, 17 Sep 2024 21:14:25 +0300
Subject: [PATCH] feat(client): use non-zero ports resolved by dns resolvers
 (#148)

If a resolved address sets the port number to something besides `0`, and the port isn't otherwise explicitly asked for, the `HttpConnector` will now use that port. This allows custom resolvers that might lookup SRVB or HTTPSrr records that include a port number.

cc https://github.com/seanmonstar/reqwest/issues/2413
---
 src/client/legacy/connect/http.rs | 33 ++++++++++++++++++++++++++++++-
 1 file changed, 32 insertions(+), 1 deletion(-)

diff --git a/src/client/legacy/connect/http.rs b/src/client/legacy/connect/http.rs
index f8c9366..bcad03f 100644
--- a/src/client/legacy/connect/http.rs
+++ b/src/client/legacy/connect/http.rs
@@ -433,7 +433,8 @@ where
                 .map_err(ConnectError::dns)?;
             let addrs = addrs
                 .map(|mut addr| {
-                    addr.set_port(port);
+                    set_port(&mut addr, port, dst.port().is_some());
+
                     addr
                 })
                 .collect();
@@ -825,9 +826,19 @@ impl ConnectingTcp<'_> {
     }
 }
 
+/// Respect explicit ports in the URI, if none, either
+/// keep non `0` ports resolved from a custom dns resolver,
+/// or use the default port for the scheme.
+fn set_port(addr: &mut SocketAddr, host_port: u16, explicit: bool) {
+    if explicit || addr.port() == 0 {
+        addr.set_port(host_port)
+    };
+}
+
 #[cfg(test)]
 mod tests {
     use std::io;
+    use std::net::SocketAddr;
 
     use ::http::Uri;
 
@@ -836,6 +847,8 @@ mod tests {
     use super::super::sealed::{Connect, ConnectSvc};
     use super::{Config, ConnectError, HttpConnector};
 
+    use super::set_port;
+
     async fn connect<C>(
         connector: C,
         dst: Uri,
@@ -1234,4 +1247,22 @@ mod tests {
             panic!("test failed");
         }
     }
+
+    #[test]
+    fn test_set_port() {
+        // Respect explicit ports no matter what the resolved port is.
+        let mut addr = SocketAddr::from(([0, 0, 0, 0], 6881));
+        set_port(&mut addr, 42, true);
+        assert_eq!(addr.port(), 42);
+
+        // Ignore default  host port, and use the socket port instead.
+        let mut addr = SocketAddr::from(([0, 0, 0, 0], 6881));
+        set_port(&mut addr, 443, false);
+        assert_eq!(addr.port(), 6881);
+
+        // Use the default port if the resolved port is `0`.
+        let mut addr = SocketAddr::from(([0, 0, 0, 0], 0));
+        set_port(&mut addr, 443, false);
+        assert_eq!(addr.port(), 443);
+    }
 }