diff --git a/src/config.rs b/src/config.rs index cb989b05..494d84cf 100644 --- a/src/config.rs +++ b/src/config.rs @@ -313,8 +313,13 @@ impl ServerConfig { } /// Get plugin address - pub fn plugin_addr(&self) -> &Option { - &self.plugin_addr + pub fn plugin_addr(&self) -> Option<&ServerAddr> { + self.plugin_addr.as_ref() + } + + /// Get server's external address + pub fn external_addr(&self) -> &ServerAddr { + self.plugin_addr.as_ref().unwrap_or(&self.addr) } /// Get URL for QRCode diff --git a/src/relay/loadbalancing/server.rs b/src/relay/loadbalancing/server.rs index 83369f7c..499ca17f 100644 --- a/src/relay/loadbalancing/server.rs +++ b/src/relay/loadbalancing/server.rs @@ -500,7 +500,7 @@ impl PingBalancer { match res { Ok(Ok(..)) => { // Got the result ... record its time - debug!( + trace!( "checked remote {} server {} latency with {} ms", server_type, stat.server_config().addr(), @@ -521,7 +521,7 @@ impl PingBalancer { } Err(..) => { // Timeout - debug!( + trace!( "checked remote {} server {} latency timeout, elapsed {} ms", server_type, stat.server_config().addr(), diff --git a/src/relay/tcprelay/client.rs b/src/relay/tcprelay/client.rs index 6773f0a0..10e8f7cc 100644 --- a/src/relay/tcprelay/client.rs +++ b/src/relay/tcprelay/client.rs @@ -21,7 +21,7 @@ use crate::relay::socks5::{ TcpResponseHeader, }; -use super::{CryptoStream, STcpStream}; +use super::ProxyStream; use crate::{config::ServerConfig, context::SharedContext}; /// Socks5 proxy client @@ -130,7 +130,7 @@ impl AsyncWrite for Socks5Client { } pub(crate) struct ServerClient { - pub stream: CryptoStream, + pub stream: ProxyStream, } impl ServerClient { @@ -139,9 +139,7 @@ impl ServerClient { addr: &Address, svr_cfg: &ServerConfig, ) -> io::Result { - let stream = super::connect_proxy_server(&*context, svr_cfg).await?; - Ok(ServerClient { - stream: super::proxy_server_handshake(context, stream, svr_cfg, addr).await?, - }) + let stream = ProxyStream::connect_proxied(context, svr_cfg, addr).await?; + Ok(ServerClient { stream }) } } diff --git a/src/relay/tcprelay/crypto_io.rs b/src/relay/tcprelay/crypto_io.rs index 673dff7b..e177a90f 100644 --- a/src/relay/tcprelay/crypto_io.rs +++ b/src/relay/tcprelay/crypto_io.rs @@ -2,7 +2,7 @@ use std::{ io, - marker::{PhantomData, Unpin}, + marker::Unpin, pin::Pin, task::{Context, Poll}, }; @@ -11,7 +11,10 @@ use byte_string::ByteStr; use bytes::Bytes; use futures::ready; use log::{debug, trace}; -use tokio::prelude::*; +use tokio::{ + io::{ReadHalf, WriteHalf}, + prelude::*, +}; use crate::{ config::ServerConfig, @@ -193,12 +196,9 @@ where /// Split connection into reader and writer /// /// The two halfs share the same `CryptoStream` - pub fn split(&mut self) -> (CryptoStreamReadHalf<'_, S>, CryptoStreamWriteHalf<'_, S>) { - let p = self as *mut _; - ( - CryptoStreamReadHalf(p, PhantomData), - CryptoStreamWriteHalf(p, PhantomData), - ) + pub fn split(self) -> (ReadHalf>, WriteHalf>) { + use tokio::io::split; + split(self) } } @@ -227,35 +227,3 @@ where self.priv_poll_shutdown(ctx) } } - -pub struct CryptoStreamReadHalf<'a, S: 'a>(*mut CryptoStream, PhantomData<&'a S>); - -unsafe impl<'a, S: Send + 'a> Send for CryptoStreamReadHalf<'a, S> {} - -impl<'a, S: AsyncRead + Unpin + 'a> AsyncRead for CryptoStreamReadHalf<'a, S> { - fn poll_read(mut self: Pin<&mut Self>, ctx: &mut Context<'_>, buf: &mut [u8]) -> Poll> { - let stream = unsafe { &mut *self.0 }; - Pin::new(stream).priv_poll_read(ctx, buf) - } -} - -pub struct CryptoStreamWriteHalf<'a, S: 'a>(*mut CryptoStream, PhantomData<&'a S>); - -unsafe impl<'a, S: Send + 'a> Send for CryptoStreamWriteHalf<'a, S> {} - -impl<'a, S: AsyncWrite + Unpin + 'a> AsyncWrite for CryptoStreamWriteHalf<'a, S> { - fn poll_write(mut self: Pin<&mut Self>, ctx: &mut Context<'_>, buf: &[u8]) -> Poll> { - let stream = unsafe { &mut *self.0 }; - Pin::new(stream).priv_poll_write(ctx, buf) - } - - fn poll_flush(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll> { - let stream = unsafe { &mut *self.0 }; - Pin::new(stream).priv_poll_flush(ctx) - } - - fn poll_shutdown(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll> { - let stream = unsafe { &mut *self.0 }; - Pin::new(stream).priv_poll_shutdown(ctx) - } -} diff --git a/src/relay/tcprelay/local.rs b/src/relay/tcprelay/local.rs index ad3f852e..b4af48b1 100644 --- a/src/relay/tcprelay/local.rs +++ b/src/relay/tcprelay/local.rs @@ -1,18 +1,10 @@ //! Relay for TCP server that running on local environment -use std::{io, time::Duration}; +use std::io; -use bytes::BytesMut; -use log::{debug, error, trace}; -use tokio::{io::AsyncWriteExt, net::TcpStream}; +use crate::{config::ConfigType, context::SharedContext}; -use crate::{ - config::{ConfigType, ServerAddr, ServerConfig}, - context::{Context, SharedContext}, - relay::{socks5::Address, utils::try_timeout}, -}; - -use super::{http_local, redir_local, socks5_local, tunnel_local, CryptoStream, STcpStream}; +use super::{http_local, redir_local, socks5_local, tunnel_local}; /// Starts a TCP local server pub async fn run(context: SharedContext) -> io::Result<()> { @@ -25,139 +17,3 @@ pub async fn run(context: SharedContext) -> io::Result<()> { ConfigType::Manager => unreachable!(), } } - -async fn connect_proxy_server_internal( - context: &Context, - orig_svr_addr: &ServerAddr, - svr_addr: &ServerAddr, - timeout: Option, -) -> io::Result { - match svr_addr { - ServerAddr::SocketAddr(ref addr) => { - let stream = try_timeout(TcpStream::connect(addr), timeout).await?; - debug!("connected proxy {} ({})", orig_svr_addr, addr); - Ok(STcpStream::new(stream, timeout)) - } - ServerAddr::DomainName(ref domain, port) => { - let result = lookup_then!(context, domain.as_str(), *port, |addr| { - match try_timeout(TcpStream::connect(addr), timeout).await { - Ok(s) => Ok(STcpStream::new(s, timeout)), - Err(e) => { - debug!( - "failed to connect proxy {} ({}:{} ({})) try another (err: {})", - orig_svr_addr, domain, port, addr, e - ); - Err(e) - } - } - }); - - match result { - Ok((addr, s)) => { - debug!("connected proxy {} ({}:{} ({}))", orig_svr_addr, domain, port, addr); - Ok(s) - } - Err(err) => { - error!( - "failed to connect proxy {} ({}:{}), {}", - orig_svr_addr, domain, port, err - ); - Err(err) - } - } - } - } -} - -/// Connect to proxy server with `ServerConfig` -pub(crate) async fn connect_proxy_server(context: &Context, svr_cfg: &ServerConfig) -> io::Result { - let timeout = svr_cfg.timeout().or(context.config().timeout); - - let svr_addr = match context.config().config_type { - ConfigType::Server => svr_cfg.addr(), - ConfigType::Socks5Local | ConfigType::TunnelLocal | ConfigType::HttpLocal | ConfigType::RedirLocal => { - svr_cfg.plugin_addr().as_ref().unwrap_or_else(|| svr_cfg.addr()) - } - ConfigType::Manager => unreachable!("ConfigType::Manager shouldn't need to connect to proxy server"), - }; - - // Retry if connect failed - // - // FIXME: This won't work if server is actually down. - // Probably we should retry with another server. - // - // Also works if plugin is starting - const RETRY_TIMES: i32 = 3; - - let orig_svr_addr = svr_cfg.addr(); - trace!( - "connecting to proxy {} ({}), timeout: {:?}", - orig_svr_addr, - svr_addr, - timeout - ); - - let mut last_err = None; - for retry_time in 0..RETRY_TIMES { - match connect_proxy_server_internal(context, orig_svr_addr, svr_addr, timeout).await { - Ok(mut s) => { - // IMPOSSIBLE, won't fail, but just a guard - if let Err(err) = s.set_nodelay(context.config().no_delay) { - error!("failed to set TCP_NODELAY on remote socket, error: {:?}", err); - } - - return Ok(s); - } - Err(err) => { - // Connection failure, retry - debug!( - "failed to connect {}, retried {} times (last err: {})", - svr_addr, retry_time, err - ); - last_err = Some(err); - - // Yield and let the others' run - // - // It may take some time for scheduler to resume this coroutine. - tokio::task::yield_now().await; - } - } - } - - let last_err = last_err.unwrap(); - error!( - "failed to connect {}, retried {} times, last_err: {}", - svr_addr, RETRY_TIMES, last_err - ); - Err(last_err) -} - -/// Handshake logic for ShadowSocks Client -pub(crate) async fn proxy_server_handshake( - context: SharedContext, - remote_stream: STcpStream, - svr_cfg: &ServerConfig, - relay_addr: &Address, -) -> io::Result> { - let mut stream = CryptoStream::new(context, remote_stream, svr_cfg); - - trace!("got encrypt stream and going to send addr: {:?}", relay_addr); - - // Send relay address to remote - // - // NOTE: `Address` handshake packets are very small in most cases, - // so it will be sent with the IV/Nonce data (implemented inside `CryptoStream`). - // - // For lower latency, first packet should be sent back quickly, - // so TCP_NODELAY should be kept enabled until the first data packet is received. - let mut addr_buf = BytesMut::with_capacity(relay_addr.serialized_len()); - relay_addr.write_to_buf(&mut addr_buf); - stream.write_all(&addr_buf).await?; - - // Here we should keep the TCP_NODELAY set until the first packet is received. - // https://github.com/shadowsocks/shadowsocks-libev/pull/746 - // - // Reset TCP_NODELAY after the first packet is received and sent back. - - Ok(stream) -} diff --git a/src/relay/tcprelay/mod.rs b/src/relay/tcprelay/mod.rs index 58a35ee5..8128c594 100644 --- a/src/relay/tcprelay/mod.rs +++ b/src/relay/tcprelay/mod.rs @@ -35,8 +35,6 @@ pub use self::{ proxy_stream::ProxyStream, }; -pub(crate) use self::local::{connect_proxy_server, proxy_server_handshake}; - const BUFFER_SIZE: usize = 8 * 1024; // 8K buffer /// Secured TcpStream diff --git a/src/relay/tcprelay/proxy_stream.rs b/src/relay/tcprelay/proxy_stream.rs index 3d3057e3..04e7eefe 100644 --- a/src/relay/tcprelay/proxy_stream.rs +++ b/src/relay/tcprelay/proxy_stream.rs @@ -7,21 +7,23 @@ use std::{ net::SocketAddr, pin::Pin, task::{Context as TaskContext, Poll}, + time::Duration, }; -use log::trace; +use bytes::BytesMut; +use log::{debug, error, trace}; use tokio::{ - io::{AsyncRead, AsyncWrite, ReadHalf, WriteHalf}, + io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadHalf, WriteHalf}, net::TcpStream, }; use crate::{ - config::ServerConfig, + config::{ConfigType, ServerAddr, ServerConfig}, context::{Context, SharedContext}, - relay::socks5::Address, + relay::{socks5::Address, utils::try_timeout}, }; -use super::{connection::Connection, CryptoStream}; +use super::{connection::Connection, CryptoStream, STcpStream}; macro_rules! forward_call { ($self:expr, $method:ident $(, $param:expr)*) => { @@ -35,8 +37,8 @@ macro_rules! forward_call { /// Stream wrapper for both direct connections and proxied connections #[allow(clippy::large_enum_variant)] pub enum ProxyStream { - Direct(Connection), - Proxied(CryptoStream>), + Direct(STcpStream), + Proxied(CryptoStream), } #[derive(Debug)] @@ -91,7 +93,7 @@ impl ProxyStream { /// /// This is used for hosts that matches ACL bypassed rules pub async fn connect_direct(context: &Context, addr: &Address) -> io::Result { - trace!("connect to {} directly (bypassed)", addr); + debug!("connect to {} directly (bypassed)", addr); // NOTE: Direct connection's timeout is controlled by the global key let timeout = context.config().timeout; @@ -113,7 +115,7 @@ impl ProxyStream { } } - /// Connect to remote through proxy server + /// Connect to remote via proxy server /// /// This is used for hosts that matches ACL proxied rules pub async fn connect_proxied( @@ -121,10 +123,15 @@ impl ProxyStream { svr_cfg: &ServerConfig, addr: &Address, ) -> io::Result { - trace!("connect to {} through {} (proxied)", addr, svr_cfg.addr()); + debug!( + "connect to {} via {} ({}) (proxied)", + addr, + svr_cfg.addr(), + svr_cfg.external_addr() + ); - let server_stream = super::local::connect_proxy_server(&*context, svr_cfg).await?; - let proxy_stream = super::local::proxy_server_handshake(context, server_stream, svr_cfg, addr).await?; + let server_stream = connect_proxy_server(&*context, svr_cfg).await?; + let proxy_stream = proxy_server_handshake(context, server_stream, svr_cfg, addr).await?; Ok(ProxyStream::Proxied(proxy_stream)) } @@ -184,3 +191,139 @@ impl AsyncWrite for ProxyStream { forward_call!(self, poll_shutdown, cx) } } + +async fn connect_proxy_server_internal( + context: &Context, + orig_svr_addr: &ServerAddr, + svr_addr: &ServerAddr, + timeout: Option, +) -> io::Result { + match svr_addr { + ServerAddr::SocketAddr(ref addr) => { + let stream = try_timeout(TcpStream::connect(addr), timeout).await?; + trace!("connected proxy {} ({})", orig_svr_addr, addr); + Ok(STcpStream::new(stream, timeout)) + } + ServerAddr::DomainName(ref domain, port) => { + let result = lookup_then!(context, domain.as_str(), *port, |addr| { + match try_timeout(TcpStream::connect(addr), timeout).await { + Ok(s) => Ok(STcpStream::new(s, timeout)), + Err(e) => { + debug!( + "failed to connect proxy {} ({}:{} ({})) try another (err: {})", + orig_svr_addr, domain, port, addr, e + ); + Err(e) + } + } + }); + + match result { + Ok((addr, s)) => { + trace!("connected proxy {} ({}:{} ({}))", orig_svr_addr, domain, port, addr); + Ok(s) + } + Err(err) => { + error!( + "failed to connect proxy {} ({}:{}), {}", + orig_svr_addr, domain, port, err + ); + Err(err) + } + } + } + } +} + +/// Connect to proxy server with `ServerConfig` +async fn connect_proxy_server(context: &Context, svr_cfg: &ServerConfig) -> io::Result { + let timeout = svr_cfg.timeout().or(context.config().timeout); + + let svr_addr = match context.config().config_type { + ConfigType::Server => svr_cfg.addr(), + ConfigType::Socks5Local | ConfigType::TunnelLocal | ConfigType::HttpLocal | ConfigType::RedirLocal => { + svr_cfg.external_addr() + } + ConfigType::Manager => unreachable!("ConfigType::Manager shouldn't need to connect to proxy server"), + }; + + // Retry if connect failed + // + // FIXME: This won't work if server is actually down. + // Probably we should retry with another server. + // + // Also works if plugin is starting + const RETRY_TIMES: i32 = 3; + + let orig_svr_addr = svr_cfg.addr(); + trace!( + "connecting to proxy {} ({}), timeout: {:?}", + orig_svr_addr, + svr_addr, + timeout + ); + + let mut last_err = None; + for retry_time in 0..RETRY_TIMES { + match connect_proxy_server_internal(context, orig_svr_addr, svr_addr, timeout).await { + Ok(mut s) => { + // IMPOSSIBLE, won't fail, but just a guard + if let Err(err) = s.set_nodelay(context.config().no_delay) { + error!("failed to set TCP_NODELAY on remote socket, error: {:?}", err); + } + + return Ok(s); + } + Err(err) => { + // Connection failure, retry + debug!( + "failed to connect {}, retried {} times (last err: {})", + svr_addr, retry_time, err + ); + last_err = Some(err); + + // Yield and let the others' run + // + // It may take some time for scheduler to resume this coroutine. + tokio::task::yield_now().await; + } + } + } + + let last_err = last_err.unwrap(); + error!( + "failed to connect {}, retried {} times, last_err: {}", + svr_addr, RETRY_TIMES, last_err + ); + Err(last_err) +} + +/// Handshake logic for ShadowSocks Client +async fn proxy_server_handshake( + context: SharedContext, + remote_stream: STcpStream, + svr_cfg: &ServerConfig, + relay_addr: &Address, +) -> io::Result> { + let mut stream = CryptoStream::new(context, remote_stream, svr_cfg); + + trace!("got encrypt stream and going to send addr: {:?}", relay_addr); + + // Send relay address to remote + // + // NOTE: `Address` handshake packets are very small in most cases, + // so it will be sent with the IV/Nonce data (implemented inside `CryptoStream`). + // + // For lower latency, first packet should be sent back quickly, + // so TCP_NODELAY should be kept enabled until the first data packet is received. + let mut addr_buf = BytesMut::with_capacity(relay_addr.serialized_len()); + relay_addr.write_to_buf(&mut addr_buf); + stream.write_all(&addr_buf).await?; + + // Here we should keep the TCP_NODELAY set until the first packet is received. + // https://github.com/shadowsocks/shadowsocks-libev/pull/746 + // + // Reset TCP_NODELAY after the first packet is received and sent back. + + Ok(stream) +} diff --git a/src/relay/tcprelay/server.rs b/src/relay/tcprelay/server.rs index 384a08b2..5d7ec217 100644 --- a/src/relay/tcprelay/server.rs +++ b/src/relay/tcprelay/server.rs @@ -162,7 +162,7 @@ pub async fn run(context: SharedContext, flow_stat: SharedMultiServerFlowStatist for (idx, svr_cfg) in context.config().server.iter().enumerate() { let mut listener = { - let addr = svr_cfg.plugin_addr().as_ref().unwrap_or_else(|| svr_cfg.addr()); + let addr = svr_cfg.external_addr(); let addr = addr.bind_addr(&*context).await?; let listener = TcpListener::bind(&addr)