Merge #76: Implement timeout for socks connection

7493630 Implement timeout for socks connection (Riccardo Casatta)

Pull request description:

  To achieve this we have to include a modified socks library, with supports for
  timeout

  Remove the error `BothSocksAndTimeout` which isn't needed anymore

ACKs for top commit:
  afilini:
    ACK 7493630

Tree-SHA512: ec1b1eba739bd954ba7811eef603cf60f3ba29d7efcf98a9eadc8dcb9b67a49fce364760718faf7b860e78b7f9eebdc4da65783e7af47792573d32cd6c76ff10

Author: afilini

Cargo.toml

@@ -23,16 +23,23 @@ serde = { version = "^1.0", features = ["derive"] }
 serde_json = { version = "^1.0" }
 
 # Optional dependencies
-socks = { version = "0.3", optional = true }
 openssl = { version = "0.10", optional = true }
 rustls = { version = "0.20", optional = true, features = ["dangerous_configuration"] }
 webpki = { version = "0.22", optional = true }
 webpki-roots = { version = "0.22", optional = true }
 
+byteorder = { version = "1.0", optional = true }
+
+[target.'cfg(unix)'.dependencies]
+libc = { version = "0.2", optional = true }
+
+[target.'cfg(windows)'.dependencies]
+winapi = { version="0.3.9", features=["winsock2"], optional = true }
+
 [features]
 default = ["proxy", "use-rustls"]
 minimal = []
 debug-calls = []
-proxy = ["socks"]
+proxy = ["byteorder", "winapi", "libc"]
 use-rustls = ["webpki", "webpki-roots", "rustls"]
 use-openssl = ["openssl"]

src/client.rs

@@ -111,9 +111,12 @@ impl ClientType {
         if url.starts_with("ssl://") {
             let url = url.replacen("ssl://", "", 1);
             let client = match config.socks5() {
-                Some(socks5) => {
-                    RawClient::new_proxy_ssl(url.as_str(), config.validate_domain(), socks5)?
-                }
+                Some(socks5) => RawClient::new_proxy_ssl(
+                    url.as_str(),
+                    config.validate_domain(),
+                    socks5,
+                    config.timeout(),
+                )?,
                 None => {
                     RawClient::new_ssl(url.as_str(), config.validate_domain(), config.timeout())?
                 }
@@ -125,7 +128,11 @@ impl ClientType {
 
             Ok(match config.socks5().as_ref() {
                 None => ClientType::TCP(RawClient::new(url.as_str(), config.timeout())?),
-                Some(socks5) => ClientType::Socks5(RawClient::new_proxy(url.as_str(), socks5)?),
+                Some(socks5) => ClientType::Socks5(RawClient::new_proxy(
+                    url.as_str(),
+                    socks5,
+                    config.timeout(),
+                )?),
             })
         }
     }

src/config.rs

@@ -45,18 +45,12 @@ impl ConfigBuilder {
     /// Set the socks5 config if Some, it accept an `Option` because it's easier for the caller to use
     /// in a method chain
     pub fn socks5(mut self, socks5_config: Option<Socks5Config>) -> Result<Self, Error> {
-        if socks5_config.is_some() && self.config.timeout.is_some() {
-            return Err(Error::BothSocksAndTimeout);
-        }
         self.config.socks5 = socks5_config;
         Ok(self)
     }
 
     /// Sets the timeout
     pub fn timeout(mut self, timeout: Option<u8>) -> Result<Self, Error> {
-        if timeout.is_some() && self.config.socks5.is_some() {
-            return Err(Error::BothSocksAndTimeout);
-        }
         self.config.timeout = timeout.map(|t| Duration::from_secs(t as u64));
         Ok(self)
     }

src/lib.rs

@@ -31,13 +31,23 @@ extern crate openssl;
 extern crate rustls;
 extern crate serde;
 extern crate serde_json;
-#[cfg(any(feature = "default", feature = "proxy"))]
-extern crate socks;
+
 #[cfg(any(feature = "use-rustls", feature = "default"))]
 extern crate webpki;
 #[cfg(any(feature = "use-rustls", feature = "default"))]
 extern crate webpki_roots;
 
+#[cfg(any(feature = "default", feature = "proxy"))]
+extern crate byteorder;
+
+#[cfg(all(unix, any(feature = "default", feature = "proxy")))]
+extern crate libc;
+#[cfg(all(windows, any(feature = "default", feature = "proxy")))]
+extern crate winapi;
+
+#[cfg(any(feature = "default", feature = "proxy"))]
+mod socks;
+
 mod api;
 mod batch;
 

src/raw_client.rs

@@ -30,7 +30,7 @@ use rustls::{
 };
 
 #[cfg(any(feature = "default", feature = "proxy"))]
-use socks::{Socks5Stream, TargetAddr, ToTargetAddr};
+use crate::socks::{Socks5Stream, TargetAddr, ToTargetAddr};
 
 use stream::ClonableStream;
 
@@ -396,16 +396,20 @@ impl RawClient<ElectrumProxyStream> {
     pub fn new_proxy<T: ToTargetAddr>(
         target_addr: T,
         proxy: &crate::Socks5Config,
+        timeout: Option<Duration>,
     ) -> Result<Self, Error> {
-        let stream = match proxy.credentials.as_ref() {
+        let mut stream = match proxy.credentials.as_ref() {
             Some(cred) => Socks5Stream::connect_with_password(
                 &proxy.addr,
                 target_addr,
                 &cred.username,
                 &cred.password,
+                timeout,
             )?,
-            None => Socks5Stream::connect(&proxy.addr, target_addr)?,
+            None => Socks5Stream::connect(&proxy.addr, target_addr, timeout)?,
         };
+        stream.get_mut().set_read_timeout(timeout)?;
+        stream.get_mut().set_write_timeout(timeout)?;
 
         Ok(stream.into())
     }
@@ -418,18 +422,22 @@ impl RawClient<ElectrumProxyStream> {
         target_addr: T,
         validate_domain: bool,
         proxy: &crate::Socks5Config,
+        timeout: Option<Duration>,
     ) -> Result<RawClient<ElectrumSslStream>, Error> {
         let target = target_addr.to_target_addr()?;
 
-        let stream = match proxy.credentials.as_ref() {
+        let mut stream = match proxy.credentials.as_ref() {
             Some(cred) => Socks5Stream::connect_with_password(
                 &proxy.addr,
                 target_addr,
                 &cred.username,
                 &cred.password,
+                timeout,
             )?,
-            None => Socks5Stream::connect(&proxy.addr, target.clone())?,
+            None => Socks5Stream::connect(&proxy.addr, target.clone(), timeout)?,
         };
+        stream.get_mut().set_read_timeout(timeout)?;
+        stream.get_mut().set_write_timeout(timeout)?;
 
         RawClient::new_ssl_from_stream(target, validate_domain, stream.into_inner())
     }

src/socks/mod.rs

@@ -0,0 +1,162 @@
+//! SOCKS proxy clients
+
+use std::io;
+use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6, ToSocketAddrs};
+use std::vec;
+
+pub use self::v4::{Socks4Listener, Socks4Stream};
+pub use self::v5::{Socks5Datagram, Socks5Listener, Socks5Stream};
+
+mod v4;
+mod v5;
+mod writev;
+
+/// A description of a connection target.
+#[derive(Debug, Clone)]
+pub enum TargetAddr {
+    /// Connect to an IP address.
+    Ip(SocketAddr),
+    /// Connect to a fully qualified domain name.
+    ///
+    /// The domain name will be passed along to the proxy server and DNS lookup
+    /// will happen there.
+    Domain(String, u16),
+}
+
+impl ToSocketAddrs for TargetAddr {
+    type Iter = Iter;
+
+    fn to_socket_addrs(&self) -> io::Result<Iter> {
+        let inner = match *self {
+            TargetAddr::Ip(addr) => IterInner::Ip(Some(addr)),
+            TargetAddr::Domain(ref domain, port) => {
+                let it = (&**domain, port).to_socket_addrs()?;
+                IterInner::Domain(it)
+            }
+        };
+        Ok(Iter(inner))
+    }
+}
+
+enum IterInner {
+    Ip(Option<SocketAddr>),
+    Domain(vec::IntoIter<SocketAddr>),
+}
+
+/// An iterator over `SocketAddr`s associated with a `TargetAddr`.
+pub struct Iter(IterInner);
+
+impl Iterator for Iter {
+    type Item = SocketAddr;
+
+    fn next(&mut self) -> Option<SocketAddr> {
+        match self.0 {
+            IterInner::Ip(ref mut addr) => addr.take(),
+            IterInner::Domain(ref mut it) => it.next(),
+        }
+    }
+}
+
+/// A trait for objects that can be converted to `TargetAddr`.
+pub trait ToTargetAddr {
+    /// Converts the value of `self` to a `TargetAddr`.
+    fn to_target_addr(&self) -> io::Result<TargetAddr>;
+}
+
+impl ToTargetAddr for TargetAddr {
+    fn to_target_addr(&self) -> io::Result<TargetAddr> {
+        Ok(self.clone())
+    }
+}
+
+impl ToTargetAddr for SocketAddr {
+    fn to_target_addr(&self) -> io::Result<TargetAddr> {
+        Ok(TargetAddr::Ip(*self))
+    }
+}
+
+impl ToTargetAddr for SocketAddrV4 {
+    fn to_target_addr(&self) -> io::Result<TargetAddr> {
+        SocketAddr::V4(*self).to_target_addr()
+    }
+}
+
+impl ToTargetAddr for SocketAddrV6 {
+    fn to_target_addr(&self) -> io::Result<TargetAddr> {
+        SocketAddr::V6(*self).to_target_addr()
+    }
+}
+
+impl ToTargetAddr for (Ipv4Addr, u16) {
+    fn to_target_addr(&self) -> io::Result<TargetAddr> {
+        SocketAddrV4::new(self.0, self.1).to_target_addr()
+    }
+}
+
+impl ToTargetAddr for (Ipv6Addr, u16) {
+    fn to_target_addr(&self) -> io::Result<TargetAddr> {
+        SocketAddrV6::new(self.0, self.1, 0, 0).to_target_addr()
+    }
+}
+
+impl<'a> ToTargetAddr for (&'a str, u16) {
+    fn to_target_addr(&self) -> io::Result<TargetAddr> {
+        // try to parse as an IP first
+        if let Ok(addr) = self.0.parse::<Ipv4Addr>() {
+            return (addr, self.1).to_target_addr();
+        }
+
+        if let Ok(addr) = self.0.parse::<Ipv6Addr>() {
+            return (addr, self.1).to_target_addr();
+        }
+
+        Ok(TargetAddr::Domain(self.0.to_owned(), self.1))
+    }
+}
+
+impl<'a> ToTargetAddr for &'a str {
+    fn to_target_addr(&self) -> io::Result<TargetAddr> {
+        // try to parse as an IP first
+        if let Ok(addr) = self.parse::<SocketAddrV4>() {
+            return addr.to_target_addr();
+        }
+
+        if let Ok(addr) = self.parse::<SocketAddrV6>() {
+            return addr.to_target_addr();
+        }
+
+        // split the string by ':' and convert the second part to u16
+        let mut parts_iter = self.rsplitn(2, ':');
+        let port_str = match parts_iter.next() {
+            Some(s) => s,
+            None => {
+                return Err(io::Error::new(
+                    io::ErrorKind::InvalidInput,
+                    "invalid socket address",
+                ))
+            }
+        };
+
+        let host = match parts_iter.next() {
+            Some(s) => s,
+            None => {
+                return Err(io::Error::new(
+                    io::ErrorKind::InvalidInput,
+                    "invalid socket address",
+                ))
+            }
+        };
+
+        let port: u16 = match port_str.parse() {
+            Ok(p) => p,
+            Err(_) => {
+                return Err(io::Error::new(
+                    io::ErrorKind::InvalidInput,
+                    "invalid port value",
+                ))
+            }
+        };
+
+        (host, port).to_target_addr()
+    }
+}