Make SOCKS5 connection respect timeout_connect.

This commit is contained in:
sklv
2020-04-11 22:00:27 +00:00
committed by Martin Algesten
parent 05ce690ebd
commit 9928559067

View File

@@ -2,6 +2,7 @@ use std::io::{Cursor, ErrorKind, Read, Result as IoResult, Write};
use std::net::SocketAddr; use std::net::SocketAddr;
use std::net::TcpStream; use std::net::TcpStream;
use std::net::ToSocketAddrs; use std::net::ToSocketAddrs;
use std::thread;
use std::time::Duration; use std::time::Duration;
#[cfg(feature = "tls")] #[cfg(feature = "tls")]
@@ -9,10 +10,15 @@ use rustls::ClientSession;
#[cfg(feature = "tls")] #[cfg(feature = "tls")]
use rustls::StreamOwned; use rustls::StreamOwned;
use crate::proxy::Proto;
use crate::proxy::Proxy;
#[cfg(feature = "socks-proxy")] #[cfg(feature = "socks-proxy")]
use socks::{Socks5Stream, ToTargetAddr}; use socks::{Socks5Stream, ToTargetAddr};
#[cfg(feature = "socks-proxy")]
use std::sync::mpsc::channel;
#[cfg(feature = "socks-proxy")]
use std::sync::{Arc, Condvar, Mutex};
use crate::proxy::Proto;
use crate::proxy::Proxy;
use crate::error::Error; use crate::error::Error;
use crate::unit::Unit; use crate::unit::Unit;
@@ -197,7 +203,13 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
// connect with a configured timeout. // connect with a configured timeout.
let mut stream = if Some(Proto::SOCKS5) == proto { let mut stream = if Some(Proto::SOCKS5) == proto {
connect_socks5(unit.proxy.as_ref().unwrap(), &sock_addr, hostname, port) connect_socks5(
unit.proxy.to_owned().unwrap(),
unit.timeout_connect,
sock_addr,
hostname.to_string(),
port,
)
} else { } else {
match unit.timeout_connect { match unit.timeout_connect {
0 => TcpStream::connect(&sock_addr), 0 => TcpStream::connect(&sock_addr),
@@ -252,9 +264,10 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
#[cfg(feature = "socks-proxy")] #[cfg(feature = "socks-proxy")]
fn connect_socks5( fn connect_socks5(
proxy: &Proxy, proxy: Proxy,
proxy_addr: &SocketAddr, timeout_connect: u64,
hostname: &str, proxy_addr: SocketAddr,
hostname: String,
port: u16, port: u16,
) -> Result<TcpStream, std::io::Error> { ) -> Result<TcpStream, std::io::Error> {
let host_addrs: Vec<SocketAddr> = format!("{}:{}", hostname, port) let host_addrs: Vec<SocketAddr> = format!("{}:{}", hostname, port)
@@ -265,32 +278,81 @@ fn connect_socks5(
if host_addrs.is_empty() { if host_addrs.is_empty() {
return Err(std::io::Error::new( return Err(std::io::Error::new(
ErrorKind::NotFound, ErrorKind::NotFound,
format!("No ip address for {}.", proxy.server), format!("DNS lookup failed for {}:{}.", hostname, port),
)); ));
} }
let host_addr = host_addrs[0].to_target_addr()?; let stream = if timeout_connect > 0 {
let master_signal = Arc::new((Mutex::new(false), Condvar::new()));
let slave_signal = master_signal.clone();
let (tx, rx) = channel();
let host_addr = host_addrs[0];
thread::spawn(move || {
let (lock, cvar) = &*slave_signal;
if tx
.send(get_socks5_stream(&proxy, &proxy_addr, &host_addr))
.is_ok()
{
let mut done = lock.lock().unwrap();
*done = true;
cvar.notify_one();
}
});
let stream = if proxy.use_authorization() { let (lock, cvar) = &*master_signal;
Socks5Stream::connect_with_password( let done = lock.lock().unwrap();
proxy_addr,
host_addr, let done_result = cvar
&proxy.user.as_ref().unwrap(), .wait_timeout(done, Duration::from_millis(timeout_connect))
&proxy.password.as_ref().unwrap(), .unwrap();
)? let done = done_result.0;
.into_inner() if *done {
rx.recv().unwrap()?
} else {
return Err(std::io::Error::new(
ErrorKind::TimedOut,
format!(
"SOCKS5 proxy: {}:{} timed out connecting after {}ms.",
hostname, port, timeout_connect
),
));
}
} else { } else {
Socks5Stream::connect(proxy_addr, host_addr)?.into_inner() get_socks5_stream(&proxy, &proxy_addr, &host_addrs[0])?
}; };
Ok(stream) Ok(stream)
} }
#[cfg(feature = "socks-proxy")]
fn get_socks5_stream(
proxy: &Proxy,
proxy_addr: &SocketAddr,
host_addr: &SocketAddr,
) -> Result<TcpStream, std::io::Error> {
if proxy.use_authorization() {
let stream = Socks5Stream::connect_with_password(
proxy_addr,
host_addr.to_target_addr()?,
&proxy.user.as_ref().unwrap(),
&proxy.password.as_ref().unwrap(),
)?
.into_inner();
Ok(stream)
} else {
match Socks5Stream::connect(proxy_addr, host_addr.to_target_addr()?) {
Ok(socks_stream) => Ok(socks_stream.into_inner()),
Err(err) => Err(err),
}
}
}
#[cfg(not(feature = "socks-proxy"))] #[cfg(not(feature = "socks-proxy"))]
fn connect_socks5( fn connect_socks5(
_proxy: &Proxy, _proxy: Proxy,
_proxy_addr: &SocketAddr, _timeout_connect: u64,
_hostname: &str, _proxy_addr: SocketAddr,
_hostname: String,
_port: u16, _port: u16,
) -> Result<TcpStream, std::io::Error> { ) -> Result<TcpStream, std::io::Error> {
Err(std::io::Error::new( Err(std::io::Error::new(