Refactored and ready for adding ACL rules

BREAKING: Removed forbidden_ip in configuration file
This commit is contained in:
zonyitoo
2020-02-23 01:43:48 +08:00
parent dfc0be7e34
commit 76bde20fe6
18 changed files with 427 additions and 473 deletions
+1 -3
View File
@@ -1,6 +1,6 @@
[package]
name = "shadowsocks-rust"
version = "1.8.10"
version = "1.8.11"
authors = ["Y. T. CHUNG <zonyitoo@gmail.com>"]
description = "shadowsocks is a fast tunnel proxy that helps you bypass firewalls."
repository = "https://github.com/zonyitoo/shadowsocks-rust"
@@ -93,8 +93,6 @@ bloomfilter = "^1.0.2"
spin = "0.5"
mio = "0.6"
serde_json = "1.0"
ipnetwork = "0.16"
regex = "1"
[target.'cfg(windows)'.dependencies]
winapi = { version = "0.3", features = ["mswsock", "winsock2"] }
+21 -1
View File
@@ -12,7 +12,16 @@ use futures::{
use log::{error, info};
use tokio::runtime::Builder;
use shadowsocks::{plugin::PluginConfig, run_local, Config, ConfigType, Mode, ServerAddr, ServerConfig};
use shadowsocks::{
acl::AccessControl,
plugin::PluginConfig,
run_local,
Config,
ConfigType,
Mode,
ServerAddr,
ServerConfig,
};
mod logging;
mod monitor;
@@ -101,6 +110,12 @@ fn main() {
.takes_value(true)
.help("Set RLIMIT_NOFILE with both soft and hard limit (only for *nix systems)"),
)
.arg(
Arg::with_name("ACL")
.long("acl")
.takes_value(true)
.help("Path to ACL (Access Control List)"),
)
.get_matches();
let debug_level = matches.occurrences_of("VERBOSE");
@@ -206,6 +221,11 @@ fn main() {
);
}
if let Some(acl_file) = matches.value_of("ACL") {
let acl = AccessControl::load_from_file(acl_file).expect("Load ACL file");
config.acl = Some(acl);
}
if config.local.is_none() {
eprintln!(
"Missing `local_address`, consider specifying it by --local-addr command line option, \
-35
View File
@@ -42,7 +42,6 @@
//! These defined server will be used with a load balancing algorithm.
use std::{
collections::HashSet,
convert::From,
default::Default,
error,
@@ -104,8 +103,6 @@ struct SSConfig {
#[serde(skip_serializing_if = "Option::is_none")]
servers: Option<Vec<SSServerExtConfig>>,
#[serde(skip_serializing_if = "Option::is_none")]
forbidden_ip: Option<Vec<String>>,
#[serde(skip_serializing_if = "Option::is_none")]
dns: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
mode: Option<String>,
@@ -681,10 +678,6 @@ pub struct Config {
pub local: Option<ClientConfig>,
/// Destination address for tunnel
pub forward: Option<Address>,
/// Ignored IPs
///
/// Suggested list: `["127.0.0.1", "::1"]`
pub forbidden_ip: HashSet<IpAddr>,
/// DNS configuration, uses system-wide DNS configuration by default
///
/// Value could be a `IpAddr`, uses UDP DNS protocol with port `53`. For example: `8.8.8.8`
@@ -779,7 +772,6 @@ impl Config {
server: Vec::new(),
local: None,
forward: None,
forbidden_ip: HashSet::new(),
dns: None,
mode: Mode::TcpOnly,
no_delay: false,
@@ -935,20 +927,6 @@ impl Config {
nconfig.manager_address = Some(manager);
}
// Forbidden IPs
if let Some(forbidden_ip) = config.forbidden_ip {
for fi in forbidden_ip {
match fi.parse::<IpAddr>() {
Ok(i) => {
nconfig.forbidden_ip.insert(i);
}
Err(err) => {
error!("Invalid forbidden_ip \"{}\", err: {}", fi, err);
}
}
}
}
// DNS
nconfig.dns = config.dns;
@@ -1045,11 +1023,6 @@ impl Config {
false
}
/// Check if IP is forbidden
pub fn check_forbidden_ip(&self, ip: &IpAddr) -> bool {
self.forbidden_ip.contains(ip)
}
/// Check if all required fields are already set
pub fn check_integrity(&self) -> Result<(), Error> {
if self.config_type.is_local() {
@@ -1183,14 +1156,6 @@ impl fmt::Display for Config {
jconf.no_delay = Some(self.no_delay);
}
if !self.forbidden_ip.is_empty() {
let mut vfi = Vec::new();
for fi in &self.forbidden_ip {
vfi.push(fi.to_string());
}
jconf.forbidden_ip = Some(vfi);
}
if let Some(ref dns) = self.dns {
jconf.dns = Some(dns.to_string());
}
+11 -7
View File
@@ -1,7 +1,8 @@
//! Shadowsocks Server Context
use std::{
net::IpAddr,
io,
net::SocketAddr,
sync::{
atomic::{AtomicBool, Ordering},
Arc,
@@ -14,9 +15,12 @@ use tokio::runtime::Handle;
#[cfg(feature = "trust-dns")]
use trust_dns_resolver::TokioAsyncResolver;
use crate::config::{Config, ConfigType, ServerConfig};
#[cfg(feature = "trust-dns")]
use crate::relay::dns_resolver::create_resolver;
use crate::{
config::{Config, ConfigType, ServerConfig},
relay::dns_resolver::resolve,
};
// Entries for server's bloom filter
//
@@ -199,6 +203,11 @@ impl Context {
self.server_state.dns_resolver()
}
/// Perform a DNS resolution
pub async fn dns_resolve(&self, host: &str, port: u16) -> io::Result<Vec<SocketAddr>> {
resolve(self, host, port).await
}
/// Check if the server is still in running state
pub fn server_running(&self) -> bool {
self.server_running.load(Ordering::Acquire)
@@ -209,11 +218,6 @@ impl Context {
self.server_running.store(false, Ordering::Release)
}
/// Check if IP is in forbidden list
pub fn check_forbidden_ip(&self, ip: &IpAddr) -> bool {
self.config.check_forbidden_ip(ip)
}
/// Check if nonce exist or not
///
/// If not, set into the current bloom filter
+3 -5
View File
@@ -25,12 +25,10 @@ cfg_if! {
/// Helper macro for resolving host and then process each addresses
#[macro_export]
macro_rules! lookup_then {
($context:expr, $addr:expr, $port:expr, $check_forbidden:expr, |$resolved_addr:ident| $body:block) => {{
use crate::relay::dns_resolver::resolve;
($context:expr, $addr:expr, $port:expr, |$resolved_addr:ident| $body:block) => {{
let mut result = None;
for $resolved_addr in resolve($context, $addr, $port, $check_forbidden).await? {
for $resolved_addr in $context.dns_resolve($addr, $port).await? {
match $body {
Ok(r) => {
result = Some(Ok(($resolved_addr, r)));
@@ -51,7 +49,7 @@ pub async fn resolve_bind_addr(context: &Context, addr: &ServerAddr) -> io::Resu
match addr {
ServerAddr::SocketAddr(ref a) => Ok(*a),
ServerAddr::DomainName(ref dname, port) => {
let addrs = self::resolve(context, dname, *port, false).await?;
let addrs = self::resolve(context, dname, *port).await?;
Ok(addrs[0])
}
}
+6 -27
View File
@@ -3,41 +3,20 @@ use std::{
net::SocketAddr,
};
use log::debug;
use tokio::net::lookup_host;
use crate::context::Context;
/// Perform a DNS resolution
pub async fn resolve(context: &Context, addr: &str, port: u16, check_forbidden: bool) -> io::Result<Vec<SocketAddr>> {
pub async fn resolve(_: &Context, addr: &str, port: u16) -> io::Result<Vec<SocketAddr>> {
match lookup_host((addr, port)).await {
Ok(v) => Ok(v.collect()),
Err(err) => {
let err = Error::new(ErrorKind::Other, format!("dns resolve {}:{}, {}", addr, port, err));
let err = Error::new(
ErrorKind::Other,
format!("dns resolve \"{}:{}\" error: {}", addr, port, err),
);
Err(err)
}
Ok(addrs) => {
let mut vaddr = Vec::new();
for addr in addrs {
let ip = addr.ip();
if check_forbidden && context.check_forbidden_ip(&ip) {
debug!("Resolved {} => {}, which is skipped by forbidden_ip", addr, ip);
continue;
}
vaddr.push(addr);
}
if vaddr.is_empty() {
let err = Error::new(
ErrorKind::Other,
format!("resolved {}:{}, but all IPs are filtered", addr, port),
);
Err(err)
} else {
debug!("Resolved {}:{} => {:?}", addr, port, vaddr);
Ok(vaddr)
}
}
}
}
+8 -26
View File
@@ -5,7 +5,7 @@ use std::{
net::SocketAddr,
};
use log::{debug, error, trace};
use log::{error, trace};
use tokio::{self, runtime::Handle};
use trust_dns_resolver::{config::ResolverConfig, TokioAsyncResolver};
@@ -79,37 +79,19 @@ pub async fn create_resolver(dns: Option<ResolverConfig>, rt: Handle) -> io::Res
}
/// Perform a DNS resolution
pub async fn resolve(context: &Context, addr: &str, port: u16, check_forbidden: bool) -> io::Result<Vec<SocketAddr>> {
pub async fn resolve(context: &Context, addr: &str, port: u16) -> io::Result<Vec<SocketAddr>> {
match context.dns_resolver() {
Some(resolver) => match resolver.lookup_ip(addr).await {
Ok(lookup_result) => {
let mut vaddr = Vec::new();
for ip in lookup_result.iter() {
if check_forbidden && context.check_forbidden_ip(&ip) {
debug!("Resolved {} => {}, which is skipped by forbidden_ip", addr, ip);
continue;
}
vaddr.push(SocketAddr::new(ip, port));
}
if vaddr.is_empty() {
let err = Error::new(
ErrorKind::Other,
format!("resolved {}:{}, but all IPs are filtered", addr, port),
);
Err(err)
} else {
debug!("Resolved {}:{} => {:?}", addr, port, vaddr);
Ok(vaddr)
}
}
Ok(lookup_result) => Ok(lookup_result.iter().map(|ip| SocketAddr::new(ip, port)).collect()),
Err(err) => {
let err = Error::new(ErrorKind::Other, format!("dns resolve {}:{}, {}", addr, port, err));
let err = Error::new(
ErrorKind::Other,
format!("dns resolve \"{}:{}\" error: {}", addr, port, err),
);
Err(err)
}
},
// Fallback to tokio's DNS resolver
None => tokio_resolve(context, addr, port, check_forbidden).await,
None => tokio_resolve(context, addr, port).await,
}
}
+2 -3
View File
@@ -124,8 +124,7 @@ impl ManagerDatagram {
match *bind_addr {
ManagerAddr::SocketAddr(ref saddr) => Ok(ManagerDatagram::UdpDatagram(create_udp_socket(saddr).await?)),
ManagerAddr::DomainName(ref dname, port) => {
let (_, socket) =
lookup_then!(context, dname, port, false, |saddr| { create_udp_socket(&saddr).await })?;
let (_, socket) = lookup_then!(context, dname, port, |saddr| { create_udp_socket(&saddr).await })?;
Ok(ManagerDatagram::UdpDatagram(socket))
}
@@ -205,7 +204,7 @@ impl ManagerDatagram {
ManagerDatagram::UdpDatagram(ref mut udp) => match *target {
ManagerAddr::SocketAddr(ref saddr) => udp.send_to(buf, saddr).await,
ManagerAddr::DomainName(ref dname, port) => {
let (_, n) = lookup_then!(context, dname, port, false, |saddr| { udp.send_to(buf, saddr).await })?;
let (_, n) = lookup_then!(context, dname, port, |saddr| { udp.send_to(buf, saddr).await })?;
Ok(n)
}
#[cfg(unix)]
+210
View File
@@ -0,0 +1,210 @@
//! TCP Connection with timeout
use std::{
future::Future,
io::{self},
ops::{Deref, DerefMut},
pin::Pin,
task::{Context, Poll},
time::Duration,
};
use futures::ready;
use log::error;
use tokio::{
io::{AsyncRead, AsyncWrite, BufReader, ReadHalf, WriteHalf},
net::TcpStream,
time::{self, Delay},
};
/// Methods required for a TCP Connection
pub trait TcpConnection {
fn set_nodelay(&self, nodelay: bool) -> io::Result<()>;
}
impl TcpConnection for TcpStream {
fn set_nodelay(&self, nodelay: bool) -> io::Result<()> {
TcpStream::set_nodelay(self, nodelay)
}
}
/// Shadowsocks' Connection
///
/// The only feature: Supports timeout
pub struct Connection<S> {
// Actual connection socket
stream: BufReader<S>,
// Timer instance
// Read and Write operations shares the same timer
timer: Option<Delay>,
// User defined server timeout
timeout: Option<Duration>,
// TCP_NODELAY
nodelay: bool,
// Written the first packet flag
//
// Connection is usually wrapped inside a `CryptoStream`, which will send IV/Nonce within the first data packet.
// `TCP_NODELAY` is already set on the internal socket for lower handshake latency.
//
// After the first packet, if `nodelay` is `false`, `TCP_NODELAY` status should be reset.
written_first_packet: bool,
}
impl<S> Connection<S>
where
S: AsyncRead + TcpConnection,
{
/// Create a Connection with a stream S
///
/// If `timeout` is Some(..), it will set a timer for both read and write operation.
pub fn new(stream: S, timeout: Option<Duration>) -> Connection<S> {
// Set `TCP_NODELAY` for quick handshaking
if let Err(err) = stream.set_nodelay(true) {
error!("Failed to set TCP_NODELAY on socket, error: {:?}", err);
}
Connection {
stream: BufReader::new(stream),
timer: None,
timeout,
nodelay: false,
written_first_packet: false,
}
}
/// Set `TCP_NODELAY` on socket
pub fn set_nodelay(&mut self, nodelay: bool) -> io::Result<()> {
self.nodelay = nodelay;
// If first packet hasn't sent, resetting nodelay is delayed
if self.written_first_packet {
self.stream.get_ref().set_nodelay(nodelay)?;
}
Ok(())
}
}
impl<S> Connection<S> {
fn make_timeout_error() -> io::Error {
use std::io::ErrorKind;
ErrorKind::TimedOut.into()
}
fn poll_timeout(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
loop {
if let Some(ref mut timer) = self.timer {
ready!(Pin::new(timer).poll(cx));
// FIXME: Clear self.timer or not?
return Poll::Ready(Err(Connection::<S>::make_timeout_error()));
} else {
match self.timeout {
Some(timeout) => self.timer = Some(time::delay_for(timeout)),
None => break,
}
}
}
Poll::Ready(Ok(()))
}
fn cancel_timeout(&mut self) {
let _ = self.timer.take();
}
}
impl<S> Connection<S>
where
S: AsyncRead + AsyncWrite + TcpConnection + Unpin,
{
pub fn split(self) -> (ReadHalf<Connection<S>>, WriteHalf<Connection<S>>) {
use tokio::io::split;
split(self)
}
}
impl<S> Deref for Connection<S>
where
S: AsyncRead,
{
type Target = S;
fn deref(&self) -> &Self::Target {
self.stream.get_ref()
}
}
impl<S> DerefMut for Connection<S>
where
S: AsyncRead,
{
fn deref_mut(&mut self) -> &mut Self::Target {
self.stream.get_mut()
}
}
impl<S: Unpin> Unpin for Connection<S> {}
impl<S> AsyncRead for Connection<S>
where
S: AsyncRead + Unpin,
{
fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
match Pin::new(&mut self.stream).poll_read(cx, buf) {
Poll::Ready(r) => {
self.cancel_timeout();
Poll::Ready(r)
}
Poll::Pending => {
ready!(self.poll_timeout(cx))?;
Poll::Pending
}
}
}
}
impl<S> AsyncWrite for Connection<S>
where
S: AsyncRead + AsyncWrite + TcpConnection + Unpin,
{
fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
match Pin::new(&mut self.stream).poll_write(cx, buf) {
Poll::Ready(r) => {
self.cancel_timeout();
if !self.written_first_packet {
self.written_first_packet = true;
if !self.nodelay {
// Reset `TCP_NODELAY`
if let Err(err) = self.stream.get_ref().set_nodelay(false) {
error!("Failed to reset TCP_NODELAY on socket, error: {:?}", err);
}
}
}
Poll::Ready(r)
}
Poll::Pending => {
ready!(self.poll_timeout(cx))?;
Poll::Pending
}
}
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
match Pin::new(&mut self.stream).poll_flush(cx) {
Poll::Ready(r) => {
self.cancel_timeout();
Poll::Ready(r)
}
Poll::Pending => {
ready!(self.poll_timeout(cx))?;
Poll::Pending
}
}
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.stream).poll_shutdown(cx)
}
}
+148 -3
View File
@@ -1,9 +1,18 @@
//! Relay for TCP server that running on local environment
use std::io;
use std::{io, time::Duration};
use super::{http_local, redir_local, socks5_local, tunnel_local};
use crate::{config::ConfigType, context::SharedContext};
use bytes::BytesMut;
use log::{debug, error, trace};
use tokio::{io::AsyncWriteExt, net::TcpStream};
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};
/// Starts a TCP local server
pub async fn run(context: SharedContext) -> io::Result<()> {
@@ -16,3 +25,139 @@ 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();
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)
}
+10 -348
View File
@@ -4,33 +4,18 @@
// Maybe removed in the future
#![allow(clippy::unnecessary_mut_passed)]
use std::{
io,
marker::Unpin,
ops::{Deref, DerefMut},
pin::Pin,
task::{self, Poll},
time::Duration,
};
use std::{io, marker::Unpin};
use bytes::BytesMut;
use futures::{future::FusedFuture, ready, select, Future};
use log::{debug, error, trace};
use futures::{future::FusedFuture, select, Future};
use tokio::{
self,
io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader, ReadHalf, WriteHalf},
io::{AsyncRead, AsyncReadExt},
net::TcpStream,
time::{self, Delay},
};
use crate::{
config::{ConfigType, ServerAddr, ServerConfig},
context::{Context, SharedContext},
relay::{socks5::Address, utils::try_timeout},
};
mod aead;
pub mod client;
mod connection;
mod crypto_io;
mod http_local;
pub mod local;
@@ -43,341 +28,18 @@ mod sys;
mod tunnel_local;
mod utils;
pub use self::crypto_io::CryptoStream;
pub use self::{
connection::{Connection, TcpConnection},
crypto_io::CryptoStream,
};
pub(crate) use self::local::{connect_proxy_server, proxy_server_handshake};
const BUFFER_SIZE: usize = 8 * 1024; // 8K buffer
/// Methods required for a TCP Connection
pub trait TcpConnection {
fn set_nodelay(&self, nodelay: bool) -> io::Result<()>;
}
impl TcpConnection for TcpStream {
fn set_nodelay(&self, nodelay: bool) -> io::Result<()> {
TcpStream::set_nodelay(self, nodelay)
}
}
/// Shadowsocks' Connection
///
/// The only feature: Supports timeout
pub struct Connection<S> {
// Actual connection socket
stream: BufReader<S>,
// Timer instance
// Read and Write operations shares the same timer
timer: Option<Delay>,
// User defined server timeout
timeout: Option<Duration>,
// TCP_NODELAY
nodelay: bool,
// Written the first packet flag
//
// Connection is usually wrapped inside a `CryptoStream`, which will send IV/Nonce within the first data packet.
// `TCP_NODELAY` is already set on the internal socket for lower handshake latency.
//
// After the first packet, if `nodelay` is `false`, `TCP_NODELAY` status should be reset.
written_first_packet: bool,
}
impl<S> Connection<S>
where
S: AsyncRead + TcpConnection,
{
/// Create a Connection with a stream S
///
/// If `timeout` is Some(..), it will set a timer for both read and write operation.
pub fn new(stream: S, timeout: Option<Duration>) -> Connection<S> {
// Set `TCP_NODELAY` for quick handshaking
if let Err(err) = stream.set_nodelay(true) {
error!("Failed to set TCP_NODELAY on socket, error: {:?}", err);
}
Connection {
stream: BufReader::new(stream),
timer: None,
timeout,
nodelay: false,
written_first_packet: false,
}
}
/// Set `TCP_NODELAY` on socket
pub fn set_nodelay(&mut self, nodelay: bool) -> io::Result<()> {
self.nodelay = nodelay;
// If first packet hasn't sent, resetting nodelay is delayed
if self.written_first_packet {
self.stream.get_ref().set_nodelay(nodelay)?;
}
Ok(())
}
}
impl<S> Connection<S> {
fn make_timeout_error() -> io::Error {
use std::io::ErrorKind;
ErrorKind::TimedOut.into()
}
fn poll_timeout(&mut self, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
loop {
if let Some(ref mut timer) = self.timer {
ready!(Pin::new(timer).poll(cx));
// FIXME: Clear self.timer or not?
return Poll::Ready(Err(Connection::<S>::make_timeout_error()));
} else {
match self.timeout {
Some(timeout) => self.timer = Some(time::delay_for(timeout)),
None => break,
}
}
}
Poll::Ready(Ok(()))
}
fn cancel_timeout(&mut self) {
let _ = self.timer.take();
}
}
impl<S> Connection<S>
where
S: AsyncRead + AsyncWrite + TcpConnection + Unpin,
{
pub fn split(self) -> (ReadHalf<Connection<S>>, WriteHalf<Connection<S>>) {
use tokio::io::split;
split(self)
}
}
impl<S> Deref for Connection<S>
where
S: AsyncRead,
{
type Target = S;
fn deref(&self) -> &Self::Target {
self.stream.get_ref()
}
}
impl<S> DerefMut for Connection<S>
where
S: AsyncRead,
{
fn deref_mut(&mut self) -> &mut Self::Target {
self.stream.get_mut()
}
}
impl<S: Unpin> Unpin for Connection<S> {}
impl<S> AsyncRead for Connection<S>
where
S: AsyncRead + Unpin,
{
fn poll_read(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
match Pin::new(&mut self.stream).poll_read(cx, buf) {
Poll::Ready(r) => {
self.cancel_timeout();
Poll::Ready(r)
}
Poll::Pending => {
ready!(self.poll_timeout(cx))?;
Poll::Pending
}
}
}
}
impl<S> AsyncWrite for Connection<S>
where
S: AsyncRead + AsyncWrite + TcpConnection + Unpin,
{
fn poll_write(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
match Pin::new(&mut self.stream).poll_write(cx, buf) {
Poll::Ready(r) => {
self.cancel_timeout();
if !self.written_first_packet {
self.written_first_packet = true;
if !self.nodelay {
// Reset `TCP_NODELAY`
if let Err(err) = self.stream.get_ref().set_nodelay(false) {
error!("Failed to reset TCP_NODELAY on socket, error: {:?}", err);
}
}
}
Poll::Ready(r)
}
Poll::Pending => {
ready!(self.poll_timeout(cx))?;
Poll::Pending
}
}
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
match Pin::new(&mut self.stream).poll_flush(cx) {
Poll::Ready(r) => {
self.cancel_timeout();
Poll::Ready(r)
}
Poll::Pending => {
ready!(self.poll_timeout(cx))?;
Poll::Pending
}
}
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.stream).poll_shutdown(cx)
}
}
/// Secured TcpStream
pub type STcpStream = Connection<TcpStream>;
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, false, |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`
async fn connect_proxy_server(context: &Context, svr_cfg: &ServerConfig) -> io::Result<STcpStream> {
let timeout = svr_cfg.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 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)
}
/// Establish tunnel between server and client
// pub fn tunnel<CF, CFI, SF, SFI>(addr: Address, c2s: CF, s2c: SF) -> impl Future<Item = (), Error = io::Error> + Send
pub async fn tunnel<CF, CFI, SF, SFI>(mut c2s: CF, mut s2c: SF) -> io::Result<()>
+2 -9
View File
@@ -73,14 +73,7 @@ async fn handle_client(
let mut remote_stream = match remote_addr {
Address::SocketAddress(ref saddr) => {
if context.check_forbidden_ip(&saddr.ip()) {
error!("{} is forbidden, failed to connect {}", saddr.ip(), saddr);
let err = io::Error::new(
io::ErrorKind::Other,
format!("{} is forbidden, failed to connect {}", saddr.ip(), saddr),
);
return Err(err);
}
// TODO: ACL
match connect_tcp_stream(saddr, &bind_addr).await {
Ok(s) => {
@@ -94,7 +87,7 @@ async fn handle_client(
}
}
Address::DomainNameAddress(ref dname, port) => {
let result = lookup_then!(&*context, dname.as_str(), port, true, |addr| {
let result = lookup_then!(&*context, dname.as_str(), port, |addr| {
match connect_tcp_stream(&addr, &bind_addr).await {
Ok(s) => Ok(s),
Err(err) => {
+1 -1
View File
@@ -65,7 +65,7 @@ impl ServerClient {
ServerAddr::SocketAddr(ref remote_addr) => {
try_timeout(self.socket.send_to(&encrypt_buf[..], remote_addr), Some(timeout)).await?
}
ServerAddr::DomainName(ref dname, port) => lookup_then!(context, dname, port, false, |addr| {
ServerAddr::DomainName(ref dname, port) => lookup_then!(context, dname, port, |addr| {
try_timeout(self.socket.send_to(&encrypt_buf[..], addr), Some(timeout)).await
})
.map(|(_, l)| l)?,
+1 -1
View File
@@ -186,7 +186,7 @@ impl UdpAssociation {
ServerAddr::SocketAddr(ref remote_addr) => {
try_timeout(remote_udp.send_to(&encrypt_buf[..], remote_addr), Some(timeout)).await?
}
ServerAddr::DomainName(ref dname, port) => lookup_then!(context, dname, *port, false, |addr| {
ServerAddr::DomainName(ref dname, port) => lookup_then!(context, dname, *port, |addr| {
try_timeout(remote_udp.send_to(&encrypt_buf[..], &addr), Some(timeout)).await
})
.map(|(_, l)| l)?,
+1 -1
View File
@@ -184,7 +184,7 @@ impl UdpAssociation {
);
try_timeout(remote_udp.send_to(body, remote_addr), Some(timeout)).await?
}
Address::DomainNameAddress(ref dname, port) => lookup_then!(context, dname, port, true, |remote_addr| {
Address::DomainNameAddress(ref dname, port) => lookup_then!(context, dname, port, |remote_addr| {
match try_timeout(remote_udp.send_to(body, &remote_addr), Some(timeout)).await {
Ok(l) => {
debug!(
+1 -1
View File
@@ -182,7 +182,7 @@ impl UdpAssociation {
ServerAddr::SocketAddr(ref remote_addr) => {
try_timeout(remote_udp.send_to(&encrypt_buf[..], remote_addr), Some(timeout)).await?
}
ServerAddr::DomainName(ref dname, port) => lookup_then!(context, dname, *port, false, |addr| {
ServerAddr::DomainName(ref dname, port) => lookup_then!(context, dname, *port, |addr| {
try_timeout(remote_udp.send_to(&encrypt_buf[..], &addr), Some(timeout)).await
})
.map(|(_, l)| l)?,
+1 -1
View File
@@ -157,7 +157,7 @@ impl UdpAssociation {
ServerAddr::SocketAddr(ref remote_addr) => {
try_timeout(remote_udp.send_to(&encrypt_buf[..], remote_addr), Some(timeout)).await?
}
ServerAddr::DomainName(ref dname, port) => lookup_then!(context, dname, *port, false, |addr| {
ServerAddr::DomainName(ref dname, port) => lookup_then!(context, dname, *port, |addr| {
try_timeout(remote_udp.send_to(&encrypt_buf[..], &addr), Some(timeout)).await
})
.map(|(_, l)| l)?,
-1
View File
@@ -1,4 +1,3 @@
use env_logger;
use tokio::{
self,
net::{TcpStream, UdpSocket},