diff --git a/Cargo.lock b/Cargo.lock index 6068a06c..e6564e99 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1208,7 +1208,7 @@ dependencies = [ [[package]] name = "shadowsocks-rust" -version = "1.7.0-alpha.19" +version = "1.7.0-alpha.20" dependencies = [ "base64 0.9.3 (registry+https://github.com/rust-lang/crates.io-index)", "byte_string 1.0.0 (registry+https://github.com/rust-lang/crates.io-index)", diff --git a/src/bin/local.rs b/src/bin/local.rs index a6063655..ec3d7d53 100644 --- a/src/bin/local.rs +++ b/src/bin/local.rs @@ -196,7 +196,7 @@ fn main() { return; } }, - None => Config::new(), + None => Config::new(ConfigType::Local), }; let mut has_provided_server_config = match ( diff --git a/src/bin/server.rs b/src/bin/server.rs index e378a0b1..cff7bbf8 100644 --- a/src/bin/server.rs +++ b/src/bin/server.rs @@ -131,6 +131,12 @@ fn main() { .takes_value(false) .help("Set no-delay option for socket"), ) + .arg( + Arg::with_name("MANAGER_ADDRESS") + .long("manager-address") + .takes_value(true) + .help("ShadowSocks Manager (ssmgr) address"), + ) .get_matches(); let mut log_builder = Builder::new(); @@ -184,7 +190,7 @@ fn main() { return; } }, - None => Config::new(), + None => Config::new(ConfigType::Server), }; let has_provided_server_config = match ( @@ -250,6 +256,13 @@ fn main() { } }; + if let Some(m) = matches.value_of("MANAGER_ADDRESS") { + config.manager_address = Some( + m.parse::() + .expect("Expecting a valid ServerAddr for manager_address"), + ); + } + info!("ShadowSocks {}", shadowsocks::VERSION); debug!("Config: {:?}", config); diff --git a/src/bin/ssdns.rs b/src/bin/ssdns.rs index b750c096..a89471e7 100644 --- a/src/bin/ssdns.rs +++ b/src/bin/ssdns.rs @@ -174,7 +174,7 @@ fn main() { return; } }, - None => Config::new(), + None => Config::new(ConfigType::Local), }; let mut has_provided_server_config = match ( diff --git a/src/bin/ssurl.rs b/src/bin/ssurl.rs index 337e506f..34e15d18 100644 --- a/src/bin/ssurl.rs +++ b/src/bin/ssurl.rs @@ -63,7 +63,7 @@ fn encode(filename: &str, need_qrcode: bool) { fn decode(encoded: &str, need_qrcode: bool) { let svrconfig = ServerConfig::from_url(encoded).unwrap(); - let mut config = Config::new(); + let mut config = Config::new(ConfigType::Server); config.server.push(svrconfig); println!("{}", config); diff --git a/src/config.rs b/src/config.rs index 6174879e..c9f3f07c 100644 --- a/src/config.rs +++ b/src/config.rs @@ -205,6 +205,8 @@ pub struct ServerConfig { plugin: Option, /// UDP timeout udp_timeout: Option, + /// Plugin address + plugin_addr: Option, } impl ServerConfig { @@ -226,6 +228,7 @@ impl ServerConfig { enc_key: enc_key, plugin: plugin, udp_timeout: None, + plugin_addr: None, } } @@ -286,6 +289,16 @@ impl ServerConfig { &self.udp_timeout } + /// Set plugin address + pub fn set_plugin_addr(&mut self, a: ServerAddr) { + self.plugin_addr = Some(a); + } + + /// Get plugin address + pub fn plugin_addr(&self) -> &Option { + &self.plugin_addr + } + /// Get URL for QRCode /// ```plain /// ss:// + base64(method:password@host:port) @@ -464,7 +477,7 @@ impl error::Error for UrlParseError { pub type ClientConfig = SocketAddr; /// Server config type -#[derive(Clone, Copy)] +#[derive(Clone, Copy, Debug)] pub enum ConfigType { /// Config for local /// @@ -483,15 +496,15 @@ pub enum Mode { } impl Mode { - pub fn enable_udp(&self) -> bool { - match *self { + pub fn enable_udp(self) -> bool { + match self { Mode::UdpOnly | Mode::TcpAndUdp => true, _ => false, } } - pub fn enable_tcp(&self) -> bool { - match *self { + pub fn enable_tcp(self) -> bool { + match self { Mode::TcpOnly | Mode::TcpAndUdp => true, _ => false, } @@ -531,16 +544,12 @@ pub struct Config { pub remote_dns: Option, pub mode: Mode, pub no_delay: bool, -} - -impl Default for Config { - fn default() -> Config { - Config::new() - } + pub manager_address: Option, + pub config_type: ConfigType, } /// Configuration parsing error kind -#[derive(Copy, Clone)] +#[derive(Copy, Clone, Debug)] pub enum ErrorKind { MissingField, Malformed, @@ -590,7 +599,7 @@ impl Debug for Error { impl Config { /// Creates an empty configuration - pub fn new() -> Config { + pub fn new(config_type: ConfigType) -> Config { Config { server: Vec::new(), local: None, @@ -599,6 +608,8 @@ impl Config { remote_dns: None, mode: Mode::TcpOnly, no_delay: false, + manager_address: None, + config_type: config_type, } } @@ -617,7 +628,7 @@ impl Config { return Err(err); } - let mut nconfig = Config::new(); + let mut nconfig = Config::new(config_type); // Standard config // Client diff --git a/src/crypto/mod.rs b/src/crypto/mod.rs index e470c4a8..60b62504 100644 --- a/src/crypto/mod.rs +++ b/src/crypto/mod.rs @@ -7,7 +7,7 @@ use openssl::symm; pub use self::{ aead::{new_aead_decryptor, new_aead_encryptor, AeadDecryptor, AeadEncryptor, BoxAeadDecryptor, BoxAeadEncryptor}, cipher::{CipherCategory, CipherResult, CipherType}, - stream::{new_stream, StreamCipher, StreamCipherVariant}, + stream::{new_stream, BoxStreamCipher, StreamCipher}, }; pub mod aead; diff --git a/src/crypto/stream.rs b/src/crypto/stream.rs index f66f32be..8134d464 100644 --- a/src/crypto/stream.rs +++ b/src/crypto/stream.rs @@ -1,5 +1,7 @@ //! Stream ciphers +use std::ops::{Deref, DerefMut}; + #[cfg(feature = "rc4")] use crypto::rc4_md5; #[cfg(feature = "sodium")] @@ -83,6 +85,25 @@ macro_rules! define_stream_ciphers { } } +/// Variant cipher which contains all possible stream ciphers +pub struct BoxStreamCipher { + cipher: Box, +} + +impl Deref for BoxStreamCipher { + type Target = StreamCipherVariant; + + fn deref(&self) -> &StreamCipherVariant { + &*self.cipher + } +} + +impl DerefMut for BoxStreamCipher { + fn deref_mut(&mut self) -> &mut StreamCipherVariant { + &mut *self.cipher + } +} + define_stream_ciphers! { pub TableCipher => table::TableCipher, pub DummyCipher => dummy::DummyCipher, @@ -94,13 +115,13 @@ define_stream_ciphers! { } /// Generate a specific Cipher with key and initialize vector -pub fn new_stream(t: CipherType, key: &[u8], iv: &[u8], mode: CryptoMode) -> StreamCipherVariant { +pub fn new_stream(t: CipherType, key: &[u8], iv: &[u8], mode: CryptoMode) -> BoxStreamCipher { assert!( t.category() == CipherCategory::Stream, "only allow initializing with stream cipher" ); - match t { + let cipher = match t { CipherType::Table => StreamCipherVariant::new(table::TableCipher::new(key, mode)), CipherType::Plain => StreamCipherVariant::new(dummy::DummyCipher), @@ -113,5 +134,8 @@ pub fn new_stream(t: CipherType, key: &[u8], iv: &[u8], mode: CryptoMode) -> Str CipherType::Rc4Md5 => StreamCipherVariant::new(rc4_md5::Rc4Md5Cipher::new(key, iv, mode)), _ => StreamCipherVariant::new(openssl::OpenSSLCipher::new(t, key, iv, mode)), + }; + BoxStreamCipher { + cipher: Box::new(cipher), } } diff --git a/src/plugin/mod.rs b/src/plugin/mod.rs index 6156fdae..0c5a62ea 100644 --- a/src/plugin/mod.rs +++ b/src/plugin/mod.rs @@ -95,7 +95,7 @@ pub fn launch_plugin(config: &mut Config, mode: PluginMode) -> io::Result, + svr_cfg: Arc, +} + +const UPDATE_INTERVAL: Duration = Duration::from_secs(5); + +pub type SharedTcpServerContext = Arc; + +impl TcpServerContext { + /// Create a new server context + pub fn new(context: SharedContext, svr_cfg: Arc) -> SharedTcpServerContext { + let stop_flag = Arc::new(AtomicBool::new(false)); + + let ctx = TcpServerContext { + tx: AtomicUsize::new(0), + rx: AtomicUsize::new(0), + context: context, + stop_flag: stop_flag.clone(), + svr_cfg: svr_cfg, + }; + + let ctx = Arc::new(ctx); + + if ctx.context.config().manager_address.is_some() { + let ctx2 = ctx.clone(); + + let fut = Interval::new_interval(UPDATE_INTERVAL) + .map_err(|e| io::Error::new(io::ErrorKind::Other, format!("timer error: {}", e))) + .for_each(move |_| { + let ctx = ctx2.clone(); + if stop_flag.load(Ordering::Acquire) { + // Finished + Err(io::Error::new(io::ErrorKind::Other, "server terminated")) + } else { + tokio::spawn(ctx.stat_interval().map_err(|_| ())); + Ok(()) + } + }); + tokio::spawn(fut.map_err(|_| ())); + } + + ctx + } + + pub fn context(&self) -> &SharedContext { + &self.context + } + + pub fn svr_cfg(&self) -> &Arc { + &self.svr_cfg + } + + pub fn incr_tx(&self, x: usize) { + self.tx.fetch_add(x, Ordering::Release); + } + + pub fn incr_rx(&self, x: usize) { + self.rx.fetch_add(x, Ordering::Release); + } + + fn stat_interval(&self) -> impl Future + Send { + let addr_fut = match self.context.config().manager_address { + Some(ServerAddr::SocketAddr(ref addr)) => boxed_future(future::ok(*addr)), + Some(ServerAddr::DomainName(ref domain, ref port)) => { + let fut = resolve(self.context.clone(), &domain[..], *port, false).map(|vec_ipaddr| vec_ipaddr[0]); + boxed_future(fut) + } + None => unreachable!(), + }; + + let svr_cfg = self.svr_cfg.clone(); + let transmission = self.tx.load(Ordering::Acquire) + self.rx.load(Ordering::Acquire); + + addr_fut + .and_then(move |addr| { + let any_addr = "0.0.0.0:0".parse::().unwrap(); + UdpSocket::bind(&any_addr).map(move |socket| (socket, addr)) + }) + .and_then(move |(socket, addr)| { + let payload = format!("stat: {{\"{}\": {}}}", svr_cfg.addr().port(), transmission); + debug!("Sending Tcp Relay to {}, payload: {}", addr, payload); + socket.send_dgram(payload, &addr).map(|_| ()) + }) + } + + pub fn close(&self) { + self.stop_flag.store(true, Ordering::Release) + } +} diff --git a/src/relay/tcprelay/mod.rs b/src/relay/tcprelay/mod.rs index b3c588c3..bca17b26 100644 --- a/src/relay/tcprelay/mod.rs +++ b/src/relay/tcprelay/mod.rs @@ -1,7 +1,7 @@ //! Relay for TCP implementation use std::{ - io::{self, BufRead, Read}, + io::{self, BufRead, Read, Write}, iter::{IntoIterator, Iterator}, mem, net::SocketAddr, @@ -9,7 +9,7 @@ use std::{ time::Duration, }; -use config::{ServerAddr, ServerConfig}; +use config::{ConfigType, ServerAddr, ServerConfig}; use context::SharedContext; use crypto::CipherCategory; use relay::{boxed_future, dns_resolver::resolve, socks5::Address}; @@ -21,6 +21,7 @@ use tokio::{ use tokio_io::{ io::{read_exact, write_all, ReadHalf, WriteHalf}, AsyncRead, + AsyncWrite, }; use futures::{self, Async, Future, Poll}; @@ -38,8 +39,10 @@ use self::{ mod aead; pub mod client; +mod context; mod crypto_io; pub mod local; +mod monitor; pub mod server; mod socks5_local; mod stream; @@ -56,13 +59,16 @@ const BUFFER_SIZE: usize = 8 * 1024; // 8K buffer // Server2Client, // } -type TcpReadHalf = ReadHalf; -type TcpWriteHalf = WriteHalf; +type TcpReadHalf = ReadHalf; +type TcpWriteHalf = WriteHalf; /// `ReadHalf `of `TcpStream` with decryption -pub enum DecryptedHalf { - Stream(StreamDecryptedReader), - Aead(AeadDecryptedReader), +pub enum DecryptedHalf +where + S: Read + Write + AsyncRead + AsyncWrite + Send, +{ + Stream(StreamDecryptedReader>), + Aead(AeadDecryptedReader>), } macro_rules! ref_half_do { @@ -83,21 +89,30 @@ macro_rules! mut_half_do { } } -impl Read for DecryptedHalf { +impl Read for DecryptedHalf +where + S: Read + Write + AsyncRead + AsyncWrite + Send, +{ fn read(&mut self, buf: &mut [u8]) -> io::Result { mut_half_do!(self, DecryptedHalf, read, buf) } } -impl DecryptedRead for DecryptedHalf { +impl DecryptedRead for DecryptedHalf +where + S: Read + Write + AsyncRead + AsyncWrite + Send, +{ fn buffer_size(&self, data: &[u8]) -> usize { ref_half_do!(self, DecryptedHalf, buffer_size, data) } } -impl AsyncRead for DecryptedHalf {} +impl AsyncRead for DecryptedHalf where S: Read + Write + AsyncRead + AsyncWrite + Send {} -impl BufRead for DecryptedHalf { +impl BufRead for DecryptedHalf +where + S: Read + Write + AsyncRead + AsyncWrite + Send, +{ fn fill_buf(&mut self) -> io::Result<&[u8]> { mut_half_do!(self, DecryptedHalf, fill_buf) } @@ -107,25 +122,37 @@ impl BufRead for DecryptedHalf { } } -impl From> for DecryptedHalf { - fn from(r: StreamDecryptedReader) -> DecryptedHalf { +impl From>> for DecryptedHalf +where + S: Read + Write + AsyncRead + AsyncWrite + Send, +{ + fn from(r: StreamDecryptedReader>) -> DecryptedHalf { DecryptedHalf::Stream(r) } } -impl From> for DecryptedHalf { - fn from(r: AeadDecryptedReader) -> DecryptedHalf { +impl From>> for DecryptedHalf +where + S: Read + Write + AsyncRead + AsyncWrite + Send, +{ + fn from(r: AeadDecryptedReader>) -> DecryptedHalf { DecryptedHalf::Aead(r) } } /// `WriteHalf` of `TcpStream` with encryption -pub enum EncryptedHalf { - Stream(StreamEncryptedWriter), - Aead(AeadEncryptedWriter), +pub enum EncryptedHalf +where + S: Read + Write + AsyncRead + AsyncWrite + Send, +{ + Stream(StreamEncryptedWriter>), + Aead(AeadEncryptedWriter>), } -impl EncryptedWrite for EncryptedHalf { +impl EncryptedWrite for EncryptedHalf +where + S: Read + Write + AsyncRead + AsyncWrite + Send, +{ fn write_raw(&mut self, data: &[u8]) -> io::Result { mut_half_do!(self, EncryptedHalf, write_raw, data) } @@ -143,14 +170,20 @@ impl EncryptedWrite for EncryptedHalf { } } -impl From> for EncryptedHalf { - fn from(d: StreamEncryptedWriter) -> EncryptedHalf { +impl From>> for EncryptedHalf +where + S: Read + Write + AsyncRead + AsyncWrite + Send, +{ + fn from(d: StreamEncryptedWriter>) -> EncryptedHalf { EncryptedHalf::Stream(d) } } -impl From> for EncryptedHalf { - fn from(d: AeadEncryptedWriter) -> EncryptedHalf { +impl From>> for EncryptedHalf +where + S: Read + Write + AsyncRead + AsyncWrite + Send, +{ + fn from(d: AeadEncryptedWriter>) -> EncryptedHalf { EncryptedHalf::Aead(d) } } @@ -232,15 +265,21 @@ fn connect_proxy_server( svr_cfg: Arc, ) -> impl Future + Send { let timeout = svr_cfg.timeout(); - trace!("Connecting to proxy {:?}, timeout: {:?}", svr_cfg.addr(), timeout); - match *svr_cfg.addr() { + + let svr_addr = match context.config().config_type { + ConfigType::Server => svr_cfg.addr(), + ConfigType::Local => svr_cfg.plugin_addr().as_ref().unwrap_or_else(|| svr_cfg.addr()), + }; + + trace!("Connecting to proxy {:?}, timeout: {:?}", svr_addr, timeout); + match svr_addr { ServerAddr::SocketAddr(ref addr) => { let fut = try_timeout(TcpStream::connect(addr), timeout); boxed_future(fut) } ServerAddr::DomainName(ref domain, port) => { let fut = { - try_timeout(resolve(context, &domain[..], port, false), timeout).and_then(move |vec_ipaddr| { + try_timeout(resolve(context, &domain[..], *port, false), timeout).and_then(move |vec_ipaddr| { let fut = TcpStreamConnect::new(vec_ipaddr.into_iter()); try_timeout(fut, timeout) }) @@ -251,17 +290,20 @@ fn connect_proxy_server( } /// Handshake logic for ShadowSocks Client -pub fn proxy_server_handshake( - remote_stream: TcpStream, +pub fn proxy_server_handshake( + remote_stream: S, svr_cfg: Arc, relay_addr: Address, ) -> impl Future< Item = ( - impl Future + Send, - impl Future + Send, + impl Future, Error = io::Error> + Send, + impl Future, Error = io::Error> + Send, ), Error = io::Error, -> + Send { +> + Send +where + S: Read + Write + AsyncRead + AsyncWrite + Send + 'static, +{ let timeout = svr_cfg.timeout(); proxy_handshake(remote_stream, svr_cfg).and_then(move |(r_fut, w_fut)| { let w_fut = w_fut.and_then(move |enc_w| { @@ -284,16 +326,19 @@ pub fn proxy_server_handshake( /// ShadowSocks Client-Server handshake protocol /// Exchange cipher IV and creates stream wrapper -pub fn proxy_handshake( - remote_stream: TcpStream, +pub fn proxy_handshake( + remote_stream: S, svr_cfg: Arc, ) -> impl Future< Item = ( - impl Future + Send, - impl Future + Send, + impl Future, Error = io::Error> + Send, + impl Future, Error = io::Error> + Send, ), Error = io::Error, -> + Send { +> + Send +where + S: Read + Write + AsyncRead + AsyncWrite + Send + 'static, +{ futures::lazy(|| Ok(remote_stream.split())).and_then(move |(r, w)| { let timeout = svr_cfg.timeout(); diff --git a/src/relay/tcprelay/monitor.rs b/src/relay/tcprelay/monitor.rs new file mode 100644 index 00000000..73072fce --- /dev/null +++ b/src/relay/tcprelay/monitor.rs @@ -0,0 +1,67 @@ +//! Server traffic monitor + +use std::{ + io::{self, Read, Write}, + ops::{Deref, DerefMut}, +}; + +use tokio::{ + io::{AsyncRead, AsyncWrite}, + net::TcpStream, + prelude::Async, +}; + +use super::context::SharedTcpServerContext; + +pub struct TcpMonStream { + stream: TcpStream, + context: SharedTcpServerContext, +} + +impl TcpMonStream { + pub fn new(c: SharedTcpServerContext, s: TcpStream) -> TcpMonStream { + TcpMonStream { stream: s, context: c } + } +} + +impl Read for TcpMonStream { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + let n = self.stream.read(buf)?; + self.context.incr_rx(n); + Ok(n) + } +} + +impl Write for TcpMonStream { + fn write(&mut self, buf: &[u8]) -> io::Result { + let n = self.stream.write(buf)?; + self.context.incr_tx(n); + Ok(n) + } + + fn flush(&mut self) -> io::Result<()> { + self.stream.flush() + } +} + +impl AsyncRead for TcpMonStream {} + +impl AsyncWrite for TcpMonStream { + fn shutdown(&mut self) -> Result, io::Error> { + AsyncWrite::shutdown(&mut self.stream) + } +} + +impl Deref for TcpMonStream { + type Target = TcpStream; + + fn deref(&self) -> &TcpStream { + &self.stream + } +} + +impl DerefMut for TcpMonStream { + fn deref_mut(&mut self) -> &mut TcpStream { + &mut self.stream + } +} diff --git a/src/relay/tcprelay/server.rs b/src/relay/tcprelay/server.rs index 0cb34129..1863b045 100644 --- a/src/relay/tcprelay/server.rs +++ b/src/relay/tcprelay/server.rs @@ -7,8 +7,6 @@ use std::{ time::Duration, }; -use config::ServerConfig; - use relay::{ boxed_future, dns_resolver::resolve, @@ -33,13 +31,21 @@ use tokio_io::{ AsyncRead, }; -use super::{proxy_handshake, try_timeout, tunnel, DecryptedHalf, EncryptedHalf, TcpStreamConnect}; +use super::{ + context::{SharedTcpServerContext, TcpServerContext}, + monitor::TcpMonStream, + proxy_handshake, + try_timeout, + tunnel, + DecryptedHalf, + EncryptedHalf, + TcpStreamConnect, +}; /// Context for doing handshake with client pub struct TcpRelayClientHandshake { - s: TcpStream, - svr_cfg: Arc, - context: SharedContext, + s: TcpMonStream, + svr_context: SharedTcpServerContext, } impl TcpRelayClientHandshake { @@ -58,16 +64,18 @@ impl TcpRelayClientHandshake { pub fn handshake( self, ) -> impl Future< - Item = TcpRelayClientPending + Send + 'static>, + Item = TcpRelayClientPending< + impl Future, Error = io::Error> + Send + 'static, + >, Error = io::Error, > + Send { - let TcpRelayClientHandshake { s, svr_cfg, context } = self; + let TcpRelayClientHandshake { s, svr_context } = self; futures::lazy(move || s.peer_addr().map(|p| (s, p))).and_then(|(s, peer_addr)| { debug!("Handshaking with peer {}", peer_addr); - let timeout = svr_cfg.timeout(); - proxy_handshake(s, svr_cfg).and_then(move |(r_fut, w_fut)| { + let timeout = svr_context.svr_cfg().timeout(); + proxy_handshake(s, svr_context.svr_cfg().clone()).and_then(move |(r_fut, w_fut)| { r_fut .and_then(move |r| { let fut = @@ -79,7 +87,7 @@ impl TcpRelayClientHandshake { addr: addr, w: w_fut, timeout: timeout, - context: context, + svr_context: svr_context, }) }) }) @@ -89,13 +97,13 @@ impl TcpRelayClientHandshake { /// Context for connecting remote pub struct TcpRelayClientPending where - E: Future + Send + 'static, + E: Future, Error = io::Error> + Send + 'static, { - r: DecryptedHalf, + r: DecryptedHalf, addr: Address, w: E, timeout: Option, - context: SharedContext, + svr_context: SharedTcpServerContext, } /// Connect to the remote server @@ -135,21 +143,25 @@ fn connect_remote( impl TcpRelayClientPending where - E: Future + Send + 'static, + E: Future, Error = io::Error> + Send + 'static, { /// Connect to the remote server pub fn connect( self, ) -> impl Future< - Item = TcpRelayClientConnected + Send + 'static>, + Item = TcpRelayClientConnected< + impl Future, Error = io::Error> + Send + 'static, + >, Error = io::Error, > + Send { let client_pair = (self.r, self.w); let timeout = self.timeout; - connect_remote(self.context, self.addr, self.timeout).map(move |stream| TcpRelayClientConnected { - server: stream.split(), - client: client_pair, - timeout: timeout, + connect_remote(self.svr_context.context().clone(), self.addr, self.timeout).map(move |stream| { + TcpRelayClientConnected { + server: stream.split(), + client: client_pair, + timeout: timeout, + } }) } } @@ -157,16 +169,16 @@ where /// Context for extablishing tunnel pub struct TcpRelayClientConnected where - E: Future + Send + 'static, + E: Future, Error = io::Error> + Send + 'static, { server: (ReadHalf, WriteHalf), - client: (DecryptedHalf, E), + client: (DecryptedHalf, E), timeout: Option, } impl TcpRelayClientConnected where - E: Future + Send + 'static, + E: Future, Error = io::Error> + Send + 'static, { /// Establish tunnel #[inline] @@ -182,12 +194,10 @@ where } } -fn handle_client( - server_cfg: Arc, - context: SharedContext, - socket: TcpStream, -) -> impl Future + Send { - if let Err(err) = socket.set_keepalive(server_cfg.timeout()) { +fn handle_client(svr_context: SharedTcpServerContext, socket: TcpStream) -> impl Future + Send { + let socket = TcpMonStream::new(svr_context.clone(), socket); + + if let Err(err) = socket.set_keepalive(svr_context.svr_cfg().timeout()) { error!("Failed to set keep alive: {:?}", err); } @@ -204,12 +214,11 @@ fn handle_client( }) .and_then(move |(socket, addr)| { trace!("Got connection, addr: {}", addr); - trace!("Picked proxy server: {:?}", server_cfg); + trace!("Picked proxy server: {:?}", svr_context.svr_cfg()); let client = TcpRelayClientHandshake { s: socket, - svr_cfg: server_cfg, - context: context, + svr_context: svr_context, }; client @@ -228,7 +237,7 @@ pub fn run(context: SharedContext) -> impl Future for svr_cfg in &context.config().server { let listener = { - let addr = svr_cfg.addr(); + let addr = svr_cfg.plugin_addr().as_ref().unwrap_or_else(|| svr_cfg.addr()); let addr = addr.listen_addr(); let listener = TcpListener::bind(&addr).unwrap_or_else(|err| panic!("Failed to listen, {}", err)); @@ -239,17 +248,32 @@ pub fn run(context: SharedContext) -> impl Future let svr_cfg = Arc::new(svr_cfg.clone()); let context = context.clone(); + let svr_context = TcpServerContext::new(context.clone(), svr_cfg.clone()); + + struct CloseGuard(SharedTcpServerContext); + impl Drop for CloseGuard { + fn drop(&mut self) { + self.0.close(); + } + } + + let close_guard = CloseGuard(svr_context.clone()); + let listening = listener .incoming() .for_each(move |socket| { - let server_cfg = svr_cfg.clone(); - let context = context.clone(); - tokio::spawn(handle_client(server_cfg, context, socket)); + let svr_context = svr_context.clone(); + tokio::spawn(handle_client(svr_context, socket)); Ok(()) }) .map_err(|err| { error!("Server run failed: {}", err); err + }) + .then(move |r| { + // Close the context to ensure reporting Future is terminated + drop(close_guard); + r }); vec_fut.push(boxed_future(listening)); diff --git a/src/relay/tcprelay/stream.rs b/src/relay/tcprelay/stream.rs index 382b9cd3..76e97b5f 100644 --- a/src/relay/tcprelay/stream.rs +++ b/src/relay/tcprelay/stream.rs @@ -6,7 +6,7 @@ use std::{ }; use bytes::{BufMut, BytesMut}; -use crypto::{new_stream, CipherType, CryptoMode, StreamCipher, StreamCipherVariant}; +use crypto::{new_stream, BoxStreamCipher, CipherType, CryptoMode, StreamCipher}; use tokio_io::{AsyncRead, AsyncWrite}; use super::{DecryptedRead, EncryptedWrite, BUFFER_SIZE}; @@ -20,7 +20,7 @@ where { reader: R, buffer: BytesMut, - cipher: StreamCipherVariant, + cipher: BoxStreamCipher, pos: usize, sent_final: bool, } @@ -142,7 +142,7 @@ where W: AsyncWrite, { writer: W, - cipher: StreamCipherVariant, + cipher: BoxStreamCipher, } impl EncryptedWriter diff --git a/tests/socks5.rs b/tests/socks5.rs index addd1169..76c69a88 100644 --- a/tests/socks5.rs +++ b/tests/socks5.rs @@ -16,7 +16,7 @@ use tokio::runtime::current_thread::Runtime; use tokio_io::io::{flush, read_to_end, write_all}; use shadowsocks::{ - config::{Config, Mode, ServerConfig}, + config::{Config, ConfigType, Mode, ServerConfig}, crypto::CipherType, relay::{socks5::Address, tcprelay::client::Socks5Client}, run_local, @@ -25,7 +25,8 @@ use shadowsocks::{ pub struct Socks5TestServer { local_addr: SocketAddr, - config: Config, + svr_config: Config, + cli_config: Config, } impl Socks5TestServer { @@ -45,8 +46,14 @@ impl Socks5TestServer { Socks5TestServer { local_addr: local_addr, - config: { - let mut cfg = Config::new(); + svr_config: { + let mut cfg = Config::new(ConfigType::Server); + cfg.server = vec![ServerConfig::basic(svr_addr, pwd.to_owned(), method)]; + cfg.mode = if enable_udp { Mode::TcpAndUdp } else { Mode::TcpOnly }; + cfg + }, + cli_config: { + let mut cfg = Config::new(ConfigType::Local); cfg.local = Some(local_addr); cfg.server = vec![ServerConfig::basic(svr_addr, pwd.to_owned(), method)]; cfg.mode = if enable_udp { Mode::TcpAndUdp } else { Mode::TcpOnly }; @@ -60,14 +67,14 @@ impl Socks5TestServer { } pub fn run(&self) { - let svr_cfg = self.config.clone(); + let svr_cfg = self.svr_config.clone(); thread::spawn(move || { let mut runtime = Runtime::new().expect("Failed to create Runtime"); let fut = run_server(svr_cfg); runtime.block_on(fut).expect("Failed to run Server"); }); - let client_cfg = self.config.clone(); + let client_cfg = self.cli_config.clone(); thread::spawn(move || { let mut runtime = Runtime::new().expect("Failed to create Runtime"); let fut = run_local(client_cfg); diff --git a/tests/udp.rs b/tests/udp.rs index 432f6564..84f89c6c 100644 --- a/tests/udp.rs +++ b/tests/udp.rs @@ -21,7 +21,7 @@ use tokio::runtime::current_thread::Runtime; use tokio_io::io::read_to_end; use shadowsocks::{ - config::{Config, Mode, ServerConfig}, + config::{Config, ConfigType, Mode, ServerConfig}, crypto::CipherType, relay::{ socks5::{Address, UdpAssociateHeader}, @@ -40,8 +40,19 @@ const UDP_LOCAL_ADDR: &'static str = "127.0.0.1:9011"; const PASSWORD: &'static str = "test-password"; const METHOD: CipherType = CipherType::Aes128Cfb; -fn get_config() -> Config { - let mut cfg = Config::new(); +fn get_svr_config() -> Config { + let mut cfg = Config::new(ConfigType::Server); + cfg.server = vec![ServerConfig::basic( + SERVER_ADDR.parse().unwrap(), + PASSWORD.to_owned(), + METHOD, + )]; + cfg.mode = Mode::UdpOnly; + cfg +} + +fn get_cli_config() -> Config { + let mut cfg = Config::new(ConfigType::Local); cfg.local = Some(LOCAL_ADDR.parse().unwrap()); cfg.server = vec![ServerConfig::basic( SERVER_ADDR.parse().unwrap(), @@ -60,7 +71,7 @@ fn start_server(bar: Arc) { thread::spawn(move || { let mut runtime = Runtime::new().expect("Failed to create Runtime"); - let fut = run_server(get_config()); + let fut = run_server(get_svr_config()); bar.wait(); runtime.block_on(fut).expect("Failed to run Server"); }); @@ -70,7 +81,7 @@ fn start_local(bar: Arc) { thread::spawn(move || { let mut runtime = Runtime::new().expect("Failed to create Runtime"); - let fut = run_local(get_config()); + let fut = run_local(get_cli_config()); bar.wait(); runtime.block_on(fut).expect("Failed to run Local"); });