starred/shadowsocks-rust
1
0
Fork 0
mirror of https://github.com/shadowsocks/shadowsocks-rust.git synced 2026-11-07 00:13:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixed CryptoStream thread safe issue, optimized code structure, hide proxy connection functions

Browse Source
This commit is contained in:
zonyitoo
2020-02-24 01:48:27 +08:00
parent d75329e149
commit f557b29b31
8 changed files with 180 additions and 212 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/config.rs
+7 -2
View File
@@ -313,8 +313,13 @@ impl ServerConfig {
}
/// Get plugin address
pub fn plugin_addr(&self) -> &Option<ServerAddr> {
&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
src/relay/loadbalancing/server.rs
+2 -2
View File
@@ -500,7 +500,7 @@ impl<S: ServerData + 'static> PingBalancer<S> {
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<S: ServerData + 'static> PingBalancer<S> {
}
Err(..) => {
// Timeout
debug!(
trace!(
"checked remote {} server {} latency timeout, elapsed {} ms",
server_type,
stat.server_config().addr(),
src/relay/tcprelay/client.rs
+4 -6
View File
@@ -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<STcpStream>,
pub stream: ProxyStream,
}
impl ServerClient {
@@ -139,9 +139,7 @@ impl ServerClient {
addr: &Address,
svr_cfg: &ServerConfig,
) -> io::Result<ServerClient> {
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 })
}
}
src/relay/tcprelay/crypto_io.rs
+8 -40
View File
@@ -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<S>`
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<CryptoStream<S>>, WriteHalf<CryptoStream<S>>) {
use tokio::io::split;
split(self)
}
}
@@ -227,35 +227,3 @@ where
self.priv_poll_shutdown(ctx)
}
}
pub struct CryptoStreamReadHalf<'a, S: 'a>(*mut CryptoStream<S>, 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<io::Result<usize>> {
let stream = unsafe { &mut *self.0 };
Pin::new(stream).priv_poll_read(ctx, buf)
}
}
pub struct CryptoStreamWriteHalf<'a, S: 'a>(*mut CryptoStream<S>, 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<io::Result<usize>> {
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<io::Result<()>> {
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<io::Result<()>> {
let stream = unsafe { &mut *self.0 };
Pin::new(stream).priv_poll_shutdown(ctx)
}
}
src/relay/tcprelay/local.rs
+3 -147
View File
@@ -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<Duration>,
) -> io::Result<STcpStream> {
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<STcpStream> {
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<CryptoStream<STcpStream>> {
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)
}
src/relay/tcprelay/mod.rs
-2
View File
@@ -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
src/relay/tcprelay/proxy_stream.rs
+155 -12
View File
@@ -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<TcpStream>),
Proxied(CryptoStream<Connection<TcpStream>>),
Direct(STcpStream),
Proxied(CryptoStream<STcpStream>),
}
#[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<ProxyStream> {
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<ProxyStream> {
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<Duration>,
) -> io::Result<STcpStream> {
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<STcpStream> {
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<CryptoStream<STcpStream>> {
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)
}
src/relay/tcprelay/server.rs
+1 -1
View File
@@ -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)