[#85] Supporting shadowsocks-manager reports

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)",

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 (

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::<ServerAddr>()
+                .expect("Expecting a valid ServerAddr for manager_address"),
+        );
+    }
+
     info!("ShadowSocks {}", shadowsocks::VERSION);
 
     debug!("Config: {:?}", config);

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 (

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);

src/config.rs

@@ -205,6 +205,8 @@ pub struct ServerConfig {
     plugin: Option<PluginConfig>,
     /// UDP timeout
     udp_timeout: Option<Duration>,
+    /// Plugin address
+    plugin_addr: Option<ServerAddr>,
 }
 
 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<ServerAddr> {
+        &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<SocketAddr>,
     pub mode: Mode,
     pub no_delay: bool,
-}
-
-impl Default for Config {
-    fn default() -> Config {
-        Config::new()
-    }
+    pub manager_address: Option<ServerAddr>,
+    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

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;

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<StreamCipherVariant>,
+}
+
+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),
     }
 }

src/plugin/mod.rs

@@ -95,7 +95,7 @@ pub fn launch_plugin(config: &mut Config, mode: PluginMode) -> io::Result<Vec<Pl
         }
 
         if let Some(svr_addr) = svr_addr_opt {
-            svr.set_addr(svr_addr);
+            svr.set_plugin_addr(svr_addr);
         }
     }
 

src/relay/local.rs

@@ -16,14 +16,14 @@ use relay::{boxed_future, tcprelay::local::run as run_tcp, udprelay::local::run
 /// extern crate shadowsocks;
 ///
 /// use shadowsocks::{
-///     config::{Config, ServerConfig},
+///     config::{Config, ConfigType, ServerConfig},
 ///     crypto::CipherType,
 ///     relay::local::run,
 /// };
 ///
 /// use tokio::prelude::*;
 ///
-/// let mut config = Config::new();
+/// let mut config = Config::new(ConfigType::Local);
 /// config.local = Some("127.0.0.1:1080".parse().unwrap());
 /// config.server = vec![ServerConfig::basic(
 ///     "127.0.0.1:8388".parse().unwrap(),

src/relay/server.rs

@@ -16,14 +16,14 @@ use relay::{boxed_future, tcprelay::server::run as run_tcp, udprelay::server::ru
 /// extern crate shadowsocks;
 ///
 /// use shadowsocks::{
-///     config::{Config, ServerConfig},
+///     config::{Config, ConfigType, ServerConfig},
 ///     crypto::CipherType,
 ///     relay::server::run,
 /// };
 ///
 /// use tokio::prelude::*;
 ///
-/// let mut config = Config::new();
+/// let mut config = Config::new(ConfigType::Server);
 /// config.server = vec![ServerConfig::basic(
 ///     "127.0.0.1:8388".parse().unwrap(),
 ///     "server-password".to_string(),

src/relay/tcprelay/context.rs

@@ -0,0 +1,113 @@
+//! TCP Relay Context
+
+use std::{
+    io,
+    net::SocketAddr,
+    sync::{
+        atomic::{AtomicBool, AtomicUsize, Ordering},
+        Arc,
+    },
+    time::Duration,
+};
+
+use config::{ServerAddr, ServerConfig};
+use context::SharedContext;
+use relay::{boxed_future, dns_resolver::resolve};
+
+use futures::{future, Future, Stream};
+use tokio::{self, net::UdpSocket, timer::Interval};
+
+/// TCP Relay Server Context
+pub struct TcpServerContext {
+    tx: AtomicUsize,
+    rx: AtomicUsize,
+    context: SharedContext,
+    stop_flag: Arc<AtomicBool>,
+    svr_cfg: Arc<ServerConfig>,
+}
+
+const UPDATE_INTERVAL: Duration = Duration::from_secs(5);
+
+pub type SharedTcpServerContext = Arc<TcpServerContext>;
+
+impl TcpServerContext {
+    /// Create a new server context
+    pub fn new(context: SharedContext, svr_cfg: Arc<ServerConfig>) -> 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<ServerConfig> {
+        &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<Item = (), Error = io::Error> + 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::<SocketAddr>().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)
+    }
+}

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<TcpStream>;
-type TcpWriteHalf = WriteHalf<TcpStream>;
+type TcpReadHalf<S /* : Read + Write + AsyncRead + AsyncWrite + Send */> = ReadHalf<S>;
+type TcpWriteHalf<S /* : Read + Write + AsyncRead + AsyncWrite + Send */> = WriteHalf<S>;
 
 /// `ReadHalf `of `TcpStream` with decryption
-pub enum DecryptedHalf {
-    Stream(StreamDecryptedReader<TcpReadHalf>),
-    Aead(AeadDecryptedReader<TcpReadHalf>),
+pub enum DecryptedHalf<S>
+where
+    S: Read + Write + AsyncRead + AsyncWrite + Send,
+{
+    Stream(StreamDecryptedReader<TcpReadHalf<S>>),
+    Aead(AeadDecryptedReader<TcpReadHalf<S>>),
 }
 
 macro_rules! ref_half_do {
@@ -83,21 +89,30 @@ macro_rules! mut_half_do {
     }
 }
 
-impl Read for DecryptedHalf {
+impl<S> Read for DecryptedHalf<S>
+where
+    S: Read + Write + AsyncRead + AsyncWrite + Send,
+{
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
         mut_half_do!(self, DecryptedHalf, read, buf)
     }
 }
 
-impl DecryptedRead for DecryptedHalf {
+impl<S> DecryptedRead for DecryptedHalf<S>
+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<S> AsyncRead for DecryptedHalf<S> where S: Read + Write + AsyncRead + AsyncWrite + Send {}
 
-impl BufRead for DecryptedHalf {
+impl<S> BufRead for DecryptedHalf<S>
+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<StreamDecryptedReader<TcpReadHalf>> for DecryptedHalf {
-    fn from(r: StreamDecryptedReader<TcpReadHalf>) -> DecryptedHalf {
+impl<S> From<StreamDecryptedReader<TcpReadHalf<S>>> for DecryptedHalf<S>
+where
+    S: Read + Write + AsyncRead + AsyncWrite + Send,
+{
+    fn from(r: StreamDecryptedReader<TcpReadHalf<S>>) -> DecryptedHalf<S> {
         DecryptedHalf::Stream(r)
     }
 }
 
-impl From<AeadDecryptedReader<TcpReadHalf>> for DecryptedHalf {
-    fn from(r: AeadDecryptedReader<TcpReadHalf>) -> DecryptedHalf {
+impl<S> From<AeadDecryptedReader<TcpReadHalf<S>>> for DecryptedHalf<S>
+where
+    S: Read + Write + AsyncRead + AsyncWrite + Send,
+{
+    fn from(r: AeadDecryptedReader<TcpReadHalf<S>>) -> DecryptedHalf<S> {
         DecryptedHalf::Aead(r)
     }
 }
 
 /// `WriteHalf` of `TcpStream` with encryption
-pub enum EncryptedHalf {
-    Stream(StreamEncryptedWriter<TcpWriteHalf>),
-    Aead(AeadEncryptedWriter<TcpWriteHalf>),
+pub enum EncryptedHalf<S>
+where
+    S: Read + Write + AsyncRead + AsyncWrite + Send,
+{
+    Stream(StreamEncryptedWriter<TcpWriteHalf<S>>),
+    Aead(AeadEncryptedWriter<TcpWriteHalf<S>>),
 }
 
-impl EncryptedWrite for EncryptedHalf {
+impl<S> EncryptedWrite for EncryptedHalf<S>
+where
+    S: Read + Write + AsyncRead + AsyncWrite + Send,
+{
     fn write_raw(&mut self, data: &[u8]) -> io::Result<usize> {
         mut_half_do!(self, EncryptedHalf, write_raw, data)
     }
@@ -143,14 +170,20 @@ impl EncryptedWrite for EncryptedHalf {
     }
 }
 
-impl From<StreamEncryptedWriter<TcpWriteHalf>> for EncryptedHalf {
-    fn from(d: StreamEncryptedWriter<TcpWriteHalf>) -> EncryptedHalf {
+impl<S> From<StreamEncryptedWriter<TcpWriteHalf<S>>> for EncryptedHalf<S>
+where
+    S: Read + Write + AsyncRead + AsyncWrite + Send,
+{
+    fn from(d: StreamEncryptedWriter<TcpWriteHalf<S>>) -> EncryptedHalf<S> {
         EncryptedHalf::Stream(d)
     }
 }
 
-impl From<AeadEncryptedWriter<TcpWriteHalf>> for EncryptedHalf {
-    fn from(d: AeadEncryptedWriter<TcpWriteHalf>) -> EncryptedHalf {
+impl<S> From<AeadEncryptedWriter<TcpWriteHalf<S>>> for EncryptedHalf<S>
+where
+    S: Read + Write + AsyncRead + AsyncWrite + Send,
+{
+    fn from(d: AeadEncryptedWriter<TcpWriteHalf<S>>) -> EncryptedHalf<S> {
         EncryptedHalf::Aead(d)
     }
 }
@@ -232,15 +265,21 @@ fn connect_proxy_server(
     svr_cfg: Arc<ServerConfig>,
 ) -> impl Future<Item = TcpStream, Error = io::Error> + 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<S>(
+    remote_stream: S,
     svr_cfg: Arc<ServerConfig>,
     relay_addr: Address,
 ) -> impl Future<
     Item = (
-        impl Future<Item = DecryptedHalf, Error = io::Error> + Send,
-        impl Future<Item = EncryptedHalf, Error = io::Error> + Send,
+        impl Future<Item = DecryptedHalf<S>, Error = io::Error> + Send,
+        impl Future<Item = EncryptedHalf<S>, 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<S>(
+    remote_stream: S,
     svr_cfg: Arc<ServerConfig>,
 ) -> impl Future<
     Item = (
-        impl Future<Item = DecryptedHalf, Error = io::Error> + Send,
-        impl Future<Item = EncryptedHalf, Error = io::Error> + Send,
+        impl Future<Item = DecryptedHalf<S>, Error = io::Error> + Send,
+        impl Future<Item = EncryptedHalf<S>, 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();
 

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<usize> {
+        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<usize> {
+        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<Async<()>, 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
+    }
+}

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<ServerConfig>,
-    context: SharedContext,
+    s: TcpMonStream,
+    svr_context: SharedTcpServerContext,
 }
 
 impl TcpRelayClientHandshake {
@@ -58,16 +64,18 @@ impl TcpRelayClientHandshake {
     pub fn handshake(
         self,
     ) -> impl Future<
-        Item = TcpRelayClientPending<impl Future<Item = EncryptedHalf, Error = io::Error> + Send + 'static>,
+        Item = TcpRelayClientPending<
+            impl Future<Item = EncryptedHalf<TcpMonStream>, 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<E>
 where
-    E: Future<Item = EncryptedHalf, Error = io::Error> + Send + 'static,
+    E: Future<Item = EncryptedHalf<TcpMonStream>, Error = io::Error> + Send + 'static,
 {
-    r: DecryptedHalf,
+    r: DecryptedHalf<TcpMonStream>,
     addr: Address,
     w: E,
     timeout: Option<Duration>,
-    context: SharedContext,
+    svr_context: SharedTcpServerContext,
 }
 
 /// Connect to the remote server
@@ -135,21 +143,25 @@ fn connect_remote(
 
 impl<E> TcpRelayClientPending<E>
 where
-    E: Future<Item = EncryptedHalf, Error = io::Error> + Send + 'static,
+    E: Future<Item = EncryptedHalf<TcpMonStream>, Error = io::Error> + Send + 'static,
 {
     /// Connect to the remote server
     pub fn connect(
         self,
     ) -> impl Future<
-        Item = TcpRelayClientConnected<impl Future<Item = EncryptedHalf, Error = io::Error> + Send + 'static>,
+        Item = TcpRelayClientConnected<
+            impl Future<Item = EncryptedHalf<TcpMonStream>, 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<E>
 where
-    E: Future<Item = EncryptedHalf, Error = io::Error> + Send + 'static,
+    E: Future<Item = EncryptedHalf<TcpMonStream>, Error = io::Error> + Send + 'static,
 {
     server: (ReadHalf<TcpStream>, WriteHalf<TcpStream>),
-    client: (DecryptedHalf, E),
+    client: (DecryptedHalf<TcpMonStream>, E),
     timeout: Option<Duration>,
 }
 
 impl<E> TcpRelayClientConnected<E>
 where
-    E: Future<Item = EncryptedHalf, Error = io::Error> + Send + 'static,
+    E: Future<Item = EncryptedHalf<TcpMonStream>, Error = io::Error> + Send + 'static,
 {
     /// Establish tunnel
     #[inline]
@@ -182,12 +194,10 @@ where
     }
 }
 
-fn handle_client(
-    server_cfg: Arc<ServerConfig>,
-    context: SharedContext,
-    socket: TcpStream,
-) -> impl Future<Item = (), Error = ()> + Send {
-    if let Err(err) = socket.set_keepalive(server_cfg.timeout()) {
+fn handle_client(svr_context: SharedTcpServerContext, socket: TcpStream) -> impl Future<Item = (), Error = ()> + 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<Item = (), Error = io::Error>
 
     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<Item = (), Error = io::Error>
 
         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));

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<W> EncryptedWriter<W>

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);

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<Barrier>) {
     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<Barrier>) {
     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");
     });