diff --git a/Cargo.toml b/Cargo.toml index 06f0700..81067ee 100644 --- a/Cargo.toml +++ b/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 } diff --git a/src/proxy.rs b/src/proxy.rs index 641f805..5828337 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -1,9 +1,10 @@ use crate::error::Error; -/// Kind of proxy connection (Basic, Digest, etc) -#[derive(Clone, Debug)] -pub(crate) enum ProxyKind { - Basic, +/// Proxy protocol +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Proto { + HTTPConnect, + SOCKS5, } /// Proxy server definition @@ -13,7 +14,7 @@ pub struct Proxy { pub(crate) port: u32, pub(crate) user: Option, pub(crate) password: Option, - 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>(proxy: S) -> Result { - let mut parts = proxy + let mut proxy_parts = proxy .as_ref() + .splitn(2, "://") + .collect::>() + .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::>() .into_iter() .rev(); - let (user, password) = if parts.len() == 2 { - Proxy::parse_creds(&parts.next())? + let (user, password) = if creds_server_port_parts.len() == 2 { + 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>(&self, host: S, port: u16) -> String { + pub(crate) fn connect>(&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) - } + let creds = base64::encode(&format!( + "{}:{}", + self.user.clone().unwrap_or_default(), + self.password.clone().unwrap_or_default() + )); + + 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")); + } } diff --git a/src/stream.rs b/src/stream.rs index b4bf0e2..aa04a40 100644 --- a/src/stream.rs +++ b/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,16 +186,26 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result TcpStream::connect(&sock_addr), - _ => TcpStream::connect_timeout( - &sock_addr, - Duration::from_millis(unit.timeout_connect as u64), - ), + 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)))?; @@ -212,27 +227,78 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result Result { + let host_addrs: Vec = 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 { + Err(std::io::Error::new( + ErrorKind::Other, + "SOCKS5 feature disabled.", + )) +} + #[cfg(test)] pub(crate) fn connect_test(unit: &Unit) -> Result { use crate::test;