Connect with a timeout will try all socket addresses

Author: harmless-tech

README.md

@@ -33,11 +33,11 @@ use socks2::Socks4Stream;
 use socks2::Socks5Stream;
 use std::io::Write;
 
-let mut connection = Socks4Stream::connect(PROXY, &TARGET, "userid").unwrap();
+let mut connection = Socks4Stream::connect(PROXY, &TARGET, "userid", None).unwrap();
 let buf = [126_u8; 50]
 connection.write(&buf);
 
-let mut connection = Socks5Stream::connect(PROXY, &TARGET).unwrap();
+let mut connection = Socks5Stream::connect(PROXY, &TARGET, None).unwrap();
 let buf = [126_u8; 50]
 connection.write(&buf);
 ```
@@ -53,11 +53,11 @@ socks2 = { version = "0.4", default-features = false, features = ["bind"] }
 use socks2::Socks4Listener;
 use socks2::Socks5Listener;
 
-let mut connection = Socks4Listener::bin(PROXY, &TARGET, "userid")
+let mut connection = Socks4Listener::bind(PROXY, &TARGET, "userid", None)
     .unwrap()
     .accept();
 
-let mut connection = Socks5Listener::bind(PROXY, &TARGET)
+let mut connection = Socks5Listener::bind(PROXY, &TARGET, None)
     .unwrap()
     .accept();
 ```
@@ -73,7 +73,7 @@ socks2 = { version = "0.4", default-features = false, features = ["udp"] }
 use socks2::Socks5Datagram;
 use std::io::Write;
 
-let mut connection = Socks5Datagram::bind(PROXY, &TARGET).unwrap();
+let mut connection = Socks5Datagram::bind(PROXY, &TARGET, None).unwrap();
 let buf = [126_u8; 50]
 connection.send_to(&buf, &OTHER_ADDR);
 ```

src/error.rs

@@ -16,8 +16,8 @@ pub enum Error {
     InvalidPortValue { addr: String, port: String },
 
     // Socks4/Socks5
-    /// Could not resolve the first socket address.
-    NoResolveSocketAddr {},
+    /// Could not resolve any of the socket address.
+    NoResolveSocketAddrs {},
     /// Response from server had an invalid version byte.
     InvalidResponseVersion { version: u8 },
     /// Unknown response code
@@ -113,7 +113,7 @@ impl PartialEq for Error {
         peq!(
             InvalidSocksAddress,
             InvalidPortValue,
-            NoResolveSocketAddr,
+            NoResolveSocketAddrs,
             InvalidResponseVersion,
             UnknownResponseCode,
             ConnectionRefused,
@@ -155,7 +155,7 @@ impl From<Error> for io::Error {
         from_error!(
             (InvalidSocksAddress, InvalidInput),
             (InvalidPortValue, InvalidInput),
-            (NoResolveSocketAddr, InvalidInput),
+            (NoResolveSocketAddrs, InvalidInput),
             (InvalidResponseVersion, InvalidData),
             (UnknownResponseCode, Other),
             (ConnectionRefused, ConnectionRefused),
@@ -193,7 +193,7 @@ impl core::fmt::Display for Error {
             Self::InvalidPortValue { addr, port } => {
                 write!(f, "invalid port value '{port}' for '{addr}'")
             },
-            Self::NoResolveSocketAddr {} => write!(f, "could not resolve a socket address"),
+            Self::NoResolveSocketAddrs {} => write!(f, "could not resolve a socket address"),
             Self::InvalidResponseVersion { version } => write!(f, "invalid response version '{version}'"),
             Self::UnknownResponseCode { code } => write!(f, "unknown response code '{code}'"),
             Self::ConnectionRefused { code } => write!(f, "connection refused or the request was rejected or failed '{code}'"),

src/lib.rs

@@ -11,9 +11,10 @@
 extern crate alloc;
 
 use alloc::vec;
+use core::time::Duration;
 use std::{
     io,
-    net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6, ToSocketAddrs},
+    net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6, TcpStream, ToSocketAddrs},
 };
 
 #[cfg(feature = "client")]
@@ -193,6 +194,30 @@ impl ToTargetAddr for &str {
     }
 }
 
+fn tcp_stream_connect<T>(proxy: T, connect_timeout: Option<Duration>) -> io::Result<TcpStream>
+where
+    T: ToSocketAddrs,
+{
+    match connect_timeout {
+        None => TcpStream::connect(proxy),
+        Some(t) => {
+            // Timeout is applied to every SocketAddr try.
+            let mut addrs = proxy.to_socket_addrs()?;
+            let mut last_err = None;
+            for addr in &mut addrs {
+                match TcpStream::connect_timeout(&addr, t) {
+                    Ok(t) => return Ok(t),
+                    Err(err) => {
+                        last_err = Some(err);
+                        continue;
+                    }
+                }
+            }
+            Err(last_err.unwrap_or_else(|| Error::NoResolveSocketAddrs {}.into_io()))
+        }
+    }
+}
+
 #[cfg(test)]
 #[allow(clippy::unwrap_used)]
 mod test {

src/v4.rs

@@ -38,6 +38,7 @@ fn read_response(socket: &mut TcpStream) -> io::Result<SocketAddrV4> {
 #[cfg(feature = "client")]
 pub mod client {
     use crate::{
+        tcp_stream_connect,
         v4::{read_response, NULL_BYTE},
         Error, TargetAddr, ToTargetAddr,
     };
@@ -59,13 +60,16 @@ pub mod client {
     impl Socks4Stream {
         /// Connects to a target server through a SOCKS4 proxy.
         ///
-        /// # Note
-        ///
+        /// # Notes
         /// If `target` is a `TargetAddr::Domain`, the domain name will be forwarded
         /// to the proxy server using the SOCKS4A protocol extension. If the proxy
         /// server does not support SOCKS4A, consider performing the DNS lookup
         /// locally and passing a `TargetAddr::Ip`.
         ///
+        /// When using `connect_timeout` the duration will apply to every socket address
+        /// tried. Only the last connection error will be returned or
+        /// `io::Error(Error::NoResolveSocketAddrs)`.
+        ///
         /// # Errors
         /// - `io::Error(std::io::ErrorKind::*, socks2::Error::*?)`
         pub fn connect<T, U>(
@@ -92,18 +96,7 @@ pub mod client {
             T: ToSocketAddrs,
             U: ToTargetAddr,
         {
-            let mut socket = match connect_timeout {
-                None => TcpStream::connect(proxy)?,
-                // TODO: Connect timeout for each address until one works?
-                Some(t) => TcpStream::connect_timeout(
-                    &proxy
-                        .to_socket_addrs()?
-                        .next()
-                        .ok_or_else(|| Error::NoResolveSocketAddr {}.into_io())?,
-                    t,
-                )?,
-            };
-
+            let mut socket = tcp_stream_connect(proxy, connect_timeout)?;
             let target = target.to_target_addr()?;
 
             let mut packet = vec![];
@@ -218,6 +211,9 @@ pub mod bind {
         /// The proxy will filter incoming connections based on the value of
         /// `target`.
         ///
+        /// # Notes
+        /// See `Socks4Stream::connect()`.
+        ///
         /// # Errors
         /// - `io::Error(std::io::ErrorKind::*, socks2::Error::*?)`
         pub fn bind<T, U>(

src/v5.rs

@@ -130,6 +130,7 @@ impl Authentication<'_> {
 #[cfg(feature = "client")]
 pub mod client {
     use crate::{
+        tcp_stream_connect,
         v5::{read_response, write_addr, Authentication, MAX_ADDR_LEN},
         Error, TargetAddr, ToTargetAddr,
     };
@@ -150,6 +151,14 @@ pub mod client {
     impl Socks5Stream {
         /// Connects to a target server through a SOCKS5 proxy.
         ///
+        /// # Notes
+        /// If `target` is a `TargetAddr::Domain`, the domain name will be forwarded
+        /// to the proxy server to be resolved there.
+        ///
+        /// When using `connect_timeout` the duration will apply to every socket address
+        /// tried. Only the last connection error will be returned or
+        /// `io::Error(Error::NoResolveSocketAddrs)`.
+        ///
         /// # Errors
         /// - `io::Error(std::io::ErrorKind::*, socks2::Error::*?)`
         pub fn connect<T, U>(
@@ -167,6 +176,9 @@ pub mod client {
         /// Connects to a target server through a SOCKS5 proxy using given
         /// username and password.
         ///
+        /// # Notes
+        /// See `Socks5Stream::connect()`.
+        ///
         /// # Errors
         /// - `io::Error(std::io::ErrorKind::*, socks2::Error::*?)`
         pub fn connect_with_password<T, U>(
@@ -195,17 +207,7 @@ pub mod client {
             T: ToSocketAddrs,
             U: ToTargetAddr,
         {
-            let mut socket = match connect_timeout {
-                None => TcpStream::connect(proxy)?,
-                // TODO: Connect timeout for each address until one works?
-                Some(t) => TcpStream::connect_timeout(
-                    &proxy
-                        .to_socket_addrs()?
-                        .next()
-                        .ok_or_else(|| Error::NoResolveSocketAddr {}.into_io())?,
-                    t,
-                )?,
-            };
+            let mut socket = tcp_stream_connect(proxy, connect_timeout)?;
 
             let target = target.to_target_addr()?;
 
@@ -390,6 +392,9 @@ pub mod bind {
         /// The proxy will filter incoming connections based on the value of
         /// `target`.
         ///
+        /// # Notes
+        /// See `Socks5Stream::connect()`.
+        ///
         /// # Errors
         /// - `io::Error(std::io::ErrorKind::*, socks2::Error::*?)`
         pub fn bind<T, U>(
@@ -410,6 +415,9 @@ pub mod bind {
         /// The proxy will filter incoming connections based on the value of
         /// `target`.
         ///
+        /// # Notes
+        /// See `Socks5Stream::connect()`.
+        ///
         /// # Errors
         /// - `io::Error(std::io::ErrorKind::*, socks2::Error::*?)`
         pub fn bind_with_password<T, U>(
@@ -481,6 +489,9 @@ pub mod udp {
         /// Creates a UDP socket bound to the specified address which will have its
         /// traffic routed through the specified proxy.
         ///
+        /// # Notes
+        /// See `Socks5Stream::connect()`.
+        ///
         /// # Errors
         /// - `io::Error(std::io::ErrorKind::*, socks2::Error::*?)`
         pub fn bind<T, U>(proxy: T, addr: U, connect_timeout: Option<Duration>) -> io::Result<Self>
@@ -495,6 +506,9 @@ pub mod udp {
         /// traffic routed through the specified proxy. The given username and password
         /// is used to authenticate to the SOCKS proxy.
         ///
+        /// # Notes
+        /// See `Socks5Stream::connect()`.
+        ///
         /// # Errors
         /// - `io::Error(std::io::ErrorKind::*, socks2::Error::*?)`
         pub fn bind_with_password<T, U>(