Add initial SOCKS5 support.

Cargo.toml

@@ -20,6 +20,7 @@ charset = ["encoding"]
 tls = ["rustls", "webpki", "webpki-roots"]
 native-certs = ["rustls-native-certs"]
 cookies = ["cookie"]
+socks-proxy = ["socks"]
 
 [dependencies]
 base64 = "0.12"
@@ -28,6 +29,7 @@ cookie = { version = "0.13", features = ["percent-encode"], optional = true}
 lazy_static = "1"
 qstring = "0.7"
 url = "2"
+socks = { version = "0.3.2", optional = true }
 rustls = { version = "0.17", optional = true, features = [] }
 webpki = { version = "0.21", optional = true }
 webpki-roots = { version = "0.19", optional = true }

src/proxy.rs

@@ -1,9 +1,10 @@
 use crate::error::Error;
 
-/// Kind of proxy connection (Basic, Digest, etc)
+/// Proxy protocol
-#[derive(Clone, Debug)]
+#[derive(Clone, Copy, Debug, PartialEq)]
-pub(crate) enum ProxyKind {
+pub enum Proto {
-    Basic,
+    HTTPConnect,
+    SOCKS5,
 }
 
 /// Proxy server definition
@@ -13,7 +14,7 @@ pub struct Proxy {
     pub(crate) port: u32,
     pub(crate) user: Option<String>,
     pub(crate) password: Option<String>,
-    pub(crate) kind: ProxyKind,
+    pub(crate) proto: Proto,
 }
 
 impl Proxy {
@@ -56,46 +57,68 @@ impl Proxy {
         }
     }
 
-    fn use_authorization(&self) -> bool {
+    pub(crate) fn use_authorization(&self) -> bool {
         self.user.is_some() && self.password.is_some()
     }
 
     pub fn new<S: AsRef<str>>(proxy: S) -> Result<Self, Error> {
-        let mut parts = proxy
+        let mut proxy_parts = proxy
             .as_ref()
+            .splitn(2, "://")
+            .collect::<Vec<&str>>()
+            .into_iter();
+
+        let proto = if proxy_parts.len() == 2 {
+            match proxy_parts.next() {
+                Some("http") => Proto::HTTPConnect,
+                Some("socks") => Proto::SOCKS5,
+                Some("socks5") => Proto::SOCKS5,
+                _ => return Err(Error::BadProxy),
+            }
+        } else {
+            Proto::HTTPConnect
+        };
+
+        let remaining_parts = proxy_parts.next();
+        if remaining_parts == None {
+            return Err(Error::BadProxy);
+        }
+
+        let mut creds_server_port_parts = remaining_parts
+            .unwrap()
             .rsplitn(2, '@')
             .collect::<Vec<&str>>()
             .into_iter()
             .rev();
 
-        let (user, password) = if parts.len() == 2 {
+        let (user, password) = if creds_server_port_parts.len() == 2 {
-            Proxy::parse_creds(&parts.next())?
+            Proxy::parse_creds(&creds_server_port_parts.next())?
         } else {
             (None, None)
         };
 
-        let (server, port) = Proxy::parse_address(&parts.next())?;
+        let (server, port) = Proxy::parse_address(&creds_server_port_parts.next())?;
 
         Ok(Self {
             server,
             user,
             password,
             port: port.unwrap_or(8080),
-            kind: ProxyKind::Basic,
+            proto,
         })
     }
 
-    pub fn connect<S: AsRef<str>>(&self, host: S, port: u16) -> String {
+    pub(crate) fn connect<S: AsRef<str>>(&self, host: S, port: u16) -> String {
         let authorization = if self.use_authorization() {
-            match self.kind {
-                ProxyKind::Basic => {
             let creds = base64::encode(&format!(
                 "{}:{}",
                 self.user.clone().unwrap_or_default(),
                 self.password.clone().unwrap_or_default()
             ));
-                    format!("Proxy-Authorization: basic {}\r\n", creds)
+
-                }
+            match self.proto {
+                Proto::HTTPConnect => format!("Proxy-Authorization: basic {}\r\n", creds),
+                Proto::SOCKS5 => String::new(),
             }
         } else {
             String::new()
@@ -131,14 +154,69 @@ Proxy-Connection: Keep-Alive\r\n\
 
 #[cfg(test)]
 mod tests {
+    use super::Proto;
     use super::Proxy;
 
     #[test]
-    fn parse_proxy() {
+    fn parse_proxy_fakeproto() {
+        assert!(Proxy::new("fakeproto://localhost").is_err());
+    }
+
+    #[test]
+    fn parse_proxy_http_user_pass_server_port() {
+        let proxy = Proxy::new("http://user:p@ssw0rd@localhost:9999").unwrap();
+        assert_eq!(proxy.user, Some(String::from("user")));
+        assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
+        assert_eq!(proxy.server, String::from("localhost"));
+        assert_eq!(proxy.port, 9999);
+        assert_eq!(proxy.proto, Proto::HTTPConnect);
+    }
+
+    #[cfg(feature = "socks-proxy")]
+    #[test]
+    fn parse_proxy_socks_user_pass_server_port() {
+        let proxy = Proxy::new("socks://user:p@ssw0rd@localhost:9999").unwrap();
+        assert_eq!(proxy.user, Some(String::from("user")));
+        assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
+        assert_eq!(proxy.server, String::from("localhost"));
+        assert_eq!(proxy.port, 9999);
+        assert_eq!(proxy.proto, Proto::SOCKS5);
+    }
+
+    #[cfg(feature = "socks-proxy")]
+    #[test]
+    fn parse_proxy_socks5_user_pass_server_port() {
+        let proxy = Proxy::new("socks5://user:p@ssw0rd@localhost:9999").unwrap();
+        assert_eq!(proxy.user, Some(String::from("user")));
+        assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
+        assert_eq!(proxy.server, String::from("localhost"));
+        assert_eq!(proxy.port, 9999);
+        assert_eq!(proxy.proto, Proto::SOCKS5);
+    }
+
+    #[test]
+    fn parse_proxy_user_pass_server_port() {
         let proxy = Proxy::new("user:p@ssw0rd@localhost:9999").unwrap();
         assert_eq!(proxy.user, Some(String::from("user")));
         assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
         assert_eq!(proxy.server, String::from("localhost"));
         assert_eq!(proxy.port, 9999);
     }
+
+    #[test]
+    fn parse_proxy_server_port() {
+        let proxy = Proxy::new("localhost:9999").unwrap();
+        assert_eq!(proxy.user, None);
+        assert_eq!(proxy.password, None);
+        assert_eq!(proxy.server, String::from("localhost"));
+        assert_eq!(proxy.port, 9999);
+    }
+
+    #[test]
+    fn parse_proxy_server() {
+        let proxy = Proxy::new("localhost").unwrap();
+        assert_eq!(proxy.user, None);
+        assert_eq!(proxy.password, None);
+        assert_eq!(proxy.server, String::from("localhost"));
+    }
 }

src/stream.rs

@@ -1,4 +1,4 @@
-use std::io::{Cursor, Read, Result as IoResult, Write};
+use std::io::{Cursor, ErrorKind, Read, Result as IoResult, Write};
 use std::net::SocketAddr;
 use std::net::TcpStream;
 use std::net::ToSocketAddrs;
@@ -9,6 +9,11 @@ use rustls::ClientSession;
 #[cfg(feature = "tls")]
 use rustls::StreamOwned;
 
+use crate::proxy::Proto;
+use crate::proxy::Proxy;
+#[cfg(feature = "socks-proxy")]
+use socks::{Socks5Stream, ToTargetAddr};
+
 use crate::error::Error;
 use crate::unit::Unit;
 
@@ -85,7 +90,7 @@ fn read_https(
 
 #[allow(deprecated)]
 fn is_close_notify(e: &std::io::Error) -> bool {
-    if e.kind() != std::io::ErrorKind::ConnectionAborted {
+    if e.kind() != ErrorKind::ConnectionAborted {
         return false;
     }
 
@@ -181,17 +186,27 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
         return Err(Error::DnsFailed(format!("No ip address for {}", hostname)));
     }
 
+    let proto = if let Some(ref proxy) = unit.proxy {
+        Some(proxy.proto)
+    } else {
+        None
+    };
+
     // pick first ip, or should we randomize?
     let sock_addr = ips[0];
 
     // connect with a configured timeout.
-    let mut stream = match unit.timeout_connect {
+    let mut stream = if Some(Proto::SOCKS5) == proto {
+        connect_socks5(unit.proxy.as_ref().unwrap(), &sock_addr, hostname, port)
+    } else {
+        match unit.timeout_connect {
             0 => TcpStream::connect(&sock_addr),
             _ => TcpStream::connect_timeout(
                 &sock_addr,
                 Duration::from_millis(unit.timeout_connect as u64),
             ),
         }
+    }
     .map_err(|err| Error::ConnectionFailed(format!("{}", err)))?;
 
     // rust's absurd api returns Err if we set 0.
@@ -212,6 +227,7 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
         stream.set_write_timeout(None).ok();
     }
 
+    if proto == Some(Proto::HTTPConnect) {
         if let Some(ref proxy) = unit.proxy {
             write!(stream, "{}", proxy.connect(hostname, port)).unwrap();
             stream.flush()?;
@@ -227,12 +243,62 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
                 }
             }
 
-        crate::Proxy::verify_response(&proxy_response)?;
+            Proxy::verify_response(&proxy_response)?;
+        }
     }
 
     Ok(stream)
 }
 
+#[cfg(feature = "socks-proxy")]
+fn connect_socks5(
+    proxy: &Proxy,
+    proxy_addr: &SocketAddr,
+    hostname: &str,
+    port: u16,
+) -> Result<TcpStream, std::io::Error> {
+    let host_addrs: Vec<SocketAddr> = format!("{}:{}", hostname, port)
+        .to_socket_addrs()
+        .map_err(|e| std::io::Error::new(ErrorKind::NotFound, format!("DNS failure: {}.", e)))?
+        .collect();
+
+    if host_addrs.is_empty() {
+        return Err(std::io::Error::new(
+            ErrorKind::NotFound,
+            format!("No ip address for {}.", proxy.server),
+        ));
+    }
+
+    let host_addr = host_addrs[0].to_target_addr()?;
+
+    let stream = if proxy.use_authorization() {
+        Socks5Stream::connect_with_password(
+            proxy_addr,
+            host_addr,
+            &proxy.user.as_ref().unwrap(),
+            &proxy.password.as_ref().unwrap(),
+        )?
+        .into_inner()
+    } else {
+        Socks5Stream::connect(proxy_addr, host_addr)?.into_inner()
+    };
+
+    Ok(stream)
+}
+
+#[cfg(not(feature = "socks-proxy"))]
+fn connect_socks5(
+    _proxy: &Proxy,
+    _proxy_addr: &SocketAddr,
+    _hostname: &str,
+    _port: u16,
+) -> Result<TcpStream, std::io::Error> {
+    Err(std::io::Error::new(
+        ErrorKind::Other,
+        "SOCKS5 feature disabled.",
+    ))
+}
+
 #[cfg(test)]
 pub(crate) fn connect_test(unit: &Unit) -> Result<Stream, Error> {
     use crate::test;