use std::sync::Arc;
use std::time::{Duration, Instant};
use anyhow::Context;
use async_trait::async_trait;
use log::{debug, info};
use socks5_proto::Address;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite};
use tokio::net::TcpStream;
use tokio::time::timeout;
use crate::address_list::{DirectList, ProxyList};
use crate::proto::padding::Padding;
use crate::proto::trojan;
use crate::proto::trojan::Command;
use crate::tls::make_server_name;
use crate::tls::make_tls_connector;
pub trait AsyncStream: AsyncRead + AsyncWrite + Unpin + Send {}
impl<T: AsyncRead + AsyncWrite + Unpin + Send> AsyncStream for T {}
#[async_trait]
pub trait Dial: Send + Sync {
async fn dial(&self, addr: Address) -> anyhow::Result<Box<dyn AsyncStream>>;
}
pub struct DirectDial {
connect_timeout: Duration,
}
impl DirectDial {
pub fn new(connect_timeout: Duration) -> Self {
Self { connect_timeout }
}
}
#[async_trait]
impl Dial for DirectDial {
async fn dial(&self, addr: Address) -> anyhow::Result<Box<dyn AsyncStream>> {
let stream: TcpStream = match addr {
Address::DomainAddress(domain, port) => {
let domain = String::from_utf8_lossy(&domain);
timeout(
self.connect_timeout,
TcpStream::connect((domain.as_ref(), port)),
)
.await
.context(format!("connect {}:{} timeout", domain, port))?
.context(format!("connect {}:{} failed", domain, port))
}
Address::SocketAddress(socket_addr) => {
timeout(self.connect_timeout, TcpStream::connect(socket_addr))
.await
.context(format!("connect {} timeout", socket_addr))?
.context(format!("connect {} failed", socket_addr))
}
}?;
Ok(Box::new(stream))
}
}
pub struct TrojanDial {
remote_addr: String,
hash: String,
insecure: bool,
padding: bool,
connect_timeout: Duration,
}
impl TrojanDial {
pub fn new(
remote_addr: String,
hash: String,
insecure: bool,
padding: bool,
connect_timeout: Duration,
) -> Self {
Self {
remote_addr,
hash,
insecure,
padding,
connect_timeout,
}
}
}
#[async_trait]
impl Dial for TrojanDial {
async fn dial(&self, addr: Address) -> anyhow::Result<Box<dyn AsyncStream>> {
let remote_ts = timeout(self.connect_timeout, TcpStream::connect(&self.remote_addr))
.await
.context(format!("connect {} timeout", self.remote_addr))?
.context(format!("connect {} failed", self.remote_addr))?;
let server_name = make_server_name(self.remote_addr.as_str())?;
let mut remote_ts_ssl = make_tls_connector(self.insecure)
.connect(server_name, remote_ts)
.await
.context("trojan can't connect tls")?;
if self.padding {
let req = trojan::Request::new(self.hash.clone(), Command::Padding, addr);
req.write_to(&mut remote_ts_ssl).await?;
Padding::read_from(&mut remote_ts_ssl).await?;
} else {
let req = trojan::Request::new(self.hash.clone(), Command::Connect, addr);
req.write_to(&mut remote_ts_ssl).await?;
}
Ok(Box::new(remote_ts_ssl))
}
}
#[cfg(feature = "websocket")]
pub struct WebSocketDial {
remote_addr: String,
hash: String,
insecure: bool,
padding: bool,
connect_timeout: Duration,
}
#[cfg(feature = "websocket")]
impl WebSocketDial {
pub fn new(
remote_addr: String,
hash: String,
insecure: bool,
padding: bool,
connect_timeout: Duration,
) -> Self {
Self {
remote_addr,
hash,
insecure,
padding,
connect_timeout,
}
}
}
#[cfg(feature = "websocket")]
#[async_trait]
impl Dial for WebSocketDial {
async fn dial(&self, addr: Address) -> anyhow::Result<Box<dyn AsyncStream>> {
use crate::stream::websocket::WebSocketCopyStream;
use crate::tls::make_tls_client_config;
use bytes::BytesMut;
use futures::SinkExt;
use futures::StreamExt;
use tokio_tungstenite::connect_async_tls_with_config;
use tokio_tungstenite::tungstenite::Message;
let (mut ws, _) = timeout(
self.connect_timeout,
connect_async_tls_with_config(
&self.remote_addr,
None,
false,
Some(tokio_tungstenite::Connector::Rustls(Arc::new(
make_tls_client_config(self.insecure),
))),
),
)
.await
.context(format!("websocket connect {} timeout", self.remote_addr))?
.context(format!("websocket connect {} failed", self.remote_addr))?;
if self.padding {
let mut buf = BytesMut::new();
let req = trojan::Request::new(self.hash.clone(), Command::Padding, addr);
req.write_to_buf(&mut buf);
ws.send(Message::Binary(buf.freeze()))
.await
.context("websocket can't send")?;
ws.flush().await?;
let _ = ws.next().await;
} else {
let mut buf = BytesMut::new();
let req = trojan::Request::new(self.hash.clone(), Command::Connect, addr);
req.write_to_buf(&mut buf);
ws.send(Message::Binary(buf.freeze()))
.await
.context("websocket can't send")?;
ws.flush().await?;
}
Ok(Box::new(WebSocketCopyStream::new(ws)))
}
}
const EXTRA_TIMEOUT_MS: u64 = 200;
pub struct SmartDial {
direct: Box<dyn Dial>,
proxy: Box<dyn Dial>,
proxy_list: Arc<ProxyList>,
direct_list: Arc<DirectList>,
connect_timeout: Duration,
}
impl SmartDial {
pub fn new(
direct: Box<dyn Dial>,
proxy: Box<dyn Dial>,
proxy_list: Arc<ProxyList>,
direct_list: Arc<DirectList>,
connect_timeout: Duration,
) -> Self {
Self {
direct,
proxy,
proxy_list,
direct_list,
connect_timeout,
}
}
async fn handle_proxy_result(
&self,
proxy_res: anyhow::Result<Box<dyn AsyncStream>>,
addr: &Address,
elapsed: Duration,
) -> anyhow::Result<Box<dyn AsyncStream>> {
match proxy_res {
Ok(mut stream) => {
let remaining = self
.connect_timeout
.saturating_add(Duration::from_millis(EXTRA_TIMEOUT_MS))
.saturating_sub(elapsed);
let check_timeout = remaining.min(Duration::from_millis(EXTRA_TIMEOUT_MS));
if !is_stream_closed(&mut stream, check_timeout).await {
self.proxy_list.add_address(addr);
info!(
"Proxy Connect to: {addr} : record [{:.3}s]",
elapsed.as_secs_f64()
);
} else {
info!(
"Proxy Connect to: {addr} : unrecord [{:.3}s]",
elapsed.as_secs_f64()
);
}
Ok(stream)
}
Err(e) => Err(e),
}
}
}
#[async_trait]
impl Dial for SmartDial {
async fn dial(&self, addr: Address) -> anyhow::Result<Box<dyn AsyncStream>> {
if self.proxy_list.contains_address(&addr) {
debug!("Address {:?} is in proxy list, using proxy dial", addr);
info!("Proxy Connect to: {addr}");
return self.proxy.dial(addr).await;
}
if self.direct_list.contains_address(&addr) {
debug!("Address {:?} is in direct list, using direct dial", addr);
match self.direct.dial(addr.clone()).await {
Ok(stream) => {
info!("Direct Connect to: {addr}");
return Ok(stream);
}
Err(e) => {
debug!(
"Direct dial failed for {:?}: {}, removing from direct list",
addr, e
);
self.direct_list.remove_address(&addr);
}
}
}
let start = Instant::now();
let direct_fut = self.direct.dial(addr.clone());
let proxy_fut = self.proxy.dial(addr.clone());
tokio::pin!(direct_fut);
tokio::pin!(proxy_fut);
tokio::select! {
direct_res = &mut direct_fut => {
match direct_res {
Ok(stream) => {
self.direct_list.add_address(&addr);
info!("Direct Connect to: {addr}");
Ok(stream)
},
Err(e) => {
debug!("Direct dial failed for {:?}: {}, using proxy", addr, e);
let proxy_res = proxy_fut.await;
self.handle_proxy_result(proxy_res, &addr, start.elapsed()).await
}
}
}
proxy_res = &mut proxy_fut => {
let direct_res = direct_fut.await;
match direct_res {
Ok(stream) => {
self.direct_list.add_address(&addr);
info!("Direct Connect to: {addr}");
Ok(stream)
}
Err(e) => {
debug!("Direct dial failed for {:?}: {}, using proxy", addr, e);
self.handle_proxy_result(proxy_res, &addr, start.elapsed()).await
}
}
}
}
}
}
pub async fn is_stream_closed<R: AsyncRead + Unpin>(
stream: &mut R,
timeout_duration: Duration,
) -> bool {
match timeout(timeout_duration, stream.read(&mut [])).await {
Ok(Ok(_)) => false,
Ok(Err(e)) => !matches!(e.kind(), std::io::ErrorKind::WouldBlock),
Err(_) => false,
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use async_trait::async_trait;
use socks5_proto::Address;
use tempfile::NamedTempFile;
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpListener,
time::Duration,
};
use super::{AsyncStream, Dial, DirectDial, SmartDial};
use crate::address_list::{DirectList, ProxyList};
struct MockDial {
succeed: bool,
delay: Duration,
}
impl MockDial {
fn succeed() -> Self {
Self {
succeed: true,
delay: Duration::ZERO,
}
}
fn fail() -> Self {
Self {
succeed: false,
delay: Duration::ZERO,
}
}
fn with_delay(delay: Duration, succeed: bool) -> Self {
Self { succeed, delay }
}
fn with_timeout_error(delay: Duration) -> Self {
Self {
succeed: false,
delay,
}
}
}
#[async_trait]
impl Dial for MockDial {
async fn dial(&self, _addr: Address) -> anyhow::Result<Box<dyn AsyncStream>> {
tokio::time::sleep(self.delay).await;
if self.succeed {
Ok(Box::new(tokio::io::duplex(64).0) as Box<dyn AsyncStream>)
} else {
Err(anyhow::anyhow!("mock dial failed"))
}
}
}
#[tokio::test]
async fn direct_dial_connects_to_local_listener() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let accept_task = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
stream.write_all(b"ok").await.unwrap();
});
let mut stream = DirectDial::new(Duration::from_secs(3))
.dial(Address::SocketAddress(addr))
.await
.unwrap();
let mut buf = [0u8; 2];
AsyncReadExt::read_exact(&mut stream, &mut buf)
.await
.unwrap();
accept_task.await.unwrap();
assert_eq!(&buf, b"ok");
}
#[tokio::test]
async fn direct_dial_returns_error_for_unreachable_port() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
drop(listener);
let result = DirectDial::new(Duration::from_secs(3))
.dial(Address::SocketAddress(addr))
.await;
assert!(result.is_err());
}
#[tokio::test]
async fn smart_dial_uses_proxy_when_domain_in_list() {
let proxy_temp_file = NamedTempFile::new().unwrap();
std::fs::write(proxy_temp_file.path(), "blocked.com\n").unwrap();
let direct_temp_file = NamedTempFile::new().unwrap();
let proxy_list = Arc::new(ProxyList::new(proxy_temp_file.path()));
let direct_list = Arc::new(DirectList::new(direct_temp_file.path()));
let direct = MockDial::succeed();
let proxy = MockDial::succeed();
let smart_dial = SmartDial::new(
Box::new(direct),
Box::new(proxy),
proxy_list,
direct_list,
Duration::from_secs(3),
);
let result = smart_dial
.dial(Address::DomainAddress(b"blocked.com".to_vec(), 443))
.await;
assert!(result.is_ok());
}
#[tokio::test]
async fn smart_dial_uses_direct_when_succeeds_first() {
let proxy_temp_file = NamedTempFile::new().unwrap();
let direct_temp_file = NamedTempFile::new().unwrap();
let proxy_list = Arc::new(ProxyList::new(proxy_temp_file.path()));
let direct_list = Arc::new(DirectList::new(direct_temp_file.path()));
let direct = MockDial::with_delay(Duration::from_millis(10), true);
let proxy = MockDial::with_delay(Duration::from_millis(100), true);
let smart_dial = SmartDial::new(
Box::new(direct),
Box::new(proxy),
proxy_list,
direct_list,
Duration::from_secs(3),
);
let result = smart_dial
.dial(Address::DomainAddress(b"fast-direct.com".to_vec(), 443))
.await;
assert!(result.is_ok());
}
#[tokio::test]
async fn smart_dial_adds_domain_when_direct_times_out_and_proxy_succeeds() {
let proxy_temp_file = NamedTempFile::new().unwrap();
let direct_temp_file = NamedTempFile::new().unwrap();
let proxy_list = Arc::new(ProxyList::new(proxy_temp_file.path()));
let direct_list = Arc::new(DirectList::new(direct_temp_file.path()));
let direct = MockDial::with_timeout_error(Duration::from_millis(100));
let proxy = MockDial::with_delay(Duration::from_millis(10), true);
let smart_dial = SmartDial::new(
Box::new(direct),
Box::new(proxy),
proxy_list.clone(),
direct_list,
Duration::from_secs(3),
);
let result = smart_dial
.dial(Address::DomainAddress(b"slow-direct.com".to_vec(), 443))
.await;
assert!(result.is_ok());
assert!(
proxy_list.contains_address(&Address::DomainAddress(b"slow-direct.com".to_vec(), 443))
);
}
#[tokio::test]
async fn smart_dial_returns_error_when_both_fail() {
let proxy_temp_file = NamedTempFile::new().unwrap();
let direct_temp_file = NamedTempFile::new().unwrap();
let proxy_list = Arc::new(ProxyList::new(proxy_temp_file.path()));
let direct_list = Arc::new(DirectList::new(direct_temp_file.path()));
let direct = MockDial::fail();
let proxy = MockDial::fail();
let smart_dial = SmartDial::new(
Box::new(direct),
Box::new(proxy),
proxy_list,
direct_list,
Duration::from_secs(3),
);
let result = smart_dial
.dial(Address::DomainAddress(b"both-fail.com".to_vec(), 443))
.await;
assert!(result.is_err());
}
#[tokio::test]
async fn smart_dial_uses_direct_when_domain_in_direct_list() {
let proxy_temp_file = NamedTempFile::new().unwrap();
let direct_temp_file = NamedTempFile::new().unwrap();
std::fs::write(direct_temp_file.path(), "direct.com\n").unwrap();
let proxy_list = Arc::new(ProxyList::new(proxy_temp_file.path()));
let direct_list = Arc::new(DirectList::new(direct_temp_file.path()));
let direct = MockDial::succeed();
let proxy = MockDial::fail();
let smart_dial = SmartDial::new(
Box::new(direct),
Box::new(proxy),
proxy_list,
direct_list,
Duration::from_secs(3),
);
let result = smart_dial
.dial(Address::DomainAddress(b"direct.com".to_vec(), 443))
.await;
assert!(result.is_ok());
}
#[tokio::test]
async fn smart_dial_removes_from_direct_list_on_failure() {
let proxy_temp_file = NamedTempFile::new().unwrap();
let direct_temp_file = NamedTempFile::new().unwrap();
std::fs::write(direct_temp_file.path(), "direct-fail.com\n").unwrap();
let proxy_list = Arc::new(ProxyList::new(proxy_temp_file.path()));
let direct_list = Arc::new(DirectList::new(direct_temp_file.path()));
let direct = MockDial::fail();
let proxy = MockDial::succeed();
let smart_dial = SmartDial::new(
Box::new(direct),
Box::new(proxy),
proxy_list,
direct_list.clone(),
Duration::from_secs(3),
);
let result = smart_dial
.dial(Address::DomainAddress(b"direct-fail.com".to_vec(), 443))
.await;
assert!(result.is_ok());
assert!(
!direct_list
.contains_address(&Address::DomainAddress(b"direct-fail.com".to_vec(), 443))
);
}
#[tokio::test]
async fn smart_dial_adds_to_direct_list_on_success() {
let proxy_temp_file = NamedTempFile::new().unwrap();
let direct_temp_file = NamedTempFile::new().unwrap();
let proxy_list = Arc::new(ProxyList::new(proxy_temp_file.path()));
let direct_list = Arc::new(DirectList::new(direct_temp_file.path()));
let direct = MockDial::with_delay(Duration::from_millis(10), true);
let proxy = MockDial::with_delay(Duration::from_millis(100), true);
let smart_dial = SmartDial::new(
Box::new(direct),
Box::new(proxy),
proxy_list,
direct_list.clone(),
Duration::from_secs(3),
);
let result = smart_dial
.dial(Address::DomainAddress(b"new-direct.com".to_vec(), 443))
.await;
assert!(result.is_ok());
assert!(
direct_list.contains_address(&Address::DomainAddress(b"new-direct.com".to_vec(), 443))
);
}
}