Migrate to std::future, async/await (#165)

* Socks5 library, main tcp relay implementation migrated with async/await

* UDP relay local migrated to async/await

* UDP relay server migrated to async/await

* Migrated UDP DNS to async/await, follows tokio-process

* Removed unused uses, add dyn keyword for TraitObjects

* [#151] Uses dl.google.com/generate_204 to check connection latency

* Make sslocal/ssserver TCP relay work successfully, disabled UDP and DNS relay temporary

* Handshake before performing actual read & write

* Wraps timeout logic into a Connection<S> struct

* Make UDP relay work again

* Local uses copy directly

* Completely removed unused ssdns

* Fixed test

* Uses tokio::main macro for main entrance

* Fixed a missing use in local.rs

* Updated travis openssl to 1.1.0L

* Windows signal

* Migrated tokio 0.2 release with std::future support

fix #173

* fixed windows monitor build

* Uses tokio on crates.io

* ctrl_c is a Future

* Run tests with single tokio runtime

* Uses spawn_blocking only with threaded runtime

* Killing all detached Tasks if server is exited

* Updated version to v1.8.0

* Removed unused pin-project

* Uses dns resolver built in tokio

* Explicitly uses threaded scheduler by default

* Reverted, it is threaded scheduler by default

* Updated bytes to v0.5

* Updated many dependencies

* Add a tunnel app (#175)

It can be used to benchmark.

Usage:
sstunnel -v -s 127.0.0.1:8488 -m plain -k password -b 127.0.0.1:1090 -f 127.0.0.1:5201

"-f xxx" : the target/forwad address xxx

* Removed unused remote-dns config field (ssdns is removed)

* Separated TCP local tunneling to its own module

* Implements FromStr for Address

* Fixed typo

* Add tunnel unit test

* Updated tokio to v0.2.2, fixed basic scheduler bug

* Removed Connection, uses try_timeout instead. Revert back to use 8K buffer

* Revert "Removed Connection, uses try_timeout instead. Revert back to use 8K buffer"

This reverts commit 82f698e520.

* Uses Delay instead of Timeout

* Delay doesnt return io::Result

* Revert "Delay doesnt return io::Result"

This reverts commit 74dc294d16.

* Fixed bug, Connection timeout should be shared between read & write operations

.travis.yml

@@ -1,45 +1,45 @@
 sudo: false
 language: rust
 matrix:
-    include:
-        - os: windows
-          rust: stable
+  include:
+    - os: windows
+      rust: stable
 
-        - os: osx
-          rust: stable
+    - os: osx
+      rust: stable
 
-        - os: linux
-          rust: stable
-        - os: linux
-          rust: beta
-        - os: linux
-          rust: nightly
+    - os: linux
+      rust: stable
+    - os: linux
+      rust: beta
+    - os: linux
+      rust: nightly
 
 env:
-    - RUSTFLAGS="-Ctarget-feature=+aes,+ssse3" RUSTDOCFLAGS="-Ctarget-feature=+aes,+ssse3" RUST_BACKTRACE=1
+  - RUSTFLAGS="-Ctarget-feature=+aes,+ssse3" RUSTDOCFLAGS="-Ctarget-feature=+aes,+ssse3" RUST_BACKTRACE=1
 
 addons:
-    apt:
-        packages:
-            - libssl-dev
+  apt:
+    packages:
+      - libssl-dev
 
 install:
-    #- wget https://github.com/jedisct1/libsodium/releases/download/1.0.16/libsodium-1.0.16.tar.gz
-    #- tar xvfz libsodium-1.0.16.tar.gz
-    #- cd libsodium-1.0.16 && ./configure --prefix=$HOME/installed_libsodium && make && make install && cd ..
-    #- export PKG_CONFIG_PATH=$HOME/installed_libsodium/lib/pkgconfig:$PKG_CONFIG_PATH
-    #- export LD_LIBRARY_PATH=$HOME/installed_libsodium/lib:$LD_LIBRARY_PATH
-    - export SODIUM_BUILD_STATIC=yes
+  #- wget https://github.com/jedisct1/libsodium/releases/download/1.0.16/libsodium-1.0.16.tar.gz
+  #- tar xvfz libsodium-1.0.16.tar.gz
+  #- cd libsodium-1.0.16 && ./configure --prefix=$HOME/installed_libsodium && make && make install && cd ..
+  #- export PKG_CONFIG_PATH=$HOME/installed_libsodium/lib/pkgconfig:$PKG_CONFIG_PATH
+  #- export LD_LIBRARY_PATH=$HOME/installed_libsodium/lib:$LD_LIBRARY_PATH
+  - export SODIUM_BUILD_STATIC=yes
 
-    # Install OpenSSL on Windows
-    - if [ "${TRAVIS_OS_NAME}" = "windows" ]; then
-          curl -O http://slproweb.com/download/Win64OpenSSL-1_1_0k.exe;
-          cmd.exe /C 'Win64OpenSSL-1_1_0k.exe /SILENT /VERYSILENT /SP- /DIR="C:\OpenSSL"';
-          export OPENSSL_LIB_DIR=/c/OpenSSL/lib;
-          export OPENSSL_INCLUDE_DIR=/c/OpenSSL/include;
-          export OPENSSL_DIR=/c/OpenSSL/bin;
-      fi
+  # Install OpenSSL on Windows
+  - if [ "${TRAVIS_OS_NAME}" = "windows" ]; then
+    curl -O http://slproweb.com/download/Win64OpenSSL-1_1_0L.exe;
+    cmd.exe /C 'Win64OpenSSL-1_1_0L.exe /SILENT /VERYSILENT /SP- /DIR="C:\OpenSSL"';
+    export OPENSSL_LIB_DIR=/c/OpenSSL/lib;
+    export OPENSSL_INCLUDE_DIR=/c/OpenSSL/include;
+    export OPENSSL_DIR=/c/OpenSSL/bin;
+    fi
 
 script:
-    - cargo build
-    - cargo test --no-fail-fast --features miscreant
+  - cargo build
+  - cargo test --no-fail-fast --features miscreant

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "shadowsocks-rust"
-version = "1.7.2"
+version = "1.8.0"
 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"
@@ -21,12 +21,12 @@ name = "ssserver"
 path = "src/bin/server.rs"
 
 [[bin]]
-name = "ssurl"
-path = "src/bin/ssurl.rs"
+name = "sstunnel"
+path = "src/bin/tunnel.rs"
 
 [[bin]]
-name = "ssdns"
-path = "src/bin/ssdns.rs"
+name = "ssurl"
+path = "src/bin/ssurl.rs"
 
 [profile.release]
 lto = true
@@ -39,41 +39,36 @@ aes-cfb = ["openssl"]
 aes-ctr = ["openssl"]
 camellia-cfb = ["openssl"]
 single-threaded = []
+trust-dns = ["trust-dns-resolver"]
 
 [dependencies]
 log = "0.4"
-byteorder = "1.3"
 rand = "0.7"
 time = "0.1"
 clap = "2.33"
-env_logger = "0.6"
+env_logger = "0.7"
 openssl = { version = "0.10", optional = true }
 libc = "0.2"
-futures = "0.1"
-tokio-io = "0.1"
-tokio = "0.1"
+tokio = { version = "0.2", features = ["full"] }
+futures = "0.3"
 json5 = "0.2"
-base64 = "0.10"
-bytes = "0.4"
-ring = "0.14"
+base64 = "0.11"
+bytes = "0.5"
+byteorder = "1"
+ring = "0.16"
 md-5 = "0.8"
 digest = "0.8"
-typenum = "1.10"
-qrcode = { version = "0.10", default-features = false }
-serde_urlencoded = "0.5"
+typenum = "1.11"
+qrcode = { version = "0.11", default-features = false }
+serde_urlencoded = "0.6"
 serde = { version = "1.0", features = ["derive"] }
-url = "2.0"
+url = "2.1"
 byte_string = "1.0"
-libsodium-ffi = { version = "0.1", optional = true }
+# libsodium-ffi = { version = "0.1", optional = true }
+libsodium-ffi = { git = "https://github.com/zonyitoo/libsodium-ffi", optional = true }
 miscreant = { version = "0.4", optional = true }
-lru-cache = "0.1"
-dns-parser = "0.8"
-trust-dns-resolver = { version = "0.11", features = ["dns-over-rustls", "dns-over-https-rustls"] }
-spin = { version = "0.5", default-features = false }
-tokio-process = "0.2"
-
-[target.'cfg(unix)'.dependencies]
-tokio-signal = "0.2"
-
-[target.'cfg(windows)'.dependencies]
-tokio-signal = "0.2"
+# trust-dns-resolver = { version = "0.11", features = ["dns-over-rustls", "dns-over-https-rustls"] }
+trust-dns-resolver = { git = "https://github.com/bluejekyll/trust-dns", features = ["dns-over-rustls", "dns-over-https-rustls"], optional = true }
+hkdf = "0.8"
+hmac = "0.7"
+sha-1 = "0.8"

src/bin/local.rs

@@ -4,22 +4,23 @@
 //! or you could specify a configuration file. The format of configuration file is defined
 //! in mod `config`.
 
-use std::{io::Result as IoResult, net::SocketAddr, process};
-
 use clap::{App, Arg};
-use futures::{future::Either, Future};
+use futures::{
+    future::{self, Either},
+    FutureExt,
+};
 use log::{debug, error, info};
-#[cfg(feature = "single-threaded")]
-use tokio::runtime::current_thread::Runtime;
-#[cfg(not(feature = "single-threaded"))]
-use tokio::runtime::Runtime;
+use std::net::SocketAddr;
+use tokio;
 
 use shadowsocks::{plugin::PluginConfig, run_local, Config, ConfigType, Mode, ServerAddr, ServerConfig};
 
 mod logging;
 mod monitor;
 
-fn main() {
+#[cfg_attr(feature = "single-threaded", tokio::main(basic_scheduler))]
+#[cfg_attr(not(feature = "single-threaded"), tokio::main)]
+async fn main() {
     let matches = App::new("shadowsocks")
         .version(shadowsocks::VERSION)
         .about("A fast tunnel proxy that helps you bypass firewalls.")
@@ -207,47 +208,11 @@ fn main() {
 
     debug!("Config: {:?}", config);
 
-    match launch_server(config) {
-        Ok(()) => {}
-        Err(err) => {
-            error!("Server exited unexpectly with error: {}", err);
-            process::exit(1);
-        }
-    }
-}
-
-#[cfg(feature = "single-threaded")]
-fn launch_server(config: Config) -> IoResult<()> {
-    let mut runtime = Runtime::new().expect("Creating runtime");
-
     let abort_signal = monitor::create_signal_monitor();
-    let result = runtime.block_on(run_local(config).select2(abort_signal));
-
-    match result {
+    match future::select(run_local(config).boxed(), abort_signal.boxed()).await {
         // Server future resolved without an error. This should never happen.
-        Ok(Either::A(_)) => panic!("Server exited unexpectly"),
-        // Server future resolved with an error.
-        Err(Either::A((err, _))) => Err(err),
+        Either::Left(_) => panic!("Server exited unexpectly"),
         // The abort signal future resolved. Means we should just exit.
-        Ok(Either::B(..)) | Err(Either::B(..)) => Ok(()),
-    }
-}
-
-#[cfg(not(feature = "single-threaded"))]
-fn launch_server(config: Config) -> IoResult<()> {
-    let mut runtime = Runtime::new().expect("Creating runtime");
-
-    let abort_signal = monitor::create_signal_monitor();
-    let result = runtime.block_on(run_local(config).select2(abort_signal));
-
-    runtime.shutdown_now().wait().unwrap();
-
-    match result {
-        // Server future resolved without an error. This should never happen.
-        Ok(Either::A(_)) => panic!("Server exited unexpectly"),
-        // Server future resolved with an error.
-        Err(Either::A((err, _))) => Err(err),
-        // The abort signal future resolved. Means we should just exit.
-        Ok(Either::B(..)) | Err(Either::B(..)) => Ok(()),
+        Either::Right(_) => (),
     }
 }

src/bin/logging/mod.rs

@@ -1,9 +1,8 @@
+use std::env;
+use std::io::{self, Write};
+
 use env_logger::{fmt::Formatter, Builder};
 use log::{LevelFilter, Record};
-use std::{
-    env,
-    io::{self, Write},
-};
 
 pub fn init(without_time: bool, debug_level: u64, bin_name: &str) {
     let mut log_builder = Builder::new();

src/bin/monitor/other.rs

@@ -1,8 +1,8 @@
-use futures::{self, Future};
+use futures;
 use std::io;
 
-pub fn create_signal_monitor() -> impl Future<Item = (), Error = io::Error> + Send {
+pub async fn create_signal_monitor() -> io::Result<()> {
     // FIXME: What can I do ...
     // Blocks forever
-    futures::empty::<(), io::Error>()
+    futures::empty::<(), io::Error>().await
 }

src/bin/monitor/unix.rs

@@ -1,34 +1,19 @@
-use futures::{Future, Stream};
-use libc;
-use log::{error, info};
+use futures::future::{self, Either, FutureExt};
+use log::info;
 use std::io;
-use tokio_signal::unix::Signal;
+use tokio::signal::unix::{signal, SignalKind};
 
-pub fn create_signal_monitor() -> impl Future<Item = (), Error = io::Error> + Send {
+pub async fn create_signal_monitor() -> io::Result<()> {
     // Future resolving to two signal streams. Can fail if setting up signal monitoring fails
-    let signals_future = Signal::new(libc::SIGTERM).join(Signal::new(libc::SIGINT));
+    let mut sigterm = signal(SignalKind::terminate())?;
+    let mut sigint = signal(SignalKind::interrupt())?;
 
-    signals_future
-       // The future resolved into our two streams of signals
-       .and_then(|(sigterms, sigints)| {
-           // Stream of all signals we care about
-           let signals = sigterms.select(sigints);
-           // Take only the first signal in the stream and log that it was triggered
-           signals.take(1).for_each(|signal| {
-               let signal_name = match signal {
-                   libc::SIGTERM => "SIGTERM",
-                   libc::SIGINT => "SIGINT",
-                   _ => unreachable!(),
-               };
-               info!("Received {}, exiting", signal_name);
-               Ok(())
-           }).map_err(|error| {
-               error!("Error while listening on process signals: {:?}", error);
-               error
-           })
-       })
-       .map_err(|error| {
-           error!("Failed to set up process signal monitoring: {:?}", error);
-           error
-       })
+    let signal_name = match future::select(sigterm.recv().boxed(), sigint.recv().boxed()).await {
+        Either::Left(..) => "SIGTERM",
+        Either::Right(..) => "SIGINT",
+    };
+
+    info!("Received {}, exiting", signal_name);
+
+    Ok(())
 }

src/bin/monitor/windows.rs

@@ -1,35 +1,19 @@
-use futures::{self, Future, Stream};
-use log::{error, info};
+use futures::future::{self, Either, FutureExt};
+use futures::StreamExt;
+use log::info;
 use std::io;
-use tokio_signal::windows::Event;
+use tokio::signal::{ctrl_c, windows::ctrl_break};
 
-pub fn create_signal_monitor() -> impl Future<Item = (), Error = io::Error> + Send {
-    // Must be wrapped with a future, because when `Event` is being created, it will spawn
-    // a task into the current executor. So if there is no current running executor, thread
-    // will panic.
-    let events = futures::lazy(|| {
-        Event::ctrl_c()
-            .join(Event::ctrl_break())
-            .and_then(|(c_signals, break_signals)| {
-                // Stream of all signals we care about
-                let signals = c_signals.select(break_signals);
-                // Take only the first signal in the stream and log that it was triggered
-                signals
-                    .take(1)
-                    .for_each(|()| {
-                        info!("Received shutdown signal, exiting");
-                        Ok(())
-                    })
-                    .map_err(|error| {
-                        error!("Error while listening on process signals: {:?}", error);
-                        error
-                    })
-            })
-            .map_err(|error| {
-                error!("Failed to set up process signal monitoring: {:?}", error);
-                error
-            })
-    });
+pub async fn create_signal_monitor() -> io::Result<()> {
+    let cc = ctrl_c();
+    let cb = ctrl_break()?;
 
-    events
+    let signal_name = match future::select(cc.boxed(), cb.into_future().boxed()).await {
+        Either::Left(..) => "CTRL-C",
+        Either::Right(..) => "CTRL-BREAK",
+    };
+
+    info!("Received {}, exiting", signal_name);
+
+    Ok(())
 }

src/bin/server.rs

@@ -7,22 +7,22 @@
 //! *It should be notice that the extented configuration file is not suitable for the server
 //! side.*
 
-use std::{io::Result as IoResult, process};
-
 use clap::{App, Arg};
-use futures::{future::Either, Future};
+use futures::{
+    future::{self, Either},
+    FutureExt,
+};
 use log::{debug, error, info};
-#[cfg(feature = "single-threaded")]
-use tokio::runtime::current_thread::Runtime;
-#[cfg(not(feature = "single-threaded"))]
-use tokio::runtime::Runtime;
+use tokio;
 
 use shadowsocks::{plugin::PluginConfig, run_server, Config, ConfigType, Mode, ServerAddr, ServerConfig};
 
 mod logging;
 mod monitor;
 
-fn main() {
+#[cfg_attr(feature = "single-threaded", tokio::main(basic_scheduler))]
+#[cfg_attr(not(feature = "single-threaded"), tokio::main)]
+async fn main() {
     let matches = App::new("shadowsocks")
         .version(shadowsocks::VERSION)
         .about("A fast tunnel proxy that helps you bypass firewalls.")
@@ -191,47 +191,11 @@ fn main() {
 
     debug!("Config: {:?}", config);
 
-    match launch_server(config) {
-        Ok(()) => {}
-        Err(err) => {
-            error!("Server exited unexpectly with error: {}", err);
-            process::exit(1);
-        }
-    }
-}
-
-#[cfg(not(feature = "single-threaded"))]
-fn launch_server(config: Config) -> IoResult<()> {
-    let mut runtime = Runtime::new().expect("Creating runtime");
-
     let abort_signal = monitor::create_signal_monitor();
-    let result = runtime.block_on(run_server(config).select2(abort_signal));
-
-    runtime.shutdown_now().wait().unwrap();
-
-    match result {
+    match future::select(run_server(config).boxed(), abort_signal.boxed()).await {
         // Server future resolved without an error. This should never happen.
-        Ok(Either::A(_)) => panic!("Server exited unexpectly"),
-        // Server future resolved with an error.
-        Err(Either::A((err, _))) => Err(err),
+        Either::Left(_) => panic!("Server exited unexpectly"),
         // The abort signal future resolved. Means we should just exit.
-        Ok(Either::B(..)) | Err(Either::B(..)) => Ok(()),
-    }
-}
-
-#[cfg(feature = "single-threaded")]
-fn launch_server(config: Config) -> IoResult<()> {
-    let mut runtime = Runtime::new().expect("Creating runtime");
-
-    let abort_signal = monitor::create_signal_monitor();
-    let result = runtime.block_on(run_server(config).select2(abort_signal));
-
-    match result {
-        // Server future resolved without an error. This should never happen.
-        Ok(Either::A(_)) => panic!("Server exited unexpectly"),
-        // Server future resolved with an error.
-        Err(Either::A((err, _))) => Err(err),
-        // The abort signal future resolved. Means we should just exit.
-        Ok(Either::B(..)) | Err(Either::B(..)) => Ok(()),
+        Either::Right(_) => (),
     }
 }

src/bin/tunnel.rs

@@ -1,26 +1,39 @@
-//! DNS over shadowsocks
-
-use std::net::SocketAddr;
+//! This is a binary running in the local environment
+//!
+//! You have to provide all needed configuration attributes via command line parameters,
+//! or you could specify a configuration file. The format of configuration file is defined
+//! in mod `config`.
 
 use clap::{App, Arg};
-use futures::Future;
+use futures::{
+    future::{self, Either},
+    FutureExt,
+};
 use log::{debug, error, info};
-use tokio::runtime::Runtime;
+use std::net::SocketAddr;
+use tokio;
 
-use shadowsocks::{run_dns, Config, ConfigType, ServerAddr, ServerConfig};
+use shadowsocks::plugin::PluginConfig;
+use shadowsocks::relay::socks5::Address;
+use shadowsocks::{run_local, Config, ConfigType, Mode, ServerAddr, ServerConfig};
 
 mod logging;
+mod monitor;
 
-fn main() {
-    let matches = App::new("ssdns")
+#[cfg_attr(feature = "single-threaded", tokio::main(basic_scheduler))]
+#[cfg_attr(not(feature = "single-threaded"), tokio::main)]
+async fn main() {
+    let matches = App::new("shadowsocks")
         .version(shadowsocks::VERSION)
-        .about("A DNS proxy that helps you bypass firewalls.")
+        .about("A fast tunnel proxy that helps you bypass firewalls.")
         .arg(
             Arg::with_name("VERBOSE")
                 .short("v")
                 .multiple(true)
                 .help("Set the level of debug"),
         )
+        .arg(Arg::with_name("UDP_ONLY").short("u").help("Server mode UDP_ONLY"))
+        .arg(Arg::with_name("TCP_AND_UDP").short("U").help("Server mode TCP_AND_UDP"))
         .arg(
             Arg::with_name("CONFIG")
                 .short("c")
@@ -42,6 +55,13 @@ fn main() {
                 .takes_value(true)
                 .help("Local address, listen only to this address if specified"),
         )
+        .arg(
+            Arg::with_name("FORWARD_ADDR")
+                .short("f")
+                .long("foward-addr")
+                .takes_value(true)
+                .help("Forward address, forward to this address"),
+        )
         .arg(
             Arg::with_name("PASSWORD")
                 .short("k")
@@ -56,6 +76,18 @@ fn main() {
                 .takes_value(true)
                 .help("Encryption method"),
         )
+        .arg(
+            Arg::with_name("PLUGIN")
+                .long("plugin")
+                .takes_value(true)
+                .help("Enable SIP003 plugin"),
+        )
+        .arg(
+            Arg::with_name("PLUGIN_OPT")
+                .long("plugin-opts")
+                .takes_value(true)
+                .help("Set SIP003 plugin options"),
+        )
         .arg(
             Arg::with_name("LOG_WITHOUT_TIME")
                 .long("log-without-time")
@@ -68,17 +100,17 @@ fn main() {
                 .help("Server address in SIP002 URL"),
         )
         .arg(
-            Arg::with_name("REMOTE_DNS")
-                .long("remote-dns")
-                .takes_value(true)
-                .help("Remote DNS server, default is 8.8.8.8:53"),
+            Arg::with_name("NO_DELAY")
+                .long("no-delay")
+                .takes_value(false)
+                .help("Set no-delay option for socket"),
         )
         .get_matches();
 
     let without_time = matches.is_present("LOG_WITHOUT_TIME");
     let debug_level = matches.occurrences_of("VERBOSE");
 
-    logging::init(without_time, debug_level, "ssdns");
+    logging::init(without_time, debug_level, "sslocal");
 
     let mut has_provided_config = false;
 
@@ -153,22 +185,49 @@ fn main() {
         return;
     }
 
-    if let Some(dns) = matches.value_of("REMOTE_DNS") {
-        let dns_addr = dns
-            .parse::<SocketAddr>()
-            .expect("`remote-dns` is not a valid SocketAddr, must be IP:Port");
-        config.remote_dns = Some(dns_addr);
+    if let Some(url) = matches.value_of("FORWARD_ADDR") {
+        let forward_addr = url.parse::<Address>().expect("Failed to parse `url`");
+
+        config.forward = Some(forward_addr);
     }
 
-    info!("ShadowSocks DNS {}", shadowsocks::VERSION);
+    if matches.is_present("UDP_ONLY") {
+        if config.mode.enable_tcp() {
+            config.mode = Mode::TcpAndUdp;
+        } else {
+            config.mode = Mode::UdpOnly;
+        }
+    }
+
+    if matches.is_present("TCP_AND_UDP") {
+        config.mode = Mode::TcpAndUdp;
+    }
+
+    if matches.is_present("NO_DELAY") {
+        config.no_delay = true;
+    }
+
+    if let Some(p) = matches.value_of("PLUGIN") {
+        let plugin = PluginConfig {
+            plugin: p.to_owned(),
+            plugin_opt: matches.value_of("PLUGIN_OPT").map(ToOwned::to_owned),
+        };
+
+        // Overrides config in file
+        for svr in config.server.iter_mut() {
+            svr.set_plugin(plugin.clone());
+        }
+    };
+
+    info!("ShadowSocks {}", shadowsocks::VERSION);
 
     debug!("Config: {:?}", config);
 
-    let mut runtime = Runtime::new().expect("Creating runtime");
-
-    let result = runtime.block_on(run_dns(config));
-
-    runtime.shutdown_now().wait().unwrap();
-
-    panic!("Server exited unexpectly with result: {:?}", result);
+    let abort_signal = monitor::create_signal_monitor();
+    match future::select(run_local(config).boxed(), abort_signal.boxed()).await {
+        // Server future resolved without an error. This should never happen.
+        Either::Left(_) => panic!("Server exited unexpectly"),
+        // The abort signal future resolved. Means we should just exit.
+        Either::Right(_) => (),
+    }
 }

src/config.rs

@@ -41,32 +41,33 @@
 //!
 //! These defined server will be used with a load balancing algorithm.
 
-use std::{
-    collections::HashSet,
-    convert::From,
-    default::Default,
-    error,
-    fmt::{self, Debug, Display, Formatter},
-    fs::OpenOptions,
-    io::Read,
-    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6},
-    option::Option,
-    path::Path,
-    str::FromStr,
-    string::ToString,
-    time::Duration,
-};
+use std::collections::HashSet;
+use std::convert::From;
+use std::default::Default;
+use std::error;
+use std::fmt::{self, Debug, Display, Formatter};
+use std::fs::OpenOptions;
+use std::io::Read;
+use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
+use std::option::Option;
+use std::path::Path;
+use std::str::FromStr;
+use std::string::ToString;
+use std::time::Duration;
 
 use base64::{decode_config, encode_config, URL_SAFE_NO_PAD};
 use bytes::Bytes;
 use json5;
-use log::{error, trace};
+use log::error;
 use serde::{Deserialize, Serialize};
 use serde_urlencoded;
+#[cfg(feature = "trust-dns")]
 use trust_dns_resolver::config::{NameServerConfigGroup, ResolverConfig};
 use url::{self, Url};
 
-use crate::{crypto::cipher::CipherType, plugin::PluginConfig};
+use crate::crypto::cipher::CipherType;
+use crate::plugin::PluginConfig;
+use crate::relay::socks5::Address;
 
 #[derive(Serialize, Deserialize, Debug, Default)]
 struct SSConfig {
@@ -97,8 +98,6 @@ struct SSConfig {
     #[serde(skip_serializing_if = "Option::is_none")]
     dns: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
-    remote_dns: Option<String>,
-    #[serde(skip_serializing_if = "Option::is_none")]
     mode: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     no_delay: Option<bool>,
@@ -220,7 +219,6 @@ impl ServerConfig {
         plugin: Option<PluginConfig>,
     ) -> ServerConfig {
         let enc_key = method.bytes_to_key(pwd.as_bytes());
-        trace!("Initialize config with pwd: {:?}, key: {:?}", pwd, enc_key);
         ServerConfig {
             addr,
             password: pwd,
@@ -461,9 +459,9 @@ impl error::Error for UrlParseError {
         }
     }
 
-    fn cause(&self) -> Option<&error::Error> {
+    fn cause(&self) -> Option<&dyn error::Error> {
         match *self {
-            UrlParseError::ParseError(ref err) => Some(err as &error::Error),
+            UrlParseError::ParseError(ref err) => Some(err as &dyn error::Error),
             UrlParseError::InvalidScheme => None,
             UrlParseError::InvalidUserInfo => None,
             UrlParseError::MissingHost => None,
@@ -540,9 +538,9 @@ impl FromStr for Mode {
 pub struct Config {
     pub server: Vec<ServerConfig>,
     pub local: Option<ClientConfig>,
+    pub forward: Option<Address>,
     pub forbidden_ip: HashSet<IpAddr>,
     pub dns: Option<String>,
-    pub remote_dns: Option<SocketAddr>,
     pub mode: Mode,
     pub no_delay: bool,
     pub manager_address: Option<ServerAddr>,
@@ -600,9 +598,9 @@ impl Config {
         Config {
             server: Vec::new(),
             local: None,
+            forward: None,
             forbidden_ip: HashSet::new(),
             dns: None,
-            remote_dns: None,
             mode: Mode::TcpOnly,
             no_delay: false,
             manager_address: None,
@@ -761,20 +759,6 @@ impl Config {
         // DNS
         nconfig.dns = config.dns;
 
-        if let Some(rdns) = config.remote_dns {
-            match rdns.parse::<SocketAddr>() {
-                Ok(r) => nconfig.remote_dns = Some(r),
-                Err(..) => {
-                    let e = Error::new(
-                        ErrorKind::Malformed,
-                        "malformed `remote_dns`, which must be a valid SocketAddr",
-                        None,
-                    );
-                    return Err(e);
-                }
-            }
-        }
-
         // Mode
         if let Some(m) = config.mode {
             match m.parse::<Mode>() {
@@ -810,6 +794,7 @@ impl Config {
         Config::load_from_str(&content[..], config_type)
     }
 
+    #[cfg(feature = "trust-dns")]
     pub fn get_dns_config(&self) -> Option<ResolverConfig> {
         self.dns.as_ref().and_then(|ds| {
             match &ds[..] {
@@ -843,11 +828,14 @@ impl Config {
         })
     }
 
-    pub fn get_remote_dns(&self) -> SocketAddr {
-        match self.remote_dns {
-            None => SocketAddr::from(SocketAddrV4::new(Ipv4Addr::new(8, 8, 8, 8), 53)),
-            Some(ip) => ip,
+    /// Check if there are any plugin are enabled with servers
+    pub fn has_server_plugins(&self) -> bool {
+        for server in &self.server {
+            if server.plugin().is_some() {
+                return true;
+            }
         }
+        false
     }
 }
 
@@ -922,10 +910,6 @@ impl fmt::Display for Config {
             jconf.dns = Some(dns.to_string());
         }
 
-        if let Some(ref remote_dns) = self.remote_dns {
-            jconf.remote_dns = Some(remote_dns.to_string());
-        }
-
         write!(f, "{}", json5::to_string(&jconf).unwrap())
     }
 }

src/context.rs

@@ -1,59 +1,93 @@
 //! Shadowsocks Server Context
 
-use std::{
-    net::SocketAddr,
-    sync::{Arc, Mutex, MutexGuard},
-    time::Instant,
-};
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
 
-use lru_cache::LruCache;
+#[cfg(futures = "trust-dns")]
 use trust_dns_resolver::AsyncResolver;
 
-use crate::{config::Config, relay::dns_resolver::create_resolver};
+use crate::config::Config;
 
-type DnsQueryCache = LruCache<u16, (SocketAddr, Instant)>;
+#[cfg(futures = "trust-dns")]
+use crate::relay::dns_resolver::create_resolver;
 
 #[derive(Clone)]
-pub struct Context {
-    config: Config,
+pub struct SharedServerState {
+    #[cfg(futures = "trust-dns")]
     dns_resolver: Arc<AsyncResolver>,
-    dns_query_cache: Option<Arc<Mutex<DnsQueryCache>>>,
+    server_running: Arc<AtomicBool>,
 }
 
+impl SharedServerState {
+    #[allow(unused_variables)]
+    pub fn new(config: &Config) -> SharedServerState {
+        SharedServerState {
+            #[cfg(futures = "trust-dns")]
+            dns_resolver: Arc::new(create_resolver(config.get_dns_config())),
+            server_running: Arc::new(AtomicBool::new(true)),
+        }
+    }
+
+    /// Check if the server is still in running state
+    pub fn server_running(&self) -> bool {
+        self.server_running.load(Ordering::Acquire)
+    }
+
+    /// Stops the server, kills all detached running tasks
+    pub fn server_stopped(&self) {
+        self.server_running.store(false, Ordering::Release)
+    }
+
+    /// Get the global shared resolver
+    #[cfg(futures = "trust-dns")]
+    pub fn dns_resolver(&self) -> &AsyncResolver {
+        &*self.dns_resolver
+    }
+}
+
+/// Shared basic configuration for the whole server
+pub struct Context {
+    config: Config,
+    server_state: SharedServerState,
+}
+
+/// Unique context thw whole server
 pub type SharedContext = Arc<Context>;
 
 impl Context {
-    pub fn new(config: Config) -> Context {
-        let resolver = create_resolver(config.get_dns_config());
-        Context {
-            config,
-            dns_resolver: Arc::new(resolver),
-            dns_query_cache: None,
-        }
+    pub fn new(config: Config, server_state: SharedServerState) -> Context {
+        Context { config, server_state }
     }
 
-    pub fn new_dns(config: Config) -> Context {
-        let resolver = create_resolver(config.get_dns_config());
-        Context {
-            config,
-            dns_resolver: Arc::new(resolver),
-            dns_query_cache: Some(Arc::new(Mutex::new(LruCache::new(1024)))),
-        }
+    pub fn new_shared(config: Config, server_state: SharedServerState) -> SharedContext {
+        SharedContext::new(Context::new(config, server_state))
     }
 
+    /// Config for TCP server
     pub fn config(&self) -> &Config {
         &self.config
     }
 
+    /// Mutable Config for TCP server
+    ///
+    /// NOTE: Only for launching plugins
     pub fn config_mut(&mut self) -> &mut Config {
         &mut self.config
     }
 
+    /// Get the global shared resolver
+    #[cfg(futures = "trust-dns")]
     pub fn dns_resolver(&self) -> &AsyncResolver {
-        &*self.dns_resolver
+        self.server_state.dns_resolver()
     }
 
-    pub fn dns_query_cache(&self) -> MutexGuard<DnsQueryCache> {
-        self.dns_query_cache.as_ref().unwrap().lock().unwrap()
+    /// Check if the server is still in running state
+    pub fn server_running(&self) -> bool {
+        self.server_state.server_running()
+    }
+
+    /// Stops the server, kills all detached running tasks
+    pub fn server_stopped(&self) {
+        self.server_state.server_stopped()
     }
 }

src/crypto/aead.rs

@@ -9,7 +9,8 @@ use crate::crypto::siv::MiscreantCipher;
 use crate::crypto::sodium::SodiumAeadCipher;
 
 use bytes::{Bytes, BytesMut};
-use ring::{digest::SHA1, hkdf, hmac::SigningKey};
+use hkdf::Hkdf;
+use sha1::Sha1;
 
 /// Encryptor API for AEAD ciphers
 pub trait AeadEncryptor {
@@ -110,14 +111,14 @@ const SUBKEY_INFO: &[u8] = b"ss-subkey";
 pub fn make_skey(t: CipherType, key: &[u8], salt: &[u8]) -> Bytes {
     assert!(t.category() == CipherCategory::Aead);
 
-    let salt = SigningKey::new(&SHA1, salt);
+    let hkdf = Hkdf::<Sha1>::new(Some(salt), key);
 
     let mut skey = BytesMut::with_capacity(key.len());
     unsafe {
         skey.set_len(key.len());
     }
 
-    hkdf::extract_and_expand(&salt, key, SUBKEY_INFO, &mut skey);
+    hkdf.expand(SUBKEY_INFO, &mut skey).unwrap();
 
     skey.freeze()
 }

src/crypto/digest.rs

@@ -1,7 +1,7 @@
 //! Message digest algorithm
 
 use md5::Md5;
-use ring::digest::{Context, SHA1, SHA1_OUTPUT_LEN};
+use sha1::Sha1;
 
 use bytes::BufMut;
 
@@ -32,14 +32,14 @@ pub enum DigestType {
 pub fn with_type(t: DigestType) -> DigestVariant {
     match t {
         DigestType::Md5 => DigestVariant::Md5(Md5::default()),
-        DigestType::Sha1 | DigestType::Sha => DigestVariant::Sha1(Context::new(&SHA1)),
+        DigestType::Sha1 | DigestType::Sha => DigestVariant::Sha1(Sha1::default()),
     }
 }
 
 /// Variant of supported digest
 pub enum DigestVariant {
     Md5(Md5),
-    Sha1(Context),
+    Sha1(Sha1),
 }
 
 impl Digest for DigestVariant {
@@ -48,7 +48,7 @@ impl Digest for DigestVariant {
 
         match *self {
             DigestVariant::Md5(ref mut d) => d.input(data),
-            DigestVariant::Sha1(ref mut d) => d.update(data),
+            DigestVariant::Sha1(ref mut d) => d.input(data),
         }
     }
 
@@ -57,7 +57,7 @@ impl Digest for DigestVariant {
 
         match *self {
             DigestVariant::Md5(ref d) => buf.put(&*d.clone().result()),
-            DigestVariant::Sha1(ref d) => buf.put(d.clone().finish().as_ref()),
+            DigestVariant::Sha1(ref d) => buf.put(&*d.clone().result()),
         }
     }
 
@@ -67,14 +67,14 @@ impl Digest for DigestVariant {
 
         match *self {
             DigestVariant::Md5(_) => <Md5 as FixedOutput>::OutputSize::to_usize(),
-            DigestVariant::Sha1(_) => SHA1_OUTPUT_LEN,
+            DigestVariant::Sha1(_) => <Md5 as FixedOutput>::OutputSize::to_usize(),
         }
     }
 
     fn reset(&mut self) {
         match *self {
             DigestVariant::Md5(ref mut d) => d.clone_from(&Md5::default()),
-            DigestVariant::Sha1(ref mut d) => d.clone_from(&Context::new(&SHA1)),
+            DigestVariant::Sha1(ref mut d) => d.clone_from(&Sha1::default()),
         }
     }
 }

src/crypto/dummy.rs

@@ -10,12 +10,12 @@ use bytes::BufMut;
 pub struct DummyCipher;
 
 impl StreamCipher for DummyCipher {
-    fn update(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()> {
+    fn update(&mut self, data: &[u8], out: &mut dyn BufMut) -> CipherResult<()> {
         out.put_slice(data);
         Ok(())
     }
 
-    fn finalize(&mut self, _: &mut BufMut) -> CipherResult<()> {
+    fn finalize(&mut self, _: &mut dyn BufMut) -> CipherResult<()> {
         Ok(())
     }
 

src/crypto/openssl.rs

@@ -109,7 +109,7 @@ impl OpenSSLCrypto {
     }
 
     /// Update data
-    pub fn update(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()> {
+    pub fn update(&mut self, data: &[u8], out: &mut dyn BufMut) -> CipherResult<()> {
         let least_reserved = data.len() + self.cipher.block_size();
         let mut buf = BytesMut::with_capacity(least_reserved); // NOTE: len() is 0 now!
         unsafe {
@@ -122,7 +122,7 @@ impl OpenSSLCrypto {
     }
 
     /// Generate the final block
-    pub fn finalize(&mut self, out: &mut BufMut) -> CipherResult<()> {
+    pub fn finalize(&mut self, out: &mut dyn BufMut) -> CipherResult<()> {
         let least_reserved = self.cipher.block_size();
         let mut buf = BytesMut::with_capacity(least_reserved); // NOTE: len() is 0 now!
         unsafe {
@@ -185,11 +185,11 @@ impl OpenSSLCipher {
 unsafe impl Send for OpenSSLCipher {}
 
 impl StreamCipher for OpenSSLCipher {
-    fn update(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()> {
+    fn update(&mut self, data: &[u8], out: &mut dyn BufMut) -> CipherResult<()> {
         self.worker.update(data, out)
     }
 
-    fn finalize(&mut self, out: &mut BufMut) -> CipherResult<()> {
+    fn finalize(&mut self, out: &mut dyn BufMut) -> CipherResult<()> {
         self.worker.finalize(out)
     }
 

src/crypto/rc4_md5.rs

@@ -28,11 +28,11 @@ impl Rc4Md5Cipher {
 }
 
 impl StreamCipher for Rc4Md5Cipher {
-    fn update(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()> {
+    fn update(&mut self, data: &[u8], out: &mut dyn BufMut) -> CipherResult<()> {
         self.crypto.update(data, out)
     }
 
-    fn finalize(&mut self, out: &mut BufMut) -> CipherResult<()> {
+    fn finalize(&mut self, out: &mut dyn BufMut) -> CipherResult<()> {
         self.crypto.finalize(out)
     }
 

src/crypto/ring.rs

@@ -1,9 +1,11 @@
 //! Cipher defined with Ring
 
-use std::{mem, ptr};
-
-use ring::aead::{
-    open_in_place, seal_in_place, Aad, Nonce, OpeningKey, SealingKey, AES_128_GCM, AES_256_GCM, CHACHA20_POLY1305,
+use ring::{
+    aead::{
+        Aad, Algorithm, Nonce, NonceSequence, OpeningKey, SealingKey, UnboundKey, AES_128_GCM, AES_256_GCM,
+        CHACHA20_POLY1305, NONCE_LEN,
+    },
+    error::Unspecified,
 };
 
 use crate::crypto::{
@@ -13,13 +15,33 @@ use crate::crypto::{
 };
 
 use byte_string::ByteStr;
-use bytes::{BufMut, Bytes, BytesMut};
+use bytes::{BufMut, BytesMut};
 use log::error;
 
 /// AEAD ciphers provided by Ring
 pub enum RingAeadCryptoVariant {
-    Seal(SealingKey, Bytes),
-    Open(OpeningKey, Bytes),
+    Seal(SealingKey<RingAeadNonceSequence>),
+    Open(OpeningKey<RingAeadNonceSequence>),
+}
+
+pub struct RingAeadNonceSequence {
+    nonce: [u8; NONCE_LEN],
+}
+
+impl RingAeadNonceSequence {
+    fn new() -> RingAeadNonceSequence {
+        RingAeadNonceSequence {
+            nonce: [0u8; NONCE_LEN],
+        }
+    }
+}
+
+impl NonceSequence for RingAeadNonceSequence {
+    fn advance(&mut self) -> Result<Nonce, Unspecified> {
+        let nonce = Nonce::assume_unique_for_key(self.nonce);
+        increase_nonce(&mut self.nonce);
+        Ok(nonce)
+    }
 }
 
 /// AEAD Cipher context
@@ -28,8 +50,6 @@ pub enum RingAeadCryptoVariant {
 pub struct RingAeadCipher {
     cipher: RingAeadCryptoVariant,
     cipher_type: CipherType,
-    key: Bytes,
-    nonce: BytesMut,
 }
 
 impl RingAeadCipher {
@@ -37,57 +57,34 @@ impl RingAeadCipher {
     pub fn new(t: CipherType, key: &[u8], salt: &[u8], is_seal: bool) -> RingAeadCipher {
         // TODO: Check if salt is duplicated
 
-        let nonce_size = t.iv_size();
-        let mut nonce = BytesMut::with_capacity(nonce_size);
-        unsafe {
-            nonce.set_len(nonce_size);
-            ptr::write_bytes(nonce.as_mut_ptr(), 0, nonce_size);
-        }
+        // Nonce is 12 bytes
+        assert_eq!(t.iv_size(), NONCE_LEN);
 
         let skey = make_skey(t, key, salt);
-        let cipher = RingAeadCipher::new_variant(t, &skey, &nonce, is_seal);
-        RingAeadCipher {
-            cipher,
-            cipher_type: t,
-            key: skey,
-            nonce,
-        }
+        let cipher = RingAeadCipher::new_variant(t, &skey, is_seal);
+        RingAeadCipher { cipher, cipher_type: t }
     }
 
-    fn new_variant(t: CipherType, key: &[u8], nonce: &[u8], is_seal: bool) -> RingAeadCryptoVariant {
-        macro_rules! seal_or_open {
-            ($item:ident, $key:ident, $crypt:ident) => {
-                RingAeadCryptoVariant::$item($key::new(&$crypt, key).unwrap(), Bytes::from(nonce))
-            };
-            ($crypt:ident) => {
-                if is_seal {
-                    seal_or_open!(Seal, SealingKey, $crypt)
-                } else {
-                    seal_or_open!(Open, OpeningKey, $crypt)
-                }
-            };
-        }
-
+    fn new_variant(t: CipherType, key: &[u8], is_seal: bool) -> RingAeadCryptoVariant {
         match t {
-            CipherType::Aes128Gcm => seal_or_open!(AES_128_GCM),
-            CipherType::Aes256Gcm => seal_or_open!(AES_256_GCM),
-            CipherType::ChaCha20IetfPoly1305 => seal_or_open!(CHACHA20_POLY1305),
+            CipherType::Aes128Gcm => RingAeadCipher::new_crypt(&AES_128_GCM, key, is_seal),
+            CipherType::Aes256Gcm => RingAeadCipher::new_crypt(&AES_256_GCM, key, is_seal),
+            CipherType::ChaCha20IetfPoly1305 => RingAeadCipher::new_crypt(&CHACHA20_POLY1305, key, is_seal),
             _ => panic!("unsupported cipher in ring {:?}", t),
         }
     }
 
-    fn increase_nonce(&mut self) {
-        increase_nonce(&mut self.nonce);
-    }
+    #[inline]
+    fn new_crypt(algorithm: &'static Algorithm, key: &[u8], is_seal: bool) -> RingAeadCryptoVariant {
+        use ring::aead::BoundKey;
 
-    fn reset(&mut self) {
-        self.increase_nonce();
-        let is_seal = match self.cipher {
-            RingAeadCryptoVariant::Seal(..) => true,
-            RingAeadCryptoVariant::Open(..) => false,
-        };
-        let var = RingAeadCipher::new_variant(self.cipher_type, &self.key, &self.nonce, is_seal);
-        mem::replace(&mut self.cipher, var);
+        let unbound_key = UnboundKey::new(algorithm, key).unwrap();
+
+        if is_seal {
+            RingAeadCryptoVariant::Seal(SealingKey::new(unbound_key, RingAeadNonceSequence::new()))
+        } else {
+            RingAeadCryptoVariant::Open(OpeningKey::new(unbound_key, RingAeadNonceSequence::new()))
+        }
     }
 }
 
@@ -99,26 +96,14 @@ impl AeadEncryptor for RingAeadCipher {
 
         let mut buf = BytesMut::with_capacity(output.len());
         buf.put_slice(input);
-        unsafe {
-            buf.set_len(output.len());
-        }
 
-        if let RingAeadCryptoVariant::Seal(ref key, ref nonce) = self.cipher {
-            seal_in_place(
-                key,
-                Nonce::try_assume_unique_for_key(nonce).expect("AEAD cipher nonce length not match"),
-                Aad::empty(),
-                &mut buf,
-                tag_len,
-            )
-            .unwrap();
+        if let RingAeadCryptoVariant::Seal(ref mut key) = self.cipher {
+            key.seal_in_place_append_tag(Aad::empty(), &mut buf).unwrap();
         } else {
             unreachable!("encrypt is called on a non-seal cipher");
         }
 
-        output.copy_from_slice(&buf);
-
-        self.reset();
+        output.copy_from_slice(&buf[..buf_len]);
     }
 }
 
@@ -130,36 +115,25 @@ impl AeadDecryptor for RingAeadCipher {
         let mut buf = BytesMut::with_capacity(input.len());
         buf.put_slice(input);
 
-        let r = if let RingAeadCryptoVariant::Open(ref key, ref nonce) = self.cipher {
-            match open_in_place(
-                key,
-                Nonce::try_assume_unique_for_key(nonce).expect("AEAD cipher nonce length not match"),
-                Aad::empty(),
-                0,
-                &mut buf,
-            ) {
+        if let RingAeadCryptoVariant::Open(ref mut key) = self.cipher {
+            match key.open_in_place(Aad::empty(), &mut buf) {
                 Ok(obuf) => {
                     output.copy_from_slice(obuf);
                     Ok(())
                 }
-                Err(err) => {
+                Err(..) => {
                     error!(
-                        "AEAD decrypt failed, nonce={:?}, input={:?}, tag={:?}, err: {:?}",
-                        ByteStr::new(nonce),
-                        ByteStr::new(&input[..tag_len]),
-                        ByteStr::new(&input[tag_len..]),
-                        err
+                        "AEAD decrypt failed, input={:?}, tag={:?}, opening: {:?}",
+                        ByteStr::new(&input[..output.len()]),
+                        ByteStr::new(&input[output.len()..]),
+                        key,
                     );
                     Err(Error::AeadDecryptFailed)
                 }
             }
         } else {
             unreachable!("decrypt is called on a non-open cipher");
-        };
-
-        self.reset();
-
-        r
+        }
     }
 }
 

src/crypto/sodium.rs

@@ -1,9 +1,6 @@
 //! Cipher defined with libsodium
 
-use std::{
-    ptr,
-    sync::{Once, ONCE_INIT},
-};
+use std::{ptr, sync::Once};
 
 use bytes::{BufMut, Bytes, BytesMut};
 
@@ -20,7 +17,7 @@ use crate::crypto::{
     AeadDecryptor, AeadEncryptor, CipherResult, CipherType, StreamCipher,
 };
 
-static SODIUM_INIT_FLAG: Once = ONCE_INIT;
+static SODIUM_INIT_FLAG: Once = Once::new();
 
 /// Cipher provided by `libsodium`
 pub struct SodiumStreamCipher {
@@ -56,7 +53,7 @@ impl SodiumStreamCipher {
         self.counter % SODIUM_BLOCK_SIZE
     }
 
-    fn process(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()> {
+    fn process(&mut self, data: &[u8], out: &mut dyn BufMut) -> CipherResult<()> {
         let padding = self.padding_len();
 
         let mut plain_text = vec![0u8; data.len() + padding];
@@ -88,12 +85,12 @@ fn crypto_stream_xor_ic<B: BufMut>(
     data: &[u8],
     out: &mut B,
 ) -> CipherResult<()> {
-    assert!(data.len() <= unsafe { out.bytes_mut().len() });
+    assert!(data.len() <= out.bytes_mut().len());
 
     let ret = unsafe {
         match t {
             CipherType::ChaCha20 => crypto_stream_chacha20_xor_ic(
-                out.bytes_mut().as_mut_ptr(),
+                out.bytes_mut().as_mut_ptr() as *mut u8,
                 data.as_ptr(),
                 data.len() as c_ulonglong,
                 iv.as_ptr(),
@@ -101,7 +98,7 @@ fn crypto_stream_xor_ic<B: BufMut>(
                 key.as_ptr(),
             ),
             CipherType::ChaCha20Ietf => crypto_stream_chacha20_ietf_xor_ic(
-                out.bytes_mut().as_mut_ptr(),
+                out.bytes_mut().as_mut_ptr() as *mut u8,
                 data.as_ptr(),
                 data.len() as c_ulonglong,
                 iv.as_ptr(),
@@ -109,7 +106,7 @@ fn crypto_stream_xor_ic<B: BufMut>(
                 key.as_ptr(),
             ),
             CipherType::Salsa20 => crypto_stream_salsa20_xor_ic(
-                out.bytes_mut().as_mut_ptr(),
+                out.bytes_mut().as_mut_ptr() as *mut u8,
                 data.as_ptr(),
                 data.len() as c_ulonglong,
                 iv.as_ptr(),
@@ -117,7 +114,7 @@ fn crypto_stream_xor_ic<B: BufMut>(
                 key.as_ptr(),
             ),
             CipherType::XSalsa20 => crypto_stream_xsalsa20_xor_ic(
-                out.bytes_mut().as_mut_ptr(),
+                out.bytes_mut().as_mut_ptr() as *mut u8,
                 data.as_ptr(),
                 data.len() as c_ulonglong,
                 iv.as_ptr(),
@@ -140,11 +137,11 @@ fn crypto_stream_xor_ic<B: BufMut>(
 }
 
 impl StreamCipher for SodiumStreamCipher {
-    fn update(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()> {
+    fn update(&mut self, data: &[u8], out: &mut dyn BufMut) -> CipherResult<()> {
         self.process(data, out)
     }
 
-    fn finalize(&mut self, _: &mut BufMut) -> CipherResult<()> {
+    fn finalize(&mut self, _: &mut dyn BufMut) -> CipherResult<()> {
         Ok(())
     }
 

src/crypto/stream.rs

@@ -18,8 +18,8 @@ use bytes::BufMut;
 /// The `update` method could be called multiple times, and the `finalize` method will
 /// encrypt the last block
 pub trait StreamCipher {
-    fn update(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()>;
-    fn finalize(&mut self, out: &mut BufMut) -> CipherResult<()>;
+    fn update(&mut self, data: &[u8], out: &mut dyn BufMut) -> CipherResult<()>;
+    fn finalize(&mut self, out: &mut dyn BufMut) -> CipherResult<()>;
     fn buffer_size(&self, data: &[u8]) -> usize;
 }
 

src/crypto/table.rs

@@ -56,7 +56,7 @@ impl TableCipher {
         }
     }
 
-    fn process(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()> {
+    fn process(&mut self, data: &[u8], out: &mut dyn BufMut) -> CipherResult<()> {
         let mut buf = BytesMut::with_capacity(self.buffer_size(data));
         unsafe {
             buf.set_len(self.buffer_size(data)); // Set length
@@ -70,11 +70,11 @@ impl TableCipher {
 }
 
 impl StreamCipher for TableCipher {
-    fn update(&mut self, data: &[u8], out: &mut BufMut) -> CipherResult<()> {
+    fn update(&mut self, data: &[u8], out: &mut dyn BufMut) -> CipherResult<()> {
         self.process(data, out)
     }
 
-    fn finalize(&mut self, _: &mut BufMut) -> CipherResult<()> {
+    fn finalize(&mut self, _: &mut dyn BufMut) -> CipherResult<()> {
         Ok(())
     }
 

src/lib.rs

@@ -75,9 +75,11 @@ pub const VERSION: &str = env!("CARGO_PKG_VERSION");
 
 pub use self::{
     config::{ClientConfig, Config, ConfigType, Mode, ServerAddr, ServerConfig},
-    relay::{dns::run as run_dns, local::run as run_local, server::run as run_server, tcprelay::client::Socks5Client},
+    relay::{local::run as run_local, server::run as run_server, tcprelay::client::Socks5Client},
 };
 
+// relay::{dns::run as run_dns},
+
 pub mod config;
 mod context;
 pub mod crypto;

src/plugin/mod.rs

@@ -13,13 +13,13 @@
 //! ```
 
 use crate::config::{Config, ServerAddr};
-use futures::{stream::futures_unordered, Future, Stream};
+use futures::stream::{FuturesUnordered, StreamExt};
 use log::{error, info};
 use std::{
     io,
     net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener},
 };
-use tokio_process::{Child, CommandExt};
+use tokio::process::Child;
 
 mod obfs_proxy;
 mod ss_plugin;
@@ -38,71 +38,78 @@ pub enum PluginMode {
     Client,
 }
 
-/// Launch plugins in config. Returns a future that completes when any plugin terminates
-/// or there were an error in watching the subprocess. Returns `None` if no plugins
-/// were launched.
-pub fn launch_plugins(
-    config: &mut Config,
-    mode: PluginMode,
-) -> io::Result<Option<impl Future<Item = (), Error = io::Error>>> {
-    let mut plugins = Vec::new();
+/// Started plugins' subprocesses carrier
+pub struct Plugins {
+    plugins: FuturesUnordered<Child>,
+}
 
-    for svr in &mut config.server {
-        let mut svr_addr_opt = None;
+impl Plugins {
+    /// Launch plugins in configuration.
+    ///
+    /// Will modify servers' listen addresses to plugins' listen addresses.
+    pub fn launch_plugins(config: &mut Config, mode: PluginMode) -> io::Result<Plugins> {
+        let plugins = FuturesUnordered::new();
 
-        if let Some(c) = svr.plugin() {
-            let loop_ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
-            let local_addr = SocketAddr::new(loop_ip, get_local_port()?);
+        for svr in &mut config.server {
+            let mut svr_addr_opt = None;
 
-            let svr_addr = match start_plugin(c, svr.addr(), &local_addr, mode) {
-                Err(err) => {
-                    error!("Failed to start plugin \"{}\", err: {}", c.plugin, err);
-                    return Err(err);
+            if let Some(c) = svr.plugin() {
+                let loop_ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
+                let local_addr = SocketAddr::new(loop_ip, get_local_port()?);
+
+                let svr_addr = match start_plugin(c, svr.addr(), &local_addr, mode) {
+                    Err(err) => {
+                        error!("Failed to start plugin \"{}\", err: {}", c.plugin, err);
+                        return Err(err);
+                    }
+                    Ok(process) => {
+                        let svr_addr = ServerAddr::SocketAddr(local_addr);
+                        plugins.push(process);
+
+                        // Replace addr with plugin
+                        svr_addr
+                    }
+                };
+
+                match mode {
+                    PluginMode::Client => info!("Started plugin \"{}\" on {} <-> {}", c.plugin, local_addr, svr.addr()),
+                    PluginMode::Server => info!("Started plugin \"{}\" on {} <-> {}", c.plugin, svr.addr(), local_addr),
                 }
-                Ok(process) => {
-                    let svr_addr = ServerAddr::SocketAddr(local_addr);
-                    plugins.push(process);
 
-                    // Replace addr with plugin
-                    svr_addr
-                }
-            };
-
-            match mode {
-                PluginMode::Client => info!("Started plugin \"{}\" on {} <-> {}", c.plugin, local_addr, svr.addr()),
-                PluginMode::Server => info!("Started plugin \"{}\" on {} <-> {}", c.plugin, svr.addr(), local_addr),
+                svr_addr_opt = Some(svr_addr); // Fuck borrow checker
             }
 
-            svr_addr_opt = Some(svr_addr); // Fuck borrow checker
+            if let Some(svr_addr) = svr_addr_opt {
+                svr.set_plugin_addr(svr_addr);
+            }
         }
 
-        if let Some(svr_addr) = svr_addr_opt {
-            svr.set_plugin_addr(svr_addr);
+        if plugins.is_empty() {
+            panic!("Didn't find any plugins to start");
         }
+
+        Ok(Plugins { plugins })
     }
 
-    if plugins.is_empty() {
-        Ok(None)
-    } else {
+    /// Returns a future that completes when any plugin terminates or there were an error in watching the subprocess.
+    pub async fn into_future(self) -> io::Result<()> {
         // Turn the vector of `Child` futures into a single future that
         // completes with an error if any of them exits or waiting for it
         // fails. When this future completes, the remaining `Child`ren will be
         // dropped and as a result the rest of the plugins will be killed
         // automatically.
-        let plugins_future =
-            futures_unordered(plugins)
-                .into_future()
-                .then(|first_plugin_result| match first_plugin_result {
-                    Ok((first_plugin_exit_status, _)) => {
-                        let msg = format!("Plugin exited unexpectedly with {}", first_plugin_exit_status.unwrap());
-                        Err(io::Error::new(io::ErrorKind::Other, msg))
-                    }
-                    Err((first_plugin_error, _)) => {
-                        error!("Error while waiting for plugin subprocess: {}", first_plugin_error);
-                        Err(first_plugin_error)
-                    }
-                });
-        Ok(Some(plugins_future))
+
+        match self.plugins.into_future().await {
+            (Some(Ok(first_plugin_exit_status)), _) => {
+                let msg = format!("Plugin exited unexpectedly with {}", first_plugin_exit_status);
+                Err(io::Error::new(io::ErrorKind::Other, msg))
+            }
+            (Some(Err(first_plugin_error)), _) => {
+                error!("Error while waiting for plugin subprocess: {}", first_plugin_error);
+                Err(first_plugin_error)
+            }
+            _ => unreachable!(),
+        }
     }
 }
 
@@ -112,7 +119,7 @@ fn start_plugin(plugin: &PluginConfig, remote: &ServerAddr, local: &SocketAddr,
     } else {
         ss_plugin::plugin_cmd(plugin, remote, local, mode)
     };
-    cmd.spawn_async()
+    cmd.spawn()
 }
 
 fn get_local_port() -> io::Result<u16> {

src/plugin/obfs_proxy.rs

@@ -1,9 +1,7 @@
 use super::{PluginConfig, PluginMode};
 use crate::config::ServerAddr;
-use std::{
-    net::SocketAddr,
-    process::{Command, Stdio},
-};
+use std::{net::SocketAddr, process::Stdio};
+use tokio::process::Command;
 
 /// For obfsproxy, we use standalone mode for now.
 /// Managed mode needs to use SOCKS5 proxy as forwarder, which is not supported

src/plugin/ss_plugin.rs

@@ -1,10 +1,8 @@
 use super::{PluginConfig, PluginMode};
 use crate::config::ServerAddr;
 use log::trace;
-use std::{
-    net::SocketAddr,
-    process::{Command, Stdio},
-};
+use std::{net::SocketAddr, process::Stdio};
+use tokio::process::Command;
 
 pub fn plugin_cmd(plugin: &PluginConfig, remote: &ServerAddr, local: &SocketAddr, _mode: PluginMode) -> Command {
     trace!("Start plugin \"{:?}\" remote: {}, local: {}", plugin, remote, local);

src/relay/dns.rs

@@ -1,16 +0,0 @@
-//! DNS relay
-
-use std::io;
-
-use futures::{self, Future};
-
-use crate::{
-    config::Config,
-    context::{Context, SharedContext},
-    relay::udprelay::dns::run as run_udp,
-};
-
-/// DNS Relay server running under local environment.
-pub fn run(config: Config) -> impl Future<Item = (), Error = io::Error> + Send {
-    futures::lazy(move || run_udp(SharedContext::new(Context::new_dns(config))))
-}

src/relay/dns_resolver/mod.rs

@@ -0,0 +1,6 @@
+//! Asynchronous DNS resolver
+
+mod trust_dns_resolver;
+
+pub use self::trust_dns_resolver::create_resolver;
+pub use self::trust_dns_resolver::resolve;

src/relay/dns_resolver/tokio_dns_resolver.rs

@@ -0,0 +1,77 @@
+use std::io::{self, ErrorKind};
+use std::net::SocketAddr;
+use std::net::ToSocketAddrs;
+
+use log::{debug, error};
+
+use crate::context::SharedContext;
+
+/// Resolve address to IP
+#[cfg(not(feature = "single-threaded"))]
+pub async fn resolve(
+    context: SharedContext,
+    addr: &str,
+    port: u16,
+    check_forbidden: bool,
+) -> io::Result<Vec<SocketAddr>> {
+    use tokio::task::spawn_blocking;
+
+    let addr = addr.to_owned();
+    spawn_blocking(move || {
+        let mut vaddr = Vec::new();
+        for addr in (addr.as_str(), port).to_socket_addrs()? {
+            let ip = addr.ip();
+
+            if check_forbidden {
+                let forbidden_ip = &context.config().forbidden_ip;
+                if forbidden_ip.contains(&ip) {
+                    debug!("Resolved {} => {}, which is skipped by forbidden_ip", addr, ip);
+                    continue;
+                }
+            }
+            vaddr.push(addr);
+        }
+
+        if vaddr.is_empty() {
+            error!("Failed to resolve {}:{}, all IPs are filtered", addr, port);
+            let err = io::Error::new(ErrorKind::Other, "resolved to empty address, all IPs are filtered");
+            Err(err)
+        } else {
+            debug!("Resolved {}:{} => {:?}", addr, port, vaddr);
+            Ok(vaddr)
+        }
+    })
+    .await?
+}
+
+/// Resolve address to IP
+#[cfg(feature = "single-threaded")]
+pub async fn resolve(
+    context: SharedContext,
+    addr: &str,
+    port: u16,
+    check_forbidden: bool,
+) -> io::Result<Vec<SocketAddr>> {
+    let mut vaddr = Vec::new();
+    for addr in (addr, port).to_socket_addrs()? {
+        let ip = addr.ip();
+
+        if check_forbidden {
+            let forbidden_ip = &context.config().forbidden_ip;
+            if forbidden_ip.contains(&ip) {
+                debug!("Resolved {} => {}, which is skipped by forbidden_ip", addr, ip);
+                continue;
+            }
+        }
+        vaddr.push(addr);
+    }
+
+    if vaddr.is_empty() {
+        error!("Failed to resolve {}:{}, all IPs are filtered", addr, port);
+        let err = io::Error::new(ErrorKind::Other, "resolved to empty address, all IPs are filtered");
+        Err(err)
+    } else {
+        debug!("Resolved {}:{} => {:?}", addr, port, vaddr);
+        Ok(vaddr)
+    }
+}

src/relay/dns_resolver/trust_dns_resolver.rs

@@ -5,7 +5,7 @@ use std::{
     net::SocketAddr,
 };
 
-use futures::Future;
+use log::{debug, error, trace};
 use tokio;
 use trust_dns_resolver::{config::ResolverConfig, AsyncResolver};
 
@@ -18,11 +18,22 @@ pub fn create_resolver(dns: Option<ResolverConfig>) -> AsyncResolver {
         {
             if let Some(conf) = dns {
                 use trust_dns_resolver::config::ResolverOpts;
+                trace!(
+                    "Initializing DNS resolver with config {:?} opts {:?}",
+                    conf,
+                    ResolverOpts::default()
+                );
                 AsyncResolver::new(conf, ResolverOpts::default())
             } else {
                 use trust_dns_resolver::system_conf::read_system_conf;
                 // use the system resolver configuration
                 let (config, opts) = read_system_conf().expect("Failed to read global dns sysconf");
+                trace!(
+                    "Initializing DNS resolver with system-config {:?} opts {:?}",
+                    config,
+                    opts
+                );
+
                 AsyncResolver::new(config, opts)
             }
         }
@@ -34,9 +45,19 @@ pub fn create_resolver(dns: Option<ResolverConfig>) -> AsyncResolver {
             use trust_dns_resolver::config::{ResolverConfig, ResolverOpts};
 
             if let Some(conf) = dns {
+                trace!(
+                    "Initializing DNS resolver with config {:?} opts {:?}",
+                    conf,
+                    ResolverOpts::default()
+                );
                 AsyncResolver::new(conf, ResolverOpts::default())
             } else {
                 // Get a new resolver with the google nameservers as the upstream recursive resolvers
+                trace!(
+                    "Initializing DNS resolver with google-config {:?} opts {:?}",
+                    ResolverConfig::google(),
+                    ResolverOpts::default()
+                );
                 AsyncResolver::new(ResolverConfig::google(), ResolverOpts::default())
             }
         }
@@ -48,17 +69,15 @@ pub fn create_resolver(dns: Option<ResolverConfig>) -> AsyncResolver {
     resolver
 }
 
-fn inner_resolve(
+async fn inner_resolve(
     context: SharedContext,
     addr: &str,
     port: u16,
     check_forbidden: bool,
-) -> impl Future<Item = Vec<SocketAddr>, Error = io::Error> + Send {
-    // let owned_addr = addr.to_owned();
-    let cloned_context = context.clone();
-    context.dns_resolver().lookup_ip(addr).then(move |r| match r {
+) -> io::Result<Vec<SocketAddr>> {
+    match context.dns_resolver().lookup_ip(addr).await {
         Err(err) => {
-            // error!("Failed to resolve {}, err: {}", owned_addr, err);
+            error!("Failed to resolve {}:{}, err: {}", addr, port, err);
             Err(io::Error::new(
                 io::ErrorKind::Other,
                 format!("dns resolve error: {}", err),
@@ -68,9 +87,9 @@ fn inner_resolve(
             let mut vaddr = Vec::new();
             for ip in lookup_result.iter() {
                 if check_forbidden {
-                    let forbidden_ip = &cloned_context.config().forbidden_ip;
+                    let forbidden_ip = &context.config().forbidden_ip;
                     if forbidden_ip.contains(&ip) {
-                        // debug!("Resolved {} => {}, which is skipped by forbidden_ip", owned_addr, ip);
+                        debug!("Resolved {} => {}, which is skipped by forbidden_ip", addr, ip);
                         continue;
                     }
                 }
@@ -78,26 +97,23 @@ fn inner_resolve(
             }
 
             if vaddr.is_empty() {
-                let err = io::Error::new(
-                    ErrorKind::Other,
-                    // format!("resolved {} to empty address, all IPs are filtered", owned_addr),
-                    "resolved to empty address, all IPs are filtered",
-                );
+                error!("Failed to resolve {}:{}, all IPs are filtered", addr, port);
+                let err = io::Error::new(ErrorKind::Other, "resolved to empty address, all IPs are filtered");
                 Err(err)
             } else {
-                // debug!("Resolved {} => {:?}", owned_addr, vaddr);
+                debug!("Resolved {}:{} => {:?}", addr, port, vaddr);
                 Ok(vaddr)
             }
         }
-    })
+    }
 }
 
 /// Resolve address to IP
-pub fn resolve(
+pub async fn resolve(
     context: SharedContext,
     addr: &str,
     port: u16,
     check_forbidden: bool,
-) -> impl Future<Item = Vec<SocketAddr>, Error = io::Error> + Send {
-    inner_resolve(context, addr, port, check_forbidden)
+) -> io::Result<Vec<SocketAddr>> {
+    inner_resolve(context, addr, port, check_forbidden).await
 }

src/relay/loadbalancing/server/ping.rs

@@ -5,22 +5,19 @@ use std::{
         atomic::{AtomicU64, AtomicUsize, Ordering},
         Arc, Mutex,
     },
-    time::{Duration, Instant},
 };
 
 use crate::{
     config::ServerConfig,
     context::SharedContext,
-    relay::tcprelay::EncryptedWrite,
     relay::{loadbalancing::server::LoadBalancer, socks5::Address, tcprelay::client::ServerClient},
 };
 
-use futures::{Future, Stream};
 use log::{debug, error, info};
 use tokio::{
     self,
-    io::read_exact,
-    timer::{Interval, Timeout},
+    io::{AsyncReadExt, AsyncWriteExt},
+    time::{self, Duration, Instant},
 };
 
 struct Server {
@@ -137,26 +134,22 @@ impl Inner {
 
             tokio::spawn(
                 // Check every DEFAULT_CHECK_INTERVAL_SEC seconds
-                Interval::new(
-                    Instant::now() + Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
-                    Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
-                )
-                .for_each(move |_| {
-                    let sc = sc.clone();
-                    let lat = latency.clone();
+                async move {
+                    // Wait until the server is ready (plugins are started)
+                    time::delay_for(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC)).await;
 
-                    Inner::check_delay(sc.clone(), context.clone()).then(move |res| {
-                        let score = match res {
-                            Ok(d) => lat.push(d),
-                            Err(..) => lat.push(DEFAULT_CHECK_TIMEOUT_SEC * 2 * 1000), // Penalty
+                    let mut interval = time::interval(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC));
+
+                    while context.server_running() {
+                        interval.tick().await;
+                        let score = match Inner::check_delay(sc.clone(), context.clone()).await {
+                            Ok(d) => latency.push(d),
+                            Err(..) => latency.push(DEFAULT_CHECK_TIMEOUT_SEC * 2 * 1000), // Penalty
                         };
                         debug!("updated remote server {} (score: {})", sc.config.addr(), score);
                         sc.score.store(score, Ordering::Release);
-
-                        Ok(())
-                    })
-                })
-                .map_err(|_| ()),
+                    }
+                },
             );
         }
 
@@ -166,52 +159,54 @@ impl Inner {
         }
     }
 
-    fn check_delay(sc: Arc<Server>, context: SharedContext) -> impl Future<Item = u64, Error = io::Error> {
-        static GET_BODY: &[u8] = b"GET / HTTP/1.1\r\nHost: www.example.com\r\nConnection: close\r\nAccept: */*\r\n\r\n";
+    async fn check_request(sc: Arc<ServerConfig>, context: SharedContext) -> io::Result<()> {
+        static GET_BODY: &[u8] =
+            b"GET /generate_204 HTTP/1.1\r\nHost: dl.google.com\r\nConnection: close\r\nAccept: */*\r\n\r\n";
 
-        let addr = Address::DomainNameAddress("www.example.com".to_owned(), 80);
+        let addr = Address::DomainNameAddress("dl.google.com".to_owned(), 80);
 
+        let ServerClient { mut stream } = ServerClient::connect(context, &addr, sc).await?;
+        stream.write_all(GET_BODY).await?;
+        stream.flush().await?;
+        let mut buf = [0u8; 1];
+        stream.read_exact(&mut buf).await?;
+
+        Ok(())
+    }
+
+    async fn check_delay(sc: Arc<Server>, context: SharedContext) -> io::Result<u64> {
         let start = Instant::now();
 
         // Send HTTP GET and read the first byte
-        let fut = ServerClient::connect(context, addr, sc.config.clone()).and_then(|ServerClient { r, w }| {
-            w.write_all(GET_BODY)
-                .and_then(|(mut w, _)| w.flush())
-                .and_then(move |_| read_exact(r, [0u8, 1]))
-        });
+        let timeout = Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC);
+        let res = time::timeout(timeout, Inner::check_request(sc.config.clone(), context.clone())).await;
 
-        Timeout::new(fut, Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC)).then(move |res| {
-            let elapsed = Instant::now() - start;
-            let elapsed = elapsed.as_secs() * 1000 + u64::from(elapsed.subsec_millis()); // Converted to ms
-            match res {
-                Ok(..) => {
-                    // Got the result ... record its time
-                    debug!("checked remote server {} latency with {} ms", sc.config.addr(), elapsed);
-                    Ok(elapsed)
-                }
-                Err(err) => {
-                    match err.into_inner() {
-                        Some(err) => {
-                            error!("failed to check server {}, error: {}", sc.config.addr(), err);
-
-                            // NOTE: connection / handshake error, server is down
-                            Err(err)
-                        }
-                        None => {
-                            // Timeout
-                            debug!(
-                                "checked remote server {} latency timeout, elapsed {} ms",
-                                sc.config.addr(),
-                                elapsed
-                            );
-
-                            // NOTE: timeout is still available, but server is too slow
-                            Ok(elapsed)
-                        }
-                    }
-                }
+        let elapsed = Instant::now() - start;
+        let elapsed = elapsed.as_secs() * 1000 + u64::from(elapsed.subsec_millis()); // Converted to ms
+        match res {
+            Ok(Ok(..)) => {
+                // Got the result ... record its time
+                debug!("checked remote server {} latency with {} ms", sc.config.addr(), elapsed);
+                Ok(elapsed)
             }
-        })
+            Ok(Err(err)) => {
+                error!("failed to check server {}, error: {}", sc.config.addr(), err);
+
+                // NOTE: connection / handshake error, server is down
+                Err(err)
+            }
+            Err(..) => {
+                // Timeout
+                debug!(
+                    "checked remote server {} latency timeout, elapsed {} ms",
+                    sc.config.addr(),
+                    elapsed
+                );
+
+                // NOTE: timeout is still available, but server is too slow
+                Ok(elapsed)
+            }
+        }
     }
 
     fn best_idx(&self) -> usize {
@@ -234,54 +229,54 @@ pub struct PingBalancer {
 
 impl PingBalancer {
     pub fn new(context: SharedContext) -> PingBalancer {
-        let inner = Arc::new(Inner::new(context));
+        let inner = Arc::new(Inner::new(context.clone()));
 
-        let inner_cloned = inner.clone();
-        tokio::spawn(
-            Interval::new(
-                Instant::now() + Duration::from_secs(2),
-                Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC),
-            )
-            .for_each(move |_| {
-                let inner = inner_cloned.clone();
+        if inner.servers.len() > 1 {
+            let cloned_inner = inner.clone();
+            tokio::spawn(async move {
+                let inner = cloned_inner;
 
-                if inner.servers.is_empty() {
-                    panic!("No server");
-                }
+                // Wait until the server is ready (plugins are started)
+                time::delay_for(Duration::from_secs(2)).await;
 
-                // Choose the best one
-                let mut svr_idx = 0;
+                let mut interval = time::interval(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC));
 
-                for (idx, svr) in inner.servers.iter().enumerate() {
-                    let choosen_svr = &inner.servers[svr_idx];
-                    if svr.score() < choosen_svr.score() {
-                        svr_idx = idx;
+                while context.server_running() {
+                    interval.tick().await;
+
+                    assert!(!inner.servers.is_empty());
+
+                    // Choose the best one
+                    let mut svr_idx = 0;
+
+                    for (idx, svr) in inner.servers.iter().enumerate() {
+                        let choosen_svr = &inner.servers[svr_idx];
+                        if svr.score() < choosen_svr.score() {
+                            svr_idx = idx;
+                        }
                     }
-                }
 
-                let choosen_svr = &inner.servers[svr_idx];
-                debug!(
-                    "chosen the best server {} (score: {})",
-                    choosen_svr.config.addr(),
-                    choosen_svr.score()
-                );
-
-                if inner.best_idx() != svr_idx {
-                    info!(
-                        "switched server from {} (score: {}) to {} (score: {})",
-                        inner.best_server().config.addr(),
-                        inner.best_server().score(),
+                    let choosen_svr = &inner.servers[svr_idx];
+                    debug!(
+                        "chosen the best server {} (score: {})",
                         choosen_svr.config.addr(),
-                        choosen_svr.score(),
+                        choosen_svr.score()
                     );
+
+                    if inner.best_idx() != svr_idx {
+                        info!(
+                            "switched server from {} (score: {}) to {} (score: {})",
+                            inner.best_server().config.addr(),
+                            inner.best_server().score(),
+                            choosen_svr.config.addr(),
+                            choosen_svr.score(),
+                        );
+                    }
+
+                    inner.set_best_idx(svr_idx);
                 }
-
-                inner.set_best_idx(svr_idx);
-
-                Ok(())
-            })
-            .map_err(|_| ()),
-        );
+            });
+        }
 
         PingBalancer { inner }
     }

src/relay/local.rs

@@ -2,67 +2,49 @@
 
 use std::io;
 
-use futures::{stream::futures_unordered, Future, Stream};
+use futures::{future::select_all, FutureExt};
+
+use log::error;
 
 use crate::{
     config::Config,
-    context::{Context, SharedContext},
-    plugin::{launch_plugins, PluginMode},
-    relay::{boxed_future, tcprelay::local::run as run_tcp, udprelay::local::run as run_udp},
+    context::{Context, SharedServerState},
+    plugin::{PluginMode, Plugins},
+    relay::{tcprelay::local::run as run_tcp, udprelay::local::run as run_udp},
 };
 
 /// Relay server running under local environment.
-///
-/// ```no_run
-/// use shadowsocks::{
-///     config::{Config, ConfigType, ServerConfig},
-///     crypto::CipherType,
-///     relay::local::run,
-/// };
-///
-/// use tokio::prelude::*;
-///
-/// 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(),
-///     "server-password".to_string(),
-///     CipherType::Aes256Cfb,
-/// )];
-/// let fut = run(config);
-/// tokio::run(fut.map_err(|err| panic!("Server run failed with error {}", err)));
-/// ```
-pub fn run(config: Config) -> impl Future<Item = (), Error = io::Error> + Send {
-    futures::lazy(move || {
-        let mut context = Context::new(config);
+pub async fn run(mut config: Config) -> io::Result<()> {
+    // Create a context containing a DNS resolver and server running state flag.
+    let state = SharedServerState::new(&config);
 
-        let mut vf = Vec::new();
+    let mut vf = Vec::new();
 
-        if context.config().mode.enable_udp() {
-            // Clone config here, because the config for TCP relay will be modified
-            // after plugins started
-            let udp_context = SharedContext::new(context.clone());
+    if config.mode.enable_udp() {
+        // Clone config here, because the config for TCP relay will be modified
+        // after plugins started.
+        // But DNS resolver and running state flag is still shared.
+        let udp_context = Context::new_shared(config.clone(), state.clone());
 
-            // Run UDP relay before starting plugins
-            // Because plugins doesn't support UDP relay
-            let udp_fut = run_udp(udp_context);
-            vf.push(boxed_future(udp_fut));
-        }
+        // Run UDP relay before starting plugins
+        // Because plugins doesn't support UDP relay
+        let udp_fut = run_udp(udp_context);
+        vf.push(udp_fut.boxed());
+    }
 
-        if let Some(plugins) =
-            launch_plugins(context.config_mut(), PluginMode::Client).expect("Failed to launch plugins")
-        {
-            vf.push(boxed_future(plugins));
-        }
+    if config.has_server_plugins() {
+        let plugins = Plugins::launch_plugins(&mut config, PluginMode::Client)?;
+        vf.push(plugins.into_future().boxed());
+    }
 
-        let tcp_fut = run_tcp(SharedContext::new(context));
+    let tcp_fut = run_tcp(Context::new_shared(config, state.clone()));
+    vf.push(tcp_fut.boxed());
 
-        vf.push(boxed_future(tcp_fut));
-        futures_unordered(vf).into_future().then(|res| -> io::Result<()> {
-            match res {
-                Ok(..) => Ok(()),
-                Err((err, ..)) => Err(err),
-            }
-        })
-    })
+    let (res, ..) = select_all(vf.into_iter()).await;
+    error!("One of TCP servers exited unexpectly, result: {:?}", res);
+
+    // Tells all detached tasks to exit
+    state.server_stopped();
+
+    Err(io::Error::new(io::ErrorKind::Other, "server exited unexpectly"))
 }

src/relay/mod.rs

@@ -1,20 +1,11 @@
 //! Relay server in local and server side implementations.
 
-use futures::Future;
-
-pub mod dns;
+#[cfg(feature = "trust-dns")]
 pub(crate) mod dns_resolver;
-mod loadbalancing;
+pub(crate) mod loadbalancing;
 pub mod local;
 pub mod server;
 pub mod socks5;
 pub mod tcprelay;
 pub mod udprelay;
-mod utils;
-
-pub fn boxed_future<T, E, F>(f: F) -> Box<Future<Item = T, Error = E> + Send + 'static>
-where
-    F: Future<Item = T, Error = E> + Send + 'static,
-{
-    Box::new(f)
-}
+pub(crate) mod utils;

src/relay/server.rs

@@ -2,69 +2,50 @@
 
 use std::io;
 
-use futures::{stream::futures_unordered, Future, Stream};
+use futures::future::{select_all, FutureExt};
+
+use log::error;
 
 use crate::{
     config::Config,
-    context::{Context, SharedContext},
-    plugin::{launch_plugins, PluginMode},
-    relay::{boxed_future, tcprelay::server::run as run_tcp, udprelay::server::run as run_udp},
+    context::{Context, SharedServerState},
+    plugin::{PluginMode, Plugins},
+    relay::{tcprelay::server::run as run_tcp, udprelay::server::run as run_udp},
 };
 
 /// Relay server running on server side.
-///
-/// ```no_run
-/// use shadowsocks::{
-///     config::{Config, ConfigType, ServerConfig},
-///     crypto::CipherType,
-///     relay::server::run,
-/// };
-///
-/// use tokio::prelude::*;
-///
-/// let mut config = Config::new(ConfigType::Server);
-/// config.server = vec![ServerConfig::basic(
-///     "127.0.0.1:8388".parse().unwrap(),
-///     "server-password".to_string(),
-///     CipherType::Aes256Cfb,
-/// )];
-///
-/// let fut = run(config);
-/// tokio::run(fut.map_err(|err| panic!("Server run failed with error {}", err)));
-/// ```
-pub fn run(config: Config) -> impl Future<Item = (), Error = io::Error> + Send {
-    futures::lazy(move || {
-        let mut context = Context::new(config);
+pub async fn run(mut config: Config) -> io::Result<()> {
+    // Create a context containing a DNS resolver and server running state flag.
+    let state = SharedServerState::new(&config);
 
-        let mut vf = Vec::new();
+    let mut vf = Vec::new();
 
-        if context.config().mode.enable_udp() {
-            // Clone config here, because the config for TCP relay will be modified
-            // after plugins started
-            let udp_context = SharedContext::new(context.clone());
+    if config.mode.enable_udp() {
+        // Clone config here, because the config for TCP relay will be modified
+        // after plugins started
+        let udp_context = Context::new_shared(config.clone(), state.clone());
 
-            // Run UDP relay before starting plugins
-            // Because plugins doesn't support UDP relay
-            let udp_fut = run_udp(udp_context);
-            vf.push(boxed_future(udp_fut));
+        // Run UDP relay before starting plugins
+        // Because plugins doesn't support UDP relay
+        let udp_fut = run_udp(udp_context);
+        vf.push(udp_fut.boxed());
+    }
+
+    if config.mode.enable_tcp() {
+        if config.has_server_plugins() {
+            let plugins = Plugins::launch_plugins(&mut config, PluginMode::Client)?;
+            vf.push(plugins.into_future().boxed());
         }
 
-        if context.config().mode.enable_tcp() {
-            if let Some(plugins) =
-                launch_plugins(context.config_mut(), PluginMode::Server).expect("Failed to launch plugins")
-            {
-                vf.push(boxed_future(plugins));
-            }
+        let tcp_fut = run_tcp(Context::new_shared(config, state.clone()));
+        vf.push(tcp_fut.boxed());
+    }
 
-            let tcp_fut = run_tcp(SharedContext::new(context));
-            vf.push(boxed_future(tcp_fut));
-        }
+    let (res, ..) = select_all(vf.into_iter()).await;
+    error!("One of TCP servers exited unexpectly, result: {:?}", res);
 
-        futures_unordered(vf).into_future().then(|res| -> io::Result<()> {
-            match res {
-                Ok(..) => Ok(()),
-                Err((err, ..)) => Err(err),
-            }
-        })
-    })
+    // Tells all detached tasks to exit
+    state.server_stopped();
+
+    Err(io::Error::new(io::ErrorKind::Other, "server exited unexpectly"))
 }

src/relay/socks5.rs

@@ -2,27 +2,24 @@
 //!
 //! Implements [SOCKS Protocol Version 5](https://www.ietf.org/rfc/rfc1928.txt) proxy protocol
 
-use std::{
-    convert::From,
-    error,
-    fmt::{self, Debug, Formatter},
-    io::{self, Cursor, Read, Write},
-    net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6, ToSocketAddrs},
-    u8, vec,
-};
+use std::convert::From;
+use std::error;
+use std::fmt::{self, Debug, Formatter};
+use std::io::{self, Cursor};
+use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6, ToSocketAddrs};
+use std::str::FromStr;
+use std::u8;
+use std::vec;
 
-use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
-use bytes::{BufMut, Bytes, BytesMut, IntoBuf};
-use futures::{try_ready, Async, Future, Poll};
+use bytes::buf::BufExt;
+use bytes::{Buf, BufMut, BytesMut};
 use log::error;
-use tokio_io::{io::read_exact, try_nb, AsyncRead, AsyncWrite};
+use tokio::prelude::*;
 
 pub use self::consts::{
     SOCKS5_AUTH_METHOD_GSSAPI, SOCKS5_AUTH_METHOD_NONE, SOCKS5_AUTH_METHOD_NOT_ACCEPTABLE, SOCKS5_AUTH_METHOD_PASSWORD,
 };
 
-use super::utils::{write_bytes, WriteBytes};
-
 #[rustfmt::skip]
 mod consts {
     pub const SOCKS5_VERSION:                          u8 = 0x05;
@@ -192,7 +189,7 @@ impl error::Error for Error {
         &self.message[..]
     }
 
-    fn cause(&self) -> Option<&error::Error> {
+    fn cause(&self) -> Option<&dyn error::Error> {
         None
     }
 }
@@ -219,17 +216,84 @@ pub enum Address {
 }
 
 impl Address {
-    #[inline]
-    pub fn read_from<R: AsyncRead>(stream: R) -> ReadAddress<R> {
-        ReadAddress::new(stream)
+    pub async fn read_from<R>(stream: &mut R) -> Result<Address, Error>
+    where
+        R: AsyncRead + Unpin,
+    {
+        let mut addr_type_buf = [0u8; 1];
+        let _ = stream.read_exact(&mut addr_type_buf).await?;
+
+        let addr_type = addr_type_buf[0];
+        match addr_type {
+            consts::SOCKS5_ADDR_TYPE_IPV4 => {
+                let mut buf = BytesMut::with_capacity(6);
+                buf.resize(6, 0);
+                let _ = stream.read_exact(&mut buf).await?;
+
+                let mut cursor = buf.to_bytes();
+                let v4addr = Ipv4Addr::new(cursor.get_u8(), cursor.get_u8(), cursor.get_u8(), cursor.get_u8());
+                let port = cursor.get_u16();
+                Ok(Address::SocketAddress(SocketAddr::V4(SocketAddrV4::new(v4addr, port))))
+            }
+            consts::SOCKS5_ADDR_TYPE_IPV6 => {
+                let mut buf = [0u8; 18];
+                let _ = stream.read_exact(&mut buf).await?;
+
+                let mut cursor = Cursor::new(&buf);
+                let v6addr = Ipv6Addr::new(
+                    cursor.get_u16(),
+                    cursor.get_u16(),
+                    cursor.get_u16(),
+                    cursor.get_u16(),
+                    cursor.get_u16(),
+                    cursor.get_u16(),
+                    cursor.get_u16(),
+                    cursor.get_u16(),
+                );
+                let port = cursor.get_u16();
+
+                Ok(Address::SocketAddress(SocketAddr::V6(SocketAddrV6::new(
+                    v6addr, port, 0, 0,
+                ))))
+            }
+            consts::SOCKS5_ADDR_TYPE_DOMAIN_NAME => {
+                let mut length_buf = [0u8; 1];
+                let _ = stream.read_exact(&mut length_buf).await?;
+                let length = length_buf[0] as usize;
+
+                // Len(Domain) + Len(Port)
+                let buf_length = length + 2;
+                let mut buf = BytesMut::with_capacity(buf_length);
+                buf.resize(buf_length, 0);
+                let _ = stream.read_exact(&mut buf).await?;
+
+                let mut cursor = buf.to_bytes();
+                let mut raw_addr = Vec::with_capacity(length);
+                raw_addr.put(&mut BufExt::take(&mut cursor, length));
+                let addr = match String::from_utf8(raw_addr) {
+                    Ok(addr) => addr,
+                    Err(..) => return Err(Error::new(Reply::GeneralFailure, "Invalid address encoding")),
+                };
+                let port = cursor.get_u16();
+
+                Ok(Address::DomainNameAddress(addr, port))
+            }
+            _ => {
+                error!("Invalid address type {}", addr_type);
+                Err(Error::new(Reply::AddressTypeNotSupported, "Not supported address type"))
+            }
+        }
     }
 
     /// Writes to writer
     #[inline]
-    pub fn write_to<W: AsyncWrite>(self, writer: W) -> WriteBytes<W, Bytes> {
+    pub async fn write_to<W>(&self, writer: &mut W) -> io::Result<()>
+    where
+        W: AsyncWrite + Unpin,
+    {
         let mut buf = BytesMut::with_capacity(self.serialized_len());
         self.write_to_buf(&mut buf);
-        write_bytes(writer, buf.freeze())
+        writer.write_all(&buf).await
     }
 
     /// Writes to buffer
@@ -287,215 +351,42 @@ impl From<(String, u16)> for Address {
     }
 }
 
-pub struct ReadAddress<R>
-where
-    R: AsyncRead,
-{
-    reader: Option<R>,
-    state: ReadAddressState,
-    buf: Option<BytesMut>,
-    already_read: usize,
-}
+/// Parse `Address` error
+#[derive(Debug)]
+pub struct AddressError;
 
-enum ReadAddressState {
-    Type,
-    Ipv4,
-    Ipv6,
-    DomainNameLength,
-    DomainName,
-}
+impl FromStr for Address {
+    type Err = AddressError;
 
-impl<R> Future for ReadAddress<R>
-where
-    R: AsyncRead,
-{
-    type Error = Error;
-    type Item = (R, Address);
-
-    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
-        debug_assert!(self.reader.is_some());
-
-        loop {
-            match self.state {
-                ReadAddressState::Type => {
-                    try_ready!(self.read_addr_type());
+    fn from_str(s: &str) -> Result<Address, AddressError> {
+        match s.parse::<SocketAddr>() {
+            Ok(addr) => Ok(Address::SocketAddress(addr)),
+            Err(..) => {
+                let mut sp = s.split(':');
+                match (sp.next(), sp.next()) {
+                    (Some(dn), Some(port)) => match port.parse::<u16>() {
+                        Ok(port) => Ok(Address::DomainNameAddress(dn.to_owned(), port)),
+                        Err(..) => Err(AddressError),
+                    },
+                    _ => Err(AddressError),
                 }
-                ReadAddressState::Ipv4 => {
-                    let addr = try_ready!(self.read_ipv4());
-                    let reader = self.reader.take().unwrap();
-                    return Ok((reader, addr).into());
-                }
-                ReadAddressState::Ipv6 => {
-                    let addr = try_ready!(self.read_ipv6());
-                    let reader = self.reader.take().unwrap();
-                    return Ok((reader, addr).into());
-                }
-                ReadAddressState::DomainNameLength => {
-                    try_ready!(self.read_domain_name_length());
-                }
-                ReadAddressState::DomainName => {
-                    let addr = try_ready!(self.read_domain_name());
-                    let reader = self.reader.take().unwrap();
-                    return Ok((reader, addr).into());
-                }
-            };
-        }
-    }
-}
-
-impl<R> ReadAddress<R>
-where
-    R: AsyncRead,
-{
-    fn new(r: R) -> ReadAddress<R> {
-        ReadAddress {
-            reader: Some(r),
-            state: ReadAddressState::Type,
-            buf: None,
-            already_read: 0,
-        }
-    }
-
-    fn read_addr_type(&mut self) -> Poll<(), Error> {
-        let addr_type = try_nb!(self.reader.as_mut().unwrap().read_u8());
-        match addr_type {
-            consts::SOCKS5_ADDR_TYPE_IPV4 => {
-                self.state = ReadAddressState::Ipv4;
-                self.alloc_buf(6);
-            }
-            consts::SOCKS5_ADDR_TYPE_IPV6 => {
-                self.state = ReadAddressState::Ipv6;
-                self.alloc_buf(18);
-            }
-            consts::SOCKS5_ADDR_TYPE_DOMAIN_NAME => {
-                self.state = ReadAddressState::DomainNameLength;
-            }
-            _ => {
-                error!("Invalid address type {}", addr_type);
-                return Err(Error::new(Reply::AddressTypeNotSupported, "Not supported address type"));
-            }
-        };
-
-        Ok(Async::Ready(()))
-    }
-
-    fn read_ipv4(&mut self) -> Poll<Address, Error> {
-        try_ready!(self.read_data());
-        let mut stream: Cursor<Bytes> = self.freeze_buf().into_buf();
-        let v4addr = Ipv4Addr::new(
-            stream.read_u8()?,
-            stream.read_u8()?,
-            stream.read_u8()?,
-            stream.read_u8()?,
-        );
-        let port = stream.read_u16::<BigEndian>()?;
-        let addr = Address::SocketAddress(SocketAddr::V4(SocketAddrV4::new(v4addr, port)));
-        Ok(Async::Ready(addr))
-    }
-
-    fn read_ipv6(&mut self) -> Poll<Address, Error> {
-        try_ready!(self.read_data());
-        let mut stream: Cursor<Bytes> = self.freeze_buf().into_buf();
-        let v6addr = Ipv6Addr::new(
-            stream.read_u16::<BigEndian>()?,
-            stream.read_u16::<BigEndian>()?,
-            stream.read_u16::<BigEndian>()?,
-            stream.read_u16::<BigEndian>()?,
-            stream.read_u16::<BigEndian>()?,
-            stream.read_u16::<BigEndian>()?,
-            stream.read_u16::<BigEndian>()?,
-            stream.read_u16::<BigEndian>()?,
-        );
-        let port = stream.read_u16::<BigEndian>()?;
-
-        let addr = Address::SocketAddress(SocketAddr::V6(SocketAddrV6::new(v6addr, port, 0, 0)));
-        Ok(Async::Ready(addr))
-    }
-
-    fn read_domain_name_length(&mut self) -> Poll<(), Error> {
-        let length = try_nb!(self.reader.as_mut().unwrap().read_u8());
-        self.state = ReadAddressState::DomainName;
-        self.alloc_buf(length as usize + 2);
-        Ok(Async::Ready(()))
-    }
-
-    fn read_domain_name(&mut self) -> Poll<Address, Error> {
-        try_ready!(self.read_data());
-        let buf = self.freeze_buf();
-        let addr_len = buf.len() - 2;
-        let mut stream: Cursor<Bytes> = buf.into_buf();
-
-        let mut raw_addr = Vec::with_capacity(addr_len);
-        unsafe {
-            raw_addr.set_len(addr_len);
-        }
-        stream.read_exact(&mut raw_addr)?;
-
-        let addr = match String::from_utf8(raw_addr) {
-            Ok(addr) => addr,
-            Err(..) => return Err(Error::new(Reply::GeneralFailure, "Invalid address encoding")),
-        };
-        let port = stream.read_u16::<BigEndian>()?;
-
-        let addr = Address::DomainNameAddress(addr, port);
-        Ok(Async::Ready(addr))
-    }
-
-    fn alloc_buf(&mut self, size: usize) {
-        let mut buf = BytesMut::with_capacity(size);
-        unsafe {
-            buf.set_len(size);
-        }
-        self.buf = Some(buf);
-    }
-
-    fn read_data(&mut self) -> Poll<(), io::Error> {
-        let buf = self.buf.as_mut().unwrap();
-
-        while self.already_read < buf.len() {
-            match self.reader.as_mut().unwrap().read(&mut buf[self.already_read..]) {
-                Ok(0) => {
-                    let err = io::Error::new(io::ErrorKind::Other, "Unexpected EOF");
-                    return Err(err);
-                }
-                Ok(n) => self.already_read += n,
-                Err(err) => return Err(err),
             }
         }
-
-        Ok(Async::Ready(()))
-    }
-
-    fn freeze_buf(&mut self) -> Bytes {
-        let buf = self.buf.take().unwrap();
-        buf.freeze()
     }
 }
 
 fn write_ipv4_address<B: BufMut>(addr: &SocketAddrV4, buf: &mut B) {
-    let mut dbuf = [0u8; 1 + 4 + 2];
-    {
-        let mut cur = Cursor::new(&mut dbuf[..]);
-        let _ = cur.write_u8(consts::SOCKS5_ADDR_TYPE_IPV4); // Address type
-        let _ = cur.write_all(&addr.ip().octets()); // Ipv4 bytes
-        let _ = cur.write_u16::<BigEndian>(addr.port());
-    }
-    buf.put_slice(&dbuf[..]);
+    buf.put_u8(consts::SOCKS5_ADDR_TYPE_IPV4); // Address type
+    buf.put_slice(&addr.ip().octets()); // Ipv4 bytes
+    buf.put_u16(addr.port()); // Port
 }
 
 fn write_ipv6_address<B: BufMut>(addr: &SocketAddrV6, buf: &mut B) {
-    let mut dbuf = [0u8; 1 + 16 + 2];
-
-    {
-        let mut cur = Cursor::new(&mut dbuf[..]);
-        let _ = cur.write_u8(consts::SOCKS5_ADDR_TYPE_IPV6);
-        for seg in &addr.ip().segments() {
-            let _ = cur.write_u16::<BigEndian>(*seg);
-        }
-        let _ = cur.write_u16::<BigEndian>(addr.port());
+    buf.put_u8(consts::SOCKS5_ADDR_TYPE_IPV6); // Address type
+    for seg in &addr.ip().segments() {
+        buf.put_u16(*seg); // Ipv6 bytes
     }
-
-    buf.put_slice(&dbuf[..]);
+    buf.put_u16(addr.port()); // Port
 }
 
 fn write_domain_name_address<B: BufMut>(dnaddr: &str, port: u16, buf: &mut B) {
@@ -504,7 +395,7 @@ fn write_domain_name_address<B: BufMut>(dnaddr: &str, port: u16, buf: &mut B) {
     buf.put_u8(consts::SOCKS5_ADDR_TYPE_DOMAIN_NAME);
     buf.put_u8(dnaddr.len() as u8);
     buf.put_slice(dnaddr[..].as_bytes());
-    buf.put_u16_be(port);
+    buf.put_u16(port);
 }
 
 fn write_socket_address<B: BufMut>(addr: &SocketAddr, buf: &mut B) {
@@ -557,42 +448,38 @@ impl TcpRequestHeader {
     }
 
     /// Read from a reader
-    pub fn read_from<R>(r: R) -> impl Future<Item = (R, TcpRequestHeader), Error = Error> + Send
+    pub async fn read_from<R>(r: &mut R) -> Result<TcpRequestHeader, Error>
     where
-        R: AsyncRead + Send + 'static,
+        R: AsyncRead + Unpin,
     {
-        read_exact(r, [0u8; 3])
-            .map_err(From::from)
-            .and_then(|(r, buf)| {
-                let ver = buf[0];
-                if ver != consts::SOCKS5_VERSION {
-                    return Err(Error::new(Reply::ConnectionRefused, "Unsupported Socks version"));
-                }
+        let mut buf = [0u8; 3];
+        let _ = r.read_exact(&mut buf).await?;
 
-                let cmd = buf[1];
-                let command = match Command::from_u8(cmd) {
-                    Some(c) => c,
-                    None => {
-                        return Err(Error::new(Reply::CommandNotSupported, "Unsupported command"));
-                    }
-                };
+        let ver = buf[0];
+        if ver != consts::SOCKS5_VERSION {
+            return Err(Error::new(Reply::ConnectionRefused, "Unsupported Socks version"));
+        }
 
-                Ok((r, command))
-            })
-            .and_then(|(r, command)| {
-                Address::read_from(r).map(move |(conn, address)| {
-                    let header = TcpRequestHeader { command, address };
+        let cmd = buf[1];
+        let command = match Command::from_u8(cmd) {
+            Some(c) => c,
+            None => {
+                return Err(Error::new(Reply::CommandNotSupported, "Unsupported command"));
+            }
+        };
 
-                    (conn, header)
-                })
-            })
+        let address = Address::read_from(r).await?;
+        Ok(TcpRequestHeader { command, address })
     }
 
     /// Write data into a writer
-    pub fn write_to<W: AsyncWrite>(self, w: W) -> WriteBytes<W, Bytes> {
+    pub async fn write_to<W>(&self, w: &mut W) -> io::Result<()>
+    where
+        W: AsyncWrite + Unpin,
+    {
         let mut buf = BytesMut::with_capacity(self.serialized_len());
         self.write_to_buf(&mut buf);
-        write_bytes(w, buf.freeze())
+        w.write_all(&buf).await
     }
 
     /// Writes to buffer
@@ -637,39 +524,36 @@ impl TcpResponseHeader {
     }
 
     /// Read from a reader
-    pub fn read_from<R>(r: R) -> impl Future<Item = (R, TcpResponseHeader), Error = Error> + Send
+    pub async fn read_from<R>(r: &mut R) -> Result<TcpResponseHeader, Error>
     where
-        R: AsyncRead + Send + 'static,
+        R: AsyncRead + Unpin,
     {
-        read_exact(r, [0u8; 3])
-            .map_err(From::from)
-            .and_then(|(r, buf)| {
-                let ver = buf[0];
-                let reply_code = buf[1];
+        let mut buf = [0u8; 3];
+        let _ = r.read_exact(&mut buf).await?;
 
-                if ver != consts::SOCKS5_VERSION {
-                    return Err(Error::new(Reply::ConnectionRefused, "Unsupported Socks version"));
-                }
+        let ver = buf[0];
+        let reply_code = buf[1];
 
-                Ok((r, reply_code))
-            })
-            .and_then(|(r, reply_code)| {
-                Address::read_from(r).map(move |(r, address)| {
-                    let rep = TcpResponseHeader {
-                        reply: Reply::from_u8(reply_code),
-                        address,
-                    };
+        if ver != consts::SOCKS5_VERSION {
+            return Err(Error::new(Reply::ConnectionRefused, "Unsupported Socks version"));
+        }
 
-                    (r, rep)
-                })
-            })
+        let address = Address::read_from(r).await?;
+
+        Ok(TcpResponseHeader {
+            reply: Reply::from_u8(reply_code),
+            address,
+        })
     }
 
     /// Write to a writer
-    pub fn write_to<W: AsyncWrite>(self, w: W) -> WriteBytes<W, Bytes> {
+    pub async fn write_to<W>(&self, w: &mut W) -> io::Result<()>
+    where
+        W: AsyncWrite + Unpin,
+    {
         let mut buf = BytesMut::with_capacity(self.serialized_len());
         self.write_to_buf(&mut buf);
-        write_bytes(w, buf.freeze())
+        w.write_all(&buf).await
     }
 
     /// Writes to buffer
@@ -707,30 +591,34 @@ impl HandshakeRequest {
     }
 
     /// Read from a reader
-    pub fn read_from<R>(r: R) -> impl Future<Item = (R, HandshakeRequest), Error = io::Error>
+    pub async fn read_from<R>(r: &mut R) -> io::Result<HandshakeRequest>
     where
-        R: AsyncRead + Send + 'static,
+        R: AsyncRead + Unpin,
     {
-        read_exact(r, [0u8, 0u8])
-            .and_then(|(r, buf)| {
-                let ver = buf[0];
-                let nmet = buf[1];
+        let mut buf = [0u8; 2];
+        let _ = r.read_exact(&mut buf).await?;
 
-                if ver != consts::SOCKS5_VERSION {
-                    return Err(io::Error::new(io::ErrorKind::Other, "Invalid Socks5 version"));
-                }
+        let ver = buf[0];
+        let nmet = buf[1];
 
-                Ok((r, nmet))
-            })
-            .and_then(|(r, nmet)| read_exact(r, vec![0u8; nmet as usize]))
-            .and_then(|(r, methods)| Ok((r, HandshakeRequest { methods })))
+        if ver != consts::SOCKS5_VERSION {
+            return Err(io::Error::new(io::ErrorKind::Other, "Invalid Socks5 version"));
+        }
+
+        let mut methods = vec![0u8; nmet as usize];
+        let _ = r.read_exact(&mut methods).await?;
+
+        Ok(HandshakeRequest { methods })
     }
 
     /// Write to a writer
-    pub fn write_to<W: AsyncWrite>(self, w: W) -> WriteBytes<W, Bytes> {
+    pub async fn write_to<W>(&self, w: &mut W) -> io::Result<()>
+    where
+        W: AsyncWrite + Unpin,
+    {
         let mut buf = BytesMut::with_capacity(self.serialized_len());
         self.write_to_buf(&mut buf);
-        write_bytes(w, buf.freeze())
+        w.write_all(&buf).await
     }
 
     /// Write to buffer
@@ -767,36 +655,40 @@ impl HandshakeResponse {
     }
 
     /// Read from a reader
-    pub fn read_from<R>(r: R) -> impl Future<Item = (R, HandshakeResponse), Error = io::Error>
+    pub async fn read_from<R>(r: &mut R) -> io::Result<HandshakeResponse>
     where
-        R: AsyncRead + Send + 'static,
+        R: AsyncRead + Unpin,
     {
-        read_exact(r, [0u8, 0u8]).and_then(|(r, buf)| {
-            let ver = buf[0];
-            let met = buf[1];
+        let mut buf = [0u8; 2];
+        let _ = r.read_exact(&mut buf).await?;
 
-            if ver != consts::SOCKS5_VERSION {
-                Err(io::Error::new(io::ErrorKind::Other, "Invalid Socks5 version"))
-            } else {
-                Ok((r, HandshakeResponse { chosen_method: met }))
-            }
-        })
+        let ver = buf[0];
+        let met = buf[1];
+
+        if ver != consts::SOCKS5_VERSION {
+            Err(io::Error::new(io::ErrorKind::Other, "Invalid Socks5 version"))
+        } else {
+            Ok(HandshakeResponse { chosen_method: met })
+        }
     }
 
     /// Write to a writer
-    pub fn write_to<W: AsyncWrite>(self, w: W) -> WriteBytes<W, Bytes> {
+    pub async fn write_to<W>(self, w: &mut W) -> io::Result<()>
+    where
+        W: AsyncWrite + Unpin,
+    {
         let mut buf = BytesMut::with_capacity(self.serialized_len());
         self.write_to_buf(&mut buf);
-        write_bytes(w, buf.freeze())
+        w.write_all(&buf).await
     }
 
     /// Write to buffer
-    pub fn write_to_buf<B: BufMut>(&self, buf: &mut B) {
+    pub fn write_to_buf<B: BufMut>(self, buf: &mut B) {
         buf.put_slice(&[consts::SOCKS5_VERSION, self.chosen_method]);
     }
 
     /// Length in bytes
-    pub fn serialized_len(&self) -> usize {
+    pub fn serialized_len(self) -> usize {
         2
     }
 }
@@ -827,24 +719,26 @@ impl UdpAssociateHeader {
     }
 
     /// Read from a reader
-    pub fn read_from<R>(r: R) -> impl Future<Item = (R, UdpAssociateHeader), Error = Error> + Send
+    pub async fn read_from<R>(r: &mut R) -> Result<UdpAssociateHeader, Error>
     where
-        R: AsyncRead + Send + 'static,
+        R: AsyncRead + Unpin,
     {
-        read_exact(r, [0u8; 3]).map_err(From::from).and_then(|(r, buf)| {
-            let frag = buf[2];
-            Address::read_from(r).map(move |(r, address)| {
-                let h = UdpAssociateHeader::new(frag, address);
-                (r, h)
-            })
-        })
+        let mut buf = [0u8; 3];
+        let _ = r.read_exact(&mut buf).await?;
+
+        let frag = buf[2];
+        let address = Address::read_from(r).await?;
+        Ok(UdpAssociateHeader::new(frag, address))
     }
 
     /// Write to a writer
-    pub fn write_to<W: AsyncWrite>(self, w: W) -> WriteBytes<W, Bytes> {
+    pub async fn write_to<W>(&self, w: &mut W) -> io::Result<()>
+    where
+        W: AsyncWrite + Unpin,
+    {
         let mut buf = BytesMut::with_capacity(self.serialized_len());
         self.write_to_buf(&mut buf);
-        write_bytes(w, buf.freeze())
+        w.write_all(&buf).await
     }
 
     /// Write to buffer

src/relay/tcprelay/aead.rs

@@ -33,295 +33,257 @@
 //! ```
 
 use std::{
-    cmp,
-    io::{self, BufRead, Read},
+    cmp, io,
+    marker::Unpin,
+    pin::Pin,
+    slice,
+    task::{Context, Poll},
     u16,
 };
 
 use byteorder::{BigEndian, ByteOrder};
 use bytes::{BufMut, BytesMut};
-use log::error;
-use tokio_io::{AsyncRead, AsyncWrite};
+use futures::ready;
+use tokio::prelude::*;
 
 use crate::crypto::{self, BoxAeadDecryptor, BoxAeadEncryptor, CipherType};
 
-use super::{DecryptedRead, EncryptedWrite, BUFFER_SIZE};
+use super::BUFFER_SIZE;
 
 /// AEAD packet payload must be smaller than 0x3FFF
 const MAX_PACKET_SIZE: usize = 0x3FFF;
 
-enum ReadingStep {
+#[derive(Debug)]
+enum DecryptReadStep {
     Length,
-    DataAndTag(usize),
-    Done,
+    Data(usize),
 }
 
 /// Reader wrapper that will decrypt data automatically
-pub struct DecryptedReader<R>
-where
-    R: AsyncRead,
-{
-    reader: R,
+pub struct DecryptedReader {
     buffer: BytesMut,
     data: BytesMut,
     cipher: BoxAeadDecryptor,
     pos: usize,
-    sent_final: bool,
     tag_size: usize,
-    read_step: ReadingStep,
+    steps: DecryptReadStep,
+    got_final: bool,
 }
 
-impl<R> DecryptedReader<R>
-where
-    R: AsyncRead,
-{
-    pub fn new(r: R, t: CipherType, key: &[u8], nounce: &[u8]) -> DecryptedReader<R> {
+impl DecryptedReader {
+    pub fn new(t: CipherType, key: &[u8], nonce: &[u8]) -> DecryptedReader {
         DecryptedReader {
-            reader: r,
             buffer: BytesMut::with_capacity(BUFFER_SIZE),
             data: BytesMut::with_capacity(BUFFER_SIZE),
-            cipher: crypto::new_aead_decryptor(t, key, nounce),
+            cipher: crypto::new_aead_decryptor(t, key, nonce),
             pos: 0,
-            sent_final: false,
             tag_size: t.tag_size(),
-            read_step: ReadingStep::Length,
+            steps: DecryptReadStep::Length,
+            got_final: false,
         }
     }
 
-    pub fn get_ref(&self) -> &R {
-        &self.reader
-    }
-
-    /// Gets a mutable reference to the underlying reader.
-    ///
-    /// # Warning
-    ///
-    /// It is inadvisable to read directly from or write directly to the
-    /// underlying reader.
-    pub fn get_mut(&mut self) -> &mut R {
-        &mut self.reader
-    }
-
-    /// Unwraps this `DecryptedReader`, returning the underlying reader.
-    ///
-    /// The internal buffer is flushed before returning the reader. Any leftover
-    /// data in the read buffer is lost.
-    pub fn into_inner(self) -> R {
-        self.reader
-    }
-
-    fn read_exact(&mut self, expect_length: usize, ignore_final: bool) -> io::Result<()> {
-        let mut incoming = [0u8; BUFFER_SIZE];
-        self.buffer.reserve(expect_length);
-
-        while self.buffer.len() < expect_length {
-            let remain = expect_length - self.buffer.len();
-            let rlen = cmp::min(remain, incoming.len());
-            match self.reader.read(&mut incoming[..rlen]) {
-                Ok(0) => {
-                    if ignore_final && self.buffer.is_empty() {
-                        self.sent_final = true;
-                        return Ok(());
-                    }
-
-                    return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "unexpected eof"));
-                }
-                Ok(n) => self.buffer.put_slice(&incoming[..n]),
-                Err(e) => return Err(e),
-            }
-        }
-        Ok(())
-    }
-
-    fn read_length(&mut self) -> io::Result<()> {
-        let expect_length = 2 + self.tag_size;
-        self.read_exact(expect_length, true)?;
-
-        if !self.sent_final {
-            // Ok, read length finished
-            {
-                let len = {
-                    let mut len_buf = [0u8; 2];
-                    self.cipher.decrypt(&self.buffer[..], &mut len_buf)?;
-                    BigEndian::read_u16(&len_buf) as usize
-                };
-
-                if len > MAX_PACKET_SIZE {
-                    use std::io::{Error, ErrorKind};
-
-                    error!(
-                        "AEAD packet size must be <= {}, but received length {}",
-                        MAX_PACKET_SIZE, len
-                    );
-
-                    let err = Error::new(
-                        ErrorKind::InvalidData,
-                        format!(
-                            "AEAD packet size must be <= {}, but received length {}",
-                            MAX_PACKET_SIZE, len
-                        ),
-                    );
-                    return Err(err);
-                }
-
-                self.read_step = ReadingStep::DataAndTag(len);
-            }
-            self.buffer.clear();
-        }
-
-        Ok(())
-    }
-
-    fn read_data(&mut self, dlen: usize) -> io::Result<()> {
-        let expect_length = dlen + self.tag_size;
-        self.read_exact(expect_length, false)?;
-
-        if !self.sent_final {
-            {
-                // Ok, got data
-                self.data.clear();
-                self.data.reserve(dlen);
-                unsafe {
-                    self.data.set_len(dlen); // Decrypted data has exactly the same length
-                }
-                self.cipher.decrypt(&self.buffer[..], &mut *self.data)?;
-            }
-
-            self.read_step = ReadingStep::Done;
-            self.buffer.clear();
-        }
-
-        Ok(())
-    }
-
-    fn read_some(&mut self) -> io::Result<()> {
-        while !self.sent_final {
-            match self.read_step {
-                ReadingStep::Length => self.read_length()?,
-                ReadingStep::DataAndTag(dlen) => {
-                    self.read_data(dlen)?;
-                    break; // Read finished! Break out
-                }
-                ReadingStep::Done => {
-                    self.read_step = ReadingStep::Length;
-                    self.data.clear();
-                }
-            }
-        }
-        Ok(())
-    }
-}
-
-impl<R> BufRead for DecryptedReader<R>
-where
-    R: AsyncRead,
-{
-    fn fill_buf(&mut self) -> io::Result<&[u8]> {
+    pub fn poll_read_decrypted<R>(
+        &mut self,
+        ctx: &mut Context<'_>,
+        r: &mut R,
+        dst: &mut [u8],
+    ) -> Poll<io::Result<usize>>
+    where
+        R: AsyncRead + Unpin,
+    {
         while self.pos >= self.data.len() {
-            if self.sent_final {
-                return Ok(&[]);
+            // Already received EOF
+            if self.got_final {
+                return Poll::Ready(Ok(0));
             }
 
-            self.read_some()?;
-            if let ReadingStep::Done = self.read_step {
-                self.pos = 0;
+            // Refill buffer
+            match self.steps {
+                DecryptReadStep::Length => ready!(self.poll_read_decrypted_length(ctx, r))?,
+                DecryptReadStep::Data(len) => ready!(self.poll_read_decrypted_data(ctx, r, len))?,
             }
         }
 
-        Ok(&self.data[self.pos..])
+        let remaining_len = self.data.len() - self.pos;
+        let n = cmp::min(dst.len(), remaining_len);
+        (&mut dst[..n]).copy_from_slice(&self.data[self.pos..self.pos + n]);
+        self.pos += n;
+        Poll::Ready(Ok(n))
     }
 
-    fn consume(&mut self, amt: usize) {
-        self.pos = cmp::min(self.pos + amt, self.data.len());
-    }
-}
+    fn poll_read_decrypted_length<R>(&mut self, ctx: &mut Context<'_>, r: &mut R) -> Poll<io::Result<()>>
+    where
+        R: AsyncRead + Unpin,
+    {
+        let buf_len = 2 + self.tag_size;
+        ready!(self.poll_read_exact(ctx, r, buf_len, true))?;
+        if self.got_final {
+            return Poll::Ready(Ok(()));
+        }
 
-impl<R> Read for DecryptedReader<R>
-where
-    R: AsyncRead,
-{
-    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        let nread = {
-            let mut available = self.fill_buf()?;
-            available.read(buf)?
+        // Done reading, decrypt it
+        let len = {
+            let mut len_buf = [0u8; 2];
+            self.cipher.decrypt(&self.buffer[..], &mut len_buf)?;
+            BigEndian::read_u16(&len_buf) as usize
         };
-        self.consume(nread);
-        Ok(nread)
+
+        // Clear buffer before overwriting it
+        self.buffer.clear();
+        self.data.clear();
+        self.pos = 0;
+
+        // Next step, read data
+        self.steps = DecryptReadStep::Data(len);
+        self.buffer.reserve(len + self.tag_size);
+        self.data.reserve(len);
+
+        Poll::Ready(Ok(()))
+    }
+
+    fn poll_read_decrypted_data<R>(&mut self, ctx: &mut Context<'_>, r: &mut R, size: usize) -> Poll<io::Result<()>>
+    where
+        R: AsyncRead + Unpin,
+    {
+        let buf_len = size + self.tag_size;
+        ready!(self.poll_read_exact(ctx, r, buf_len, false))?;
+
+        // Done reading data, decrypt it
+        unsafe {
+            // It has enough space, I am sure about that
+            let buffer = slice::from_raw_parts_mut(self.data.bytes_mut().as_mut_ptr() as *mut u8, size);
+            self.cipher.decrypt(&self.buffer[..], buffer)?;
+
+            // Move forward the pointer
+            self.data.advance_mut(size);
+        }
+
+        // Clear buffer before overwriting it
+        self.buffer.clear();
+
+        // Reset read position
+        self.pos = 0;
+
+        // Next step, read length
+        self.steps = DecryptReadStep::Length;
+        self.buffer.reserve(2 + self.tag_size);
+
+        Poll::Ready(Ok(()))
+    }
+
+    fn poll_read_exact<R>(
+        &mut self,
+        ctx: &mut Context<'_>,
+        r: &mut R,
+        size: usize,
+        allow_eof: bool,
+    ) -> Poll<io::Result<()>>
+    where
+        R: AsyncRead + Unpin,
+    {
+        while self.buffer.len() < size {
+            let remaining = size - self.buffer.len();
+            unsafe {
+                // It has enough space, I am sure about that
+                let buffer = slice::from_raw_parts_mut(self.buffer.bytes_mut().as_mut_ptr() as *mut u8, remaining);
+                let n = ready!(Pin::new(&mut *r).poll_read(ctx, buffer))?;
+                if n == 0 {
+                    if self.buffer.is_empty() && allow_eof && !self.got_final {
+                        // Read nothing
+                        self.got_final = true;
+                        return Poll::Ready(Ok(()));
+                    } else {
+                        use std::io::ErrorKind;
+                        return Poll::Ready(Err(ErrorKind::UnexpectedEof.into()));
+                    }
+                }
+                self.buffer.advance_mut(n);
+            }
+        }
+
+        Poll::Ready(Ok(()))
     }
 }
 
-impl<R> DecryptedRead for DecryptedReader<R>
-where
-    R: AsyncRead,
-{
-    fn buffer_size(&self, data: &[u8]) -> usize {
-        2 + self.tag_size // len and len_tag
-        + data.len() + self.tag_size // data and data_tag
-    }
+enum EncryptWriteStep {
+    Nothing,
+    Writing(BytesMut, usize),
 }
 
-impl<R> AsyncRead for DecryptedReader<R> where R: AsyncRead {}
-
 /// Writer wrapper that will encrypt data automatically
-pub struct EncryptedWriter<W>
-where
-    W: AsyncWrite,
-{
-    writer: W,
+pub struct EncryptedWriter {
     cipher: BoxAeadEncryptor,
     tag_size: usize,
+    steps: EncryptWriteStep,
 }
 
-impl<W> EncryptedWriter<W>
-where
-    W: AsyncWrite,
-{
+impl EncryptedWriter {
     /// Creates a new EncryptedWriter
-    pub fn new(w: W, t: CipherType, key: &[u8], nonce: &[u8]) -> EncryptedWriter<W> {
+    pub fn new(t: CipherType, key: &[u8], nonce: &[u8]) -> EncryptedWriter {
         EncryptedWriter {
-            writer: w,
             cipher: crypto::new_aead_encryptor(t, key, nonce),
             tag_size: t.tag_size(),
+            steps: EncryptWriteStep::Nothing,
         }
     }
-}
 
-impl<W> EncryptedWrite for EncryptedWriter<W>
-where
-    W: AsyncWrite,
-{
-    fn write_raw(&mut self, data: &[u8]) -> io::Result<usize> {
-        self.writer.write(data)
+    pub fn poll_write_encrypted<W>(&mut self, ctx: &mut Context<'_>, w: &mut W, data: &[u8]) -> Poll<io::Result<usize>>
+    where
+        W: AsyncWrite + Unpin,
+    {
+        ready!(self.poll_write_all_encrypted(ctx, w, data))?;
+        Poll::Ready(Ok(data.len()))
     }
 
-    fn flush(&mut self) -> io::Result<()> {
-        self.writer.flush()
-    }
-
-    fn encrypt<B: BufMut>(&mut self, data: &[u8], buf: &mut B) -> io::Result<()> {
+    pub fn poll_write_all_encrypted<W>(&mut self, ctx: &mut Context<'_>, w: &mut W, data: &[u8]) -> Poll<io::Result<()>>
+    where
+        W: AsyncWrite + Unpin,
+    {
         // Data.Len is a 16-bit big-endian integer indicating the length of Data. It should be smaller than 0x3FFF.
         assert!(
             data.len() <= MAX_PACKET_SIZE,
             "Buffer size too large, AEAD encryption protocol requires buffer to be smaller than 0x3FFF"
         );
 
-        let output_length = self.buffer_size(data);
-        let data_length = data.len() as u16;
+        loop {
+            match self.steps {
+                EncryptWriteStep::Nothing => {
+                    let output_length = self.buffer_size(data);
+                    let data_length = data.len() as u16;
 
-        let mut data_len_buf = [0u8; 2];
-        BigEndian::write_u16(&mut data_len_buf, data_length);
+                    let mut buf = BytesMut::with_capacity(output_length);
 
-        let output_length_size = 2 + self.tag_size;
-        self.cipher
-            .encrypt(&data_len_buf, unsafe { &mut buf.bytes_mut()[..output_length_size] });
-        self.cipher
-            .encrypt(data, unsafe { &mut buf.bytes_mut()[output_length_size..output_length] });
+                    let mut data_len_buf = [0u8; 2];
+                    BigEndian::write_u16(&mut data_len_buf, data_length);
 
-        unsafe {
-            buf.advance_mut(output_length);
+                    unsafe {
+                        let b = slice::from_raw_parts_mut(buf.bytes_mut().as_mut_ptr() as *mut u8, output_length);
+
+                        let output_length_size = 2 + self.tag_size;
+                        self.cipher.encrypt(&data_len_buf, &mut b[..output_length_size]);
+                        self.cipher.encrypt(data, &mut b[output_length_size..output_length]);
+
+                        buf.advance_mut(output_length);
+                    }
+
+                    self.steps = EncryptWriteStep::Writing(buf, 0);
+                }
+                EncryptWriteStep::Writing(ref mut buf, ref mut pos) => {
+                    while *pos < buf.len() {
+                        let n = ready!(Pin::new(&mut *w).poll_write(ctx, &buf[*pos..]))?;
+                        if n == 0 {
+                            use std::io::ErrorKind;
+                            return Poll::Ready(Err(ErrorKind::UnexpectedEof.into()));
+                        }
+                        *pos += n;
+                    }
+
+                    self.steps = EncryptWriteStep::Nothing;
+                    return Poll::Ready(Ok(()));
+                }
+            }
         }
-
-        Ok(())
     }
 
     fn buffer_size(&self, data: &[u8]) -> usize {

src/relay/tcprelay/client.rs

@@ -1,20 +1,21 @@
 //! TCP relay client implementation
 
 use std::{
-    io::{self, Read, Write},
+    io,
     net::SocketAddr,
+    pin::Pin,
     sync::Arc,
+    task::{Context, Poll},
 };
 
-use futures::{self, Async, Future, Poll};
 use log::trace;
-use tokio::net::TcpStream;
-use tokio_io::{io::flush, AsyncRead, AsyncWrite};
+use tokio::{net::TcpStream, prelude::*};
 
 use crate::relay::socks5::{
     self, Address, Command, HandshakeRequest, HandshakeResponse, Reply, TcpRequestHeader, TcpResponseHeader,
 };
 
+use super::{CryptoStream, STcpStream};
 use crate::{config::ServerConfig, context::SharedContext};
 
 /// Socks5 proxy client
@@ -24,131 +25,117 @@ pub struct Socks5Client {
 
 impl Socks5Client {
     /// Connects to `addr` via `proxy`
-    pub fn connect<A>(addr: A, proxy: SocketAddr) -> impl Future<Item = Socks5Client, Error = io::Error> + Send
+    pub async fn connect<A>(addr: A, proxy: &SocketAddr) -> io::Result<Socks5Client>
     where
         Address: From<A>,
-        A: Send + 'static,
     {
-        futures::lazy(move || TcpStream::connect(&proxy))
-            .and_then(move |s| {
-                // 1. Handshake
-                let hs = HandshakeRequest::new(vec![socks5::SOCKS5_AUTH_METHOD_NONE]);
-                trace!("Client connected, going to send handshake: {:?}", hs);
+        let mut s = TcpStream::connect(proxy).await?;
 
-                hs.write_to(s)
-                    .and_then(flush)
-                    .and_then(HandshakeResponse::read_from)
-                    .and_then(|(s, rsp)| {
-                        trace!("Got handshake response: {:?}", rsp);
-                        assert_eq!(rsp.chosen_method, socks5::SOCKS5_AUTH_METHOD_NONE);
-                        Ok(s)
-                    })
-            })
-            .and_then(move |s| {
-                // 2. Send request header
-                let h = TcpRequestHeader::new(Command::TcpConnect, From::from(addr));
-                trace!("Going to connect, req: {:?}", h);
-                h.write_to(s)
-                    .and_then(flush)
-                    .and_then(|s| TcpResponseHeader::read_from(s).map_err(From::from))
-                    .and_then(|(s, rsp)| {
-                        trace!("Got response: {:?}", rsp);
-                        match rsp.reply {
-                            Reply::Succeeded => Ok(s),
-                            r => {
-                                let err = io::Error::new(io::ErrorKind::Other, format!("{}", r));
-                                Err(err)
-                            }
-                        }
-                    })
-            })
-            .map(|s| Socks5Client { stream: s })
+        // 1. Handshake
+        let hs = HandshakeRequest::new(vec![socks5::SOCKS5_AUTH_METHOD_NONE]);
+        trace!("Client connected, going to send handshake: {:?}", hs);
+
+        hs.write_to(&mut s).await?;
+        s.flush().await?;
+
+        let hsp = HandshakeResponse::read_from(&mut s).await?;
+
+        trace!("Got handshake response: {:?}", hsp);
+        assert_eq!(hsp.chosen_method, socks5::SOCKS5_AUTH_METHOD_NONE);
+
+        // 2. Send request header
+        let h = TcpRequestHeader::new(Command::TcpConnect, From::from(addr));
+        trace!("Going to connect, req: {:?}", h);
+
+        h.write_to(&mut s).await?;
+        s.flush().await?;
+        let hp = TcpResponseHeader::read_from(&mut s).await?;
+
+        trace!("Got response: {:?}", hp);
+        match hp.reply {
+            Reply::Succeeded => (),
+            r => {
+                let err = io::Error::new(io::ErrorKind::Other, format!("{}", r));
+                return Err(err);
+            }
+        }
+
+        Ok(Socks5Client { stream: s })
     }
 
     /// UDP Associate `addr` via `proxy`
-    pub fn udp_associate<A>(
-        addr: A,
-        proxy: SocketAddr,
-    ) -> impl Future<Item = (Socks5Client, Address), Error = io::Error> + Send
+    pub async fn udp_associate<A>(addr: A, proxy: &SocketAddr) -> io::Result<(Socks5Client, Address)>
     where
         Address: From<A>,
-        A: Send + 'static,
     {
-        futures::lazy(move || TcpStream::connect(&proxy))
-            .and_then(move |s| {
-                // 1. Handshake
-                let hs = HandshakeRequest::new(vec![socks5::SOCKS5_AUTH_METHOD_NONE]);
-                trace!("Client connected, going to send handshake: {:?}", hs);
+        let mut s = TcpStream::connect(proxy).await?;
 
-                hs.write_to(s)
-                    .and_then(flush)
-                    .and_then(HandshakeResponse::read_from)
-                    .and_then(|(s, rsp)| {
-                        trace!("Got handshake response: {:?}", rsp);
-                        assert_eq!(rsp.chosen_method, socks5::SOCKS5_AUTH_METHOD_NONE);
-                        Ok(s)
-                    })
-            })
-            .and_then(move |s| {
-                // 2. Send request header
-                let h = TcpRequestHeader::new(Command::UdpAssociate, From::from(addr));
-                trace!("Going to connect, req: {:?}", h);
-                h.write_to(s)
-                    .and_then(flush)
-                    .and_then(|s| TcpResponseHeader::read_from(s).map_err(From::from))
-                    .and_then(|(s, rsp)| {
-                        trace!("Got response: {:?}", rsp);
-                        match rsp.reply {
-                            Reply::Succeeded => Ok((s, rsp.address)),
-                            r => {
-                                let err = io::Error::new(io::ErrorKind::Other, format!("{}", r));
-                                Err(err)
-                            }
-                        }
-                    })
-            })
-            .map(|(s, a)| (Socks5Client { stream: s }, a))
+        // 1. Handshake
+        let hs = HandshakeRequest::new(vec![socks5::SOCKS5_AUTH_METHOD_NONE]);
+        trace!("Client connected, going to send handshake: {:?}", hs);
+
+        hs.write_to(&mut s).await?;
+        s.flush().await?;
+
+        let hsp = HandshakeResponse::read_from(&mut s).await?;
+
+        trace!("Got handshake response: {:?}", hsp);
+        assert_eq!(hsp.chosen_method, socks5::SOCKS5_AUTH_METHOD_NONE);
+
+        // 2. Send request header
+        let h = TcpRequestHeader::new(Command::UdpAssociate, From::from(addr));
+        trace!("Going to connect, req: {:?}", h);
+
+        h.write_to(&mut s).await?;
+        s.flush().await?;
+        let hp = TcpResponseHeader::read_from(&mut s).await?;
+
+        trace!("Got response: {:?}", hp);
+        match hp.reply {
+            Reply::Succeeded => (),
+            r => {
+                let err = io::Error::new(io::ErrorKind::Other, format!("{}", r));
+                return Err(err);
+            }
+        }
+
+        Ok((Socks5Client { stream: s }, hp.address))
     }
 }
 
-impl Read for Socks5Client {
-    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        self.stream.read(buf)
+impl AsyncRead for Socks5Client {
+    fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context, buf: &mut [u8]) -> Poll<Result<usize, io::Error>> {
+        Pin::new(&mut self.stream).poll_read(cx, buf)
     }
 }
 
-impl Write for Socks5Client {
-    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        self.stream.write(buf)
-    }
-
-    fn flush(&mut self) -> io::Result<()> {
-        self.stream.flush()
-    }
-}
-
-impl AsyncRead for Socks5Client {}
 impl AsyncWrite for Socks5Client {
-    fn shutdown(&mut self) -> Poll<(), io::Error> {
-        // FIXME: Finalize the internal cipher
-        Ok(Async::Ready(()))
+    fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<Result<usize, io::Error>> {
+        Pin::new(&mut self.stream).poll_write(cx, buf)
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), io::Error>> {
+        Pin::new(&mut self.stream).poll_flush(cx)
+    }
+
+    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), io::Error>> {
+        Pin::new(&mut self.stream).poll_shutdown(cx)
     }
 }
 
 pub(crate) struct ServerClient {
-    pub r: super::DecryptedHalf<TcpStream>,
-    pub w: super::EncryptedHalf<TcpStream>,
+    pub stream: CryptoStream<STcpStream>,
 }
 
 impl ServerClient {
-    pub(crate) fn connect(
+    pub(crate) async fn connect(
         context: SharedContext,
-        addr: Address,
+        addr: &Address,
         svr_cfg: Arc<ServerConfig>,
-    ) -> impl Future<Item = ServerClient, Error = io::Error> {
-        super::connect_proxy_server(context, svr_cfg.clone())
-            .and_then(|svr_s| super::proxy_server_handshake(svr_s, svr_cfg, addr))
-            .and_then(|(svr_r, svr_w)| svr_r.join(svr_w))
-            .map(|(r, w)| ServerClient { r, w })
+    ) -> io::Result<ServerClient> {
+        let stream = super::connect_proxy_server(context, svr_cfg.clone()).await?;
+        Ok(ServerClient {
+            stream: super::proxy_server_handshake(stream, svr_cfg, addr).await?,
+        })
     }
 }

src/relay/tcprelay/context.rs

@@ -4,7 +4,7 @@ use std::{
     io,
     net::SocketAddr,
     sync::{
-        atomic::{AtomicBool, AtomicUsize, Ordering},
+        atomic::{AtomicUsize, Ordering},
         Arc,
     },
     time::Duration,
@@ -13,19 +13,16 @@ use std::{
 use crate::{
     config::{ServerAddr, ServerConfig},
     context::SharedContext,
-    relay::{boxed_future, dns_resolver::resolve},
 };
 
-use futures::{future, Future, Stream};
 use log::debug;
-use tokio::{self, net::UdpSocket, timer::Interval};
+use tokio::{self, net::UdpSocket, time};
 
 /// TCP Relay Server Context
 pub struct TcpServerContext {
     tx: AtomicUsize,
     rx: AtomicUsize,
     context: SharedContext,
-    stop_flag: Arc<AtomicBool>,
     svr_cfg: Arc<ServerConfig>,
 }
 
@@ -36,34 +33,28 @@ 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,
-            stop_flag: stop_flag.clone(),
             svr_cfg,
         };
 
         let ctx = Arc::new(ctx);
 
         if ctx.context.config().manager_address.is_some() {
-            let ctx2 = ctx.clone();
+            let ctx = ctx.clone();
+            tokio::spawn(async move {
+                let mut interval = time::interval(UPDATE_INTERVAL);
+                while ctx.context.server_running() {
+                    interval.tick().await;
 
-            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(|_| ()));
+                    let ctx = ctx.clone();
+                    tokio::spawn(async move {
+                        let _ = ctx.stat_interval().await;
+                    });
+                }
+            });
         }
 
         ctx
@@ -85,32 +76,41 @@ impl TcpServerContext {
         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!(),
-        };
-
+    async fn stat_interval(&self) -> io::Result<()> {
         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(|_| ())
-            })
-    }
+        let any_addr = "0.0.0.0:0".parse::<SocketAddr>().unwrap();
+        let mut socket = UdpSocket::bind(&any_addr).await?;
 
-    pub fn close(&self) {
-        self.stop_flag.store(true, Ordering::Release)
+        let payload = format!("stat: {{\"{}\": {}}}", svr_cfg.addr().port(), transmission);
+
+        let maddr = self
+            .context
+            .config()
+            .manager_address
+            .as_ref()
+            .expect("manager_address must not be None");
+
+        debug!("Sending Tcp Relay to {}, payload: {}", maddr, payload);
+
+        match maddr {
+            ServerAddr::SocketAddr(ref addr) => {
+                socket.send_to(payload.as_ref(), addr).await?;
+            }
+            #[cfg(feature = "trust-dns")]
+            ServerAddr::DomainName(ref domain, ref port) => {
+                use crate::relay::dns_resolver::resolve;
+
+                let addrs = resolve(self.context.clone(), &domain[..], *port, false).await?;
+                socket.send_to(payload.as_ref(), addrs[0]).await?;
+            }
+            #[cfg(not(feature = "trust-dns"))]
+            ServerAddr::DomainName(ref domain, ref port) => {
+                socket.send_to(payload.as_ref(), (domain.as_str(), *port)).await?;
+            }
+        }
+
+        Ok(())
     }
 }

src/relay/tcprelay/crypto_io.rs

@@ -1,450 +1,254 @@
 //! IO facilities for TCP relay
 
 use std::{
-    io::{self, BufRead, Read},
-    mem,
-    time::{Duration, Instant},
+    io,
+    marker::{PhantomData, Unpin},
+    pin::Pin,
+    sync::Arc,
+    task::{Context, Poll},
 };
 
-use bytes::{BufMut, BytesMut};
-use futures::{Async, Future, Poll};
-use tokio::timer::Delay;
-use tokio_io::{
-    io::{copy, Copy},
-    try_nb, AsyncRead, AsyncWrite,
-};
+use byte_string::ByteStr;
+use bytes::Bytes;
+use futures::ready;
+use log::trace;
+use tokio::prelude::*;
 
 use super::{
-    utils::{copy_timeout, copy_timeout_opt, CopyTimeout, CopyTimeoutOpt},
-    BUFFER_SIZE,
+    aead::{DecryptedReader as AeadDecryptedReader, EncryptedWriter as AeadEncryptedWriter},
+    stream::{DecryptedReader as StreamDecryptedReader, EncryptedWriter as StreamEncryptedWriter},
 };
 
-static DUMMY_BUFFER: [u8; BUFFER_SIZE] = [0u8; BUFFER_SIZE];
+use crate::{config::ServerConfig, crypto::CipherCategory};
 
-/// Reader to read data from ShadowSocks protocol
-///
-/// This trait requires `BufRead`, because obviously this reader has to contain a buffer inside,
-/// which stores the decrypted data.
-pub trait DecryptedRead: BufRead + AsyncRead {
-    /// Decrypt buffer size
-    fn buffer_size(&self, data: &[u8]) -> usize;
-
-    /// Copies all data to `w`
-    fn copy<W>(self, w: W) -> Copy<Self, W>
-    where
-        Self: Sized,
-        W: AsyncWrite,
-    {
-        copy(self, w)
-    }
-
-    /// Copies all data to `w`, return `TimedOut` if timeout reaches
-    fn copy_timeout<W>(self, w: W, timeout: Duration) -> CopyTimeout<Self, W>
-    where
-        Self: Sized,
-        W: AsyncWrite,
-    {
-        copy_timeout(self, w, timeout)
-    }
-
-    /// The same as `copy_timeout`, but has optional `timeout`
-    fn copy_timeout_opt<W>(self, w: W, timeout: Option<Duration>) -> CopyTimeoutOpt<Self, W>
-    where
-        Self: Sized,
-        W: AsyncWrite,
-    {
-        copy_timeout_opt(self, w, timeout)
-    }
+enum DecryptedReader {
+    Aead(AeadDecryptedReader),
+    Stream(StreamDecryptedReader),
 }
 
-/// Writer that encrypt data and write it as ShadowSocks protocol
-///
-/// The writer cannot implement `io::Write`, because you cannot ensure that you can write all the data in a batch
-/// in non-blocking I/O environment.
-pub trait EncryptedWrite {
-    /// Writes raw bytes directly to the writer
-    fn write_raw(&mut self, data: &[u8]) -> io::Result<usize>;
-    /// Flush the writer
-    fn flush(&mut self) -> io::Result<()>;
-    /// Encrypt data into buffer for writing
-    fn encrypt<B: BufMut>(&mut self, data: &[u8], buf: &mut B) -> io::Result<()>;
-    /// Encrypt buffer size
-    fn buffer_size(&self, data: &[u8]) -> usize;
-
-    /// Encrypt data in `buf` and write all to the writer
-    fn write_all<B: AsRef<[u8]>>(self, buf: B) -> EncryptedWriteAll<Self, B>
-    where
-        Self: Sized,
-    {
-        EncryptedWriteAll::new(self, buf)
-    }
-
-    /// Copies all data from `r`
-    fn copy<R: Read>(self, r: R) -> EncryptedCopy<R, Self>
-    where
-        Self: Sized,
-    {
-        EncryptedCopy::new(r, self)
-    }
-
-    /// Copies all data from `r` with timeout
-    fn copy_timeout<R: Read>(self, r: R, timeout: Duration) -> EncryptedCopyTimeout<R, Self>
-    where
-        Self: Sized,
-    {
-        EncryptedCopyTimeout::new(r, self, timeout)
-    }
-
-    /// Copies all data from `r` with optional timeout
-    fn copy_timeout_opt<R: Read>(self, r: R, timeout: Option<Duration>) -> EncryptedCopyOpt<R, Self>
-    where
-        Self: Sized,
-    {
-        match timeout {
-            Some(t) => EncryptedCopyOpt::CopyTimeout(self.copy_timeout(r, t)),
-            None => EncryptedCopyOpt::Copy(self.copy(r)),
-        }
-    }
+enum EncryptedWriter {
+    Aead(AeadEncryptedWriter),
+    Stream(StreamEncryptedWriter),
 }
 
-/// Write all data encrypted
-pub enum EncryptedWriteAll<W, B>
-where
-    W: EncryptedWrite,
-    B: AsRef<[u8]>,
-{
-    Writing {
-        writer: W,
-        buf: B,
-        pos: usize,
-        enc_buf: BytesMut,
-        encrypted: bool,
-    },
-    Empty,
+enum ReadStatus {
+    WaitIv(Vec<u8>, usize),
+    Established,
 }
 
-impl<W, B> EncryptedWriteAll<W, B>
-where
-    W: EncryptedWrite,
-    B: AsRef<[u8]>,
-{
-    fn new(w: W, buf: B) -> EncryptedWriteAll<W, B> {
-        let buffer_size = w.buffer_size(&DUMMY_BUFFER);
-        EncryptedWriteAll::Writing {
-            writer: w,
-            buf,
-            pos: 0,
-            enc_buf: BytesMut::with_capacity(buffer_size),
-            encrypted: false,
-        }
-    }
+enum WriteStatus {
+    SendIv(Bytes, usize),
+    Established,
 }
 
-impl<W, B> Future for EncryptedWriteAll<W, B>
-where
-    W: EncryptedWrite,
-    B: AsRef<[u8]>,
-{
-    type Error = io::Error;
-    type Item = (W, B);
+pub struct CryptoStream<S> {
+    stream: S,
+    dec: Option<DecryptedReader>,
+    enc: Option<EncryptedWriter>,
+    svr_cfg: Arc<ServerConfig>,
+    read_status: ReadStatus,
+    write_status: WriteStatus,
+}
 
-    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
-        match *self {
-            EncryptedWriteAll::Empty => panic!("poll after EncryptedWriteAll finished"),
-            EncryptedWriteAll::Writing {
-                ref mut writer,
-                ref buf,
-                ref mut pos,
-                ref mut enc_buf,
-                ref mut encrypted,
-            } => {
-                if !*encrypted {
-                    *encrypted = true;
+impl<S: Unpin> Unpin for CryptoStream<S> {}
 
-                    // Ensure buffer has enough space
-                    let buffer_len = writer.buffer_size(buf.as_ref());
-                    enc_buf.reserve(buffer_len);
+impl<S> CryptoStream<S> {
+    pub fn new(stream: S, svr_cfg: Arc<ServerConfig>) -> CryptoStream<S> {
+        let method = svr_cfg.method();
+        let prev_len = match method.category() {
+            CipherCategory::Stream => method.iv_size(),
+            CipherCategory::Aead => method.salt_size(),
+        };
 
-                    writer.encrypt(buf.as_ref(), enc_buf)?;
-                }
-
-                while *pos < enc_buf.len() {
-                    let n = try_nb!(writer.write_raw(&enc_buf[*pos..]));
-                    *pos += n;
-                    if n == 0 {
-                        let err = io::Error::new(io::ErrorKind::Other, "zero-length write");
-                        return Err(err);
-                    }
-                }
+        let local_iv = match method.category() {
+            CipherCategory::Stream => {
+                let local_iv = method.gen_init_vec();
+                trace!("Generated Stream cipher IV {:?}", local_iv);
+                local_iv
             }
-        }
+            CipherCategory::Aead => {
+                let local_salt = method.gen_salt();
+                trace!("Generated AEAD cipher salt {:?}", local_salt);
+                local_salt
+            }
+        };
 
-        match mem::replace(self, EncryptedWriteAll::Empty) {
-            EncryptedWriteAll::Writing { writer, buf, .. } => Ok((writer, buf).into()),
-            EncryptedWriteAll::Empty => unreachable!(),
+        CryptoStream {
+            stream,
+            dec: None,
+            enc: None,
+            svr_cfg,
+            read_status: ReadStatus::WaitIv(vec![0u8; prev_len], 0usize),
+            write_status: WriteStatus::SendIv(local_iv, 0usize),
         }
     }
 }
 
-/// Encrypted copy
-pub struct EncryptedCopy<R, W>
+impl<S> CryptoStream<S>
 where
-    R: Read,
-    W: EncryptedWrite,
+    S: AsyncRead + Unpin,
 {
-    reader: Option<R>,
-    writer: Option<W>,
-    read_done: bool,
-    amt: u64,
-    pos: usize,
-    cap: usize,
-    buf: BytesMut,
-}
-
-impl<R, W> EncryptedCopy<R, W>
-where
-    R: Read,
-    W: EncryptedWrite,
-{
-    fn new(r: R, w: W) -> EncryptedCopy<R, W> {
-        let buffer_size = w.buffer_size(&DUMMY_BUFFER);
-        EncryptedCopy {
-            reader: Some(r),
-            writer: Some(w),
-            read_done: false,
-            amt: 0,
-            pos: 0,
-            cap: 0,
-            buf: BytesMut::with_capacity(buffer_size),
-        }
-    }
-}
-
-impl<R, W> Future for EncryptedCopy<R, W>
-where
-    R: Read,
-    W: EncryptedWrite,
-{
-    type Error = io::Error;
-    type Item = (u64, R, W);
-
-    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
-        let mut local_buf = [0u8; BUFFER_SIZE];
-        loop {
-            // If our buffer is empty, then we need to read some data to
-            // continue.
-            if self.pos == self.cap && !self.read_done {
-                let n = try_nb!(self.reader.as_mut().unwrap().read(&mut local_buf[..]));
-                self.buf.clear();
+    fn poll_read_handshake(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        if let ReadStatus::WaitIv(ref mut buf, ref mut pos) = self.read_status {
+            while *pos < buf.len() {
+                let n = ready!(Pin::new(&mut self.stream).poll_read(cx, &mut buf[*pos..]))?;
                 if n == 0 {
-                    self.read_done = true;
-                } else {
-                    let data = &local_buf[..n];
-                    // Ensure we have enough space
-                    let buffer_len = self.writer.as_mut().unwrap().buffer_size(data);
-                    self.buf.reserve(buffer_len);
-
-                    self.writer.as_mut().unwrap().encrypt(data, &mut self.buf)?;
+                    use std::io::ErrorKind;
+                    return Poll::Ready(Err(ErrorKind::UnexpectedEof.into()));
                 }
-                self.pos = 0;
-                self.cap = self.buf.len();
+                *pos += n;
             }
 
-            // If our buffer has some data, let's write it out!
-            while self.pos < self.cap {
-                let i = try_nb!(self.writer.as_mut().unwrap().write_raw(&self.buf[self.pos..self.cap]));
-                self.pos += i;
-                self.amt += i as u64;
-            }
+            let method = self.svr_cfg.method();
+            let dec = match method.category() {
+                CipherCategory::Stream => {
+                    trace!("Got Stream cipher IV {:?}", ByteStr::new(&buf));
+                    DecryptedReader::Stream(StreamDecryptedReader::new(method, self.svr_cfg.key(), &buf))
+                }
+                CipherCategory::Aead => {
+                    trace!("Got AEAD cipher salt {:?}", ByteStr::new(&buf));
+                    DecryptedReader::Aead(AeadDecryptedReader::new(method, self.svr_cfg.key(), &buf))
+                }
+            };
 
-            // If we've written al the data and we've seen EOF, flush out the
-            // data and finish the transfer.
-            // done with the entire transfer.
-            if self.pos == self.cap && self.read_done {
-                try_nb!(self.writer.as_mut().unwrap().flush());
-                return Ok((self.amt, self.reader.take().unwrap(), self.writer.take().unwrap()).into());
-            }
+            self.dec = Some(dec);
+            self.read_status = ReadStatus::Established;
+        }
+
+        Poll::Ready(Ok(()))
+    }
+
+    fn priv_poll_read(mut self: Pin<&mut Self>, ctx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
+        ready!(self.poll_read_handshake(ctx))?;
+
+        let stream = unsafe { &mut *(&mut self.stream as *mut _) };
+        match *self.dec.as_mut().unwrap() {
+            DecryptedReader::Aead(ref mut r) => r.poll_read_decrypted(ctx, stream, buf),
+            DecryptedReader::Stream(ref mut r) => r.poll_read_decrypted(ctx, stream, buf),
         }
     }
 }
 
-/// Encrypted copy
-pub struct EncryptedCopyTimeout<R, W>
+impl<S> CryptoStream<S>
 where
-    R: Read,
-    W: EncryptedWrite,
+    S: AsyncWrite + Unpin,
 {
-    reader: Option<R>,
-    writer: Option<W>,
-    read_done: bool,
-    amt: u64,
-    pos: usize,
-    cap: usize,
-    timeout: Duration,
-    timer: Option<Delay>,
-    read_buf: [u8; BUFFER_SIZE],
-    write_buf: BytesMut,
-}
-
-impl<R, W> EncryptedCopyTimeout<R, W>
-where
-    R: Read,
-    W: EncryptedWrite,
-{
-    fn new(r: R, w: W, dur: Duration) -> EncryptedCopyTimeout<R, W> {
-        let buffer_size = w.buffer_size(&DUMMY_BUFFER);
-        EncryptedCopyTimeout {
-            reader: Some(r),
-            writer: Some(w),
-            read_done: false,
-            amt: 0,
-            pos: 0,
-            cap: 0,
-            timeout: dur,
-            timer: None,
-            read_buf: [0u8; BUFFER_SIZE],
-            write_buf: BytesMut::with_capacity(buffer_size),
-        }
-    }
-
-    fn try_poll_timeout(&mut self) -> io::Result<()> {
-        match self.timer.as_mut() {
-            None => Ok(()),
-            Some(t) => match t.poll() {
-                Err(err) => panic!("Failed to poll on timer, err: {}", err),
-                Ok(Async::Ready(..)) => Err(io::Error::new(io::ErrorKind::TimedOut, "connection timed out")),
-                Ok(Async::NotReady) => Ok(()),
-            },
-        }
-    }
-
-    fn clear_timer(&mut self) {
-        let _ = self.timer.take();
-    }
-
-    fn read_or_set_timeout(&mut self) -> io::Result<usize> {
-        // First, return if timeout
-        self.try_poll_timeout()?;
-
-        // Then, unset the previous timeout
-        self.clear_timer();
-
-        self.write_buf.clear();
-        match self.reader.as_mut().unwrap().read(&mut self.read_buf) {
-            Ok(0) => {
-                self.cap = 0;
-                self.pos = 0;
-                Ok(0)
-            }
-            Ok(n) => {
-                let data = &self.read_buf[..n];
-                // Ensoure we have enough space
-                let buffer_len = self.writer.as_mut().unwrap().buffer_size(data);
-                self.write_buf.reserve(buffer_len);
-
-                self.writer.as_mut().unwrap().encrypt(data, &mut self.write_buf)?;
-                self.cap = self.write_buf.len();
-                self.pos = 0;
-                Ok(n)
-            }
-            Err(e) => {
-                if e.kind() == io::ErrorKind::WouldBlock {
-                    self.timer = Some(Delay::new(Instant::now() + self.timeout));
-                }
-                Err(e)
-            }
-        }
-    }
-
-    fn write_or_set_timeout(&mut self) -> io::Result<usize> {
-        // First, return if timeout
-        self.try_poll_timeout()?;
-
-        // Then, unset the previous timeout
-        self.clear_timer();
-
-        match self
-            .writer
-            .as_mut()
-            .unwrap()
-            .write_raw(&self.write_buf[self.pos..self.cap])
-        {
-            Ok(n) => {
-                self.pos += n;
-                Ok(n)
-            }
-            Err(e) => {
-                if e.kind() == io::ErrorKind::WouldBlock {
-                    self.timer = Some(Delay::new(Instant::now() + self.timeout));
-                }
-                Err(e)
-            }
-        }
-    }
-}
-
-impl<R, W> Future for EncryptedCopyTimeout<R, W>
-where
-    R: Read,
-    W: EncryptedWrite,
-{
-    type Error = io::Error;
-    type Item = (u64, R, W);
-
-    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
-        loop {
-            // If our buffer is empty, then we need to read some data to
-            // continue.
-            if self.pos == self.cap && !self.read_done {
-                let n = try_nb!(self.read_or_set_timeout());
+    fn poll_write_handshake(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        if let WriteStatus::SendIv(ref iv, ref mut pos) = self.write_status {
+            while *pos < iv.len() {
+                let n = ready!(Pin::new(&mut self.stream).poll_write(cx, &iv[*pos..]))?;
                 if n == 0 {
-                    self.read_done = true;
+                    use std::io::ErrorKind;
+                    return Poll::Ready(Err(ErrorKind::UnexpectedEof.into()));
                 }
+                *pos += n;
             }
 
-            // If our buffer has some data, let's write it out!
-            while self.pos < self.cap {
-                let i = try_nb!(self.write_or_set_timeout());
-                if i == 0 {
-                    let err = io::Error::new(io::ErrorKind::WriteZero, "copied zero bytes into writer");
-                    return Err(err);
+            let method = self.svr_cfg.method();
+            let enc = match method.category() {
+                CipherCategory::Stream => {
+                    trace!("Sent Stream cipher IV {:?}", ByteStr::new(&iv));
+                    EncryptedWriter::Stream(StreamEncryptedWriter::new(method, self.svr_cfg.key(), &iv))
                 }
-                self.amt += i as u64;
-            }
+                CipherCategory::Aead => {
+                    trace!("Sent AEAD cipher salt {:?}", ByteStr::new(&iv));
+                    EncryptedWriter::Aead(AeadEncryptedWriter::new(method, self.svr_cfg.key(), &iv))
+                }
+            };
 
-            // If we've written al the data and we've seen EOF, flush out the
-            // data and finish the transfer.
-            // done with the entire transfer.
-            if self.pos == self.cap && self.read_done {
-                try_nb!(self.writer.as_mut().unwrap().flush());
-                return Ok((self.amt, self.reader.take().unwrap(), self.writer.take().unwrap()).into());
-            }
+            self.enc = Some(enc);
+            self.write_status = WriteStatus::Established;
         }
+
+        Poll::Ready(Ok(()))
+    }
+
+    fn priv_poll_write(mut self: Pin<&mut Self>, ctx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
+        ready!(self.poll_write_handshake(ctx))?;
+
+        let stream = unsafe { &mut *(&mut self.stream as *mut _) };
+        match *self.enc.as_mut().unwrap() {
+            EncryptedWriter::Aead(ref mut w) => w.poll_write_encrypted(ctx, stream, buf),
+            EncryptedWriter::Stream(ref mut w) => w.poll_write_encrypted(ctx, stream, buf),
+        }
+    }
+
+    fn priv_poll_flush(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        AsyncWrite::poll_flush(Pin::new(&mut self.stream), ctx)
+    }
+
+    fn priv_poll_shutdown(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        AsyncWrite::poll_shutdown(Pin::new(&mut self.stream), ctx)
     }
 }
 
-/// Work for both timeout or no timeout
-pub enum EncryptedCopyOpt<R, W>
+impl<S> CryptoStream<S>
 where
-    R: Read,
-    W: EncryptedWrite,
+    S: AsyncRead + AsyncWrite + Unpin,
 {
-    Copy(EncryptedCopy<R, W>),
-    CopyTimeout(EncryptedCopyTimeout<R, W>),
+    pub fn split(&mut self) -> (CryptoStreamReadHalf<'_, S>, CryptoStreamWriteHalf<'_, S>) {
+        let p = self as *mut _;
+        (
+            CryptoStreamReadHalf(p, PhantomData),
+            CryptoStreamWriteHalf(p, PhantomData),
+        )
+    }
 }
 
-impl<R, W> Future for EncryptedCopyOpt<R, W>
+impl<S> AsyncRead for CryptoStream<S>
 where
-    R: Read,
-    W: EncryptedWrite,
+    S: AsyncRead + Unpin,
 {
-    type Error = io::Error;
-    type Item = (u64, R, W);
-
-    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
-        match *self {
-            EncryptedCopyOpt::Copy(ref mut c) => c.poll(),
-            EncryptedCopyOpt::CopyTimeout(ref mut c) => c.poll(),
-        }
+    fn poll_read(self: Pin<&mut Self>, ctx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
+        self.priv_poll_read(ctx, buf)
+    }
+}
+
+impl<S> AsyncWrite for CryptoStream<S>
+where
+    S: AsyncWrite + Unpin,
+{
+    fn poll_write(self: Pin<&mut Self>, ctx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
+        self.priv_poll_write(ctx, buf)
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        self.priv_poll_flush(ctx)
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        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

@@ -2,12 +2,13 @@
 
 use std::io;
 
-use futures::Future;
-
-use super::socks5_local;
+use super::{socks5_local, tunnel_local};
 use crate::context::SharedContext;
 
 /// Starts a TCP local server
-pub fn run(context: SharedContext) -> impl Future<Item = (), Error = io::Error> + Send {
-    socks5_local::run(context)
+pub async fn run(context: SharedContext) -> io::Result<()> {
+    match context.config().forward {
+        Some(..) => tunnel_local::run(context).await,
+        None => socks5_local::run(context).await,
+    }
 }

src/relay/tcprelay/mod.rs

@@ -1,40 +1,30 @@
 //! Relay for TCP implementation
 
-use std::{
-    io::{self, BufRead, Read, Write},
-    iter::{IntoIterator, Iterator},
-    mem,
-    net::SocketAddr,
-    sync::Arc,
-    time::Duration,
-};
+// Allow for futures
+// Maybe removed in the future
+#![allow(clippy::unnecessary_mut_passed)]
 
-use crate::{
-    config::{ConfigType, ServerAddr, ServerConfig},
-    context::SharedContext,
-    crypto::CipherCategory,
-    relay::{boxed_future, dns_resolver::resolve, socks5::Address},
-};
+use std::io;
+use std::marker::Unpin;
+use std::ops::{Deref, DerefMut};
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+use std::time::Duration;
 
-use byte_string::ByteStr;
-use bytes::{BufMut, BytesMut};
-use futures::{self, Async, Future, Poll};
-use log::{error, trace};
-use tokio::{
-    net::{tcp::ConnectFuture, TcpStream},
-    timer::Timeout,
-};
-use tokio_io::{
-    io::{read_exact, write_all, ReadHalf, WriteHalf},
-    try_nb, AsyncRead, AsyncWrite,
-};
+use crate::config::{ConfigType, ServerAddr, ServerConfig};
+use crate::context::SharedContext;
+use crate::relay::socks5::Address;
+use crate::relay::utils::try_timeout;
 
-pub use self::crypto_io::{DecryptedRead, EncryptedWrite};
-
-use self::{
-    aead::{DecryptedReader as AeadDecryptedReader, EncryptedWriter as AeadEncryptedWriter},
-    stream::{DecryptedReader as StreamDecryptedReader, EncryptedWriter as StreamEncryptedWriter},
-};
+use bytes::BytesMut;
+use futures::future::FusedFuture;
+use futures::{ready, select, Future};
+use log::trace;
+use tokio::io::{ReadHalf, WriteHalf};
+use tokio::net::TcpStream;
+use tokio::prelude::*;
+use tokio::time::{self, Delay};
 
 mod aead;
 pub mod client;
@@ -45,224 +35,144 @@ mod monitor;
 pub mod server;
 mod socks5_local;
 mod stream;
+mod tunnel_local;
 mod utils;
 
+pub use self::crypto_io::CryptoStream;
+
 const BUFFER_SIZE: usize = 8 * 1024; // 8K buffer
 
-// /// Directions in the tunnel
-// #[derive(Debug, Copy, Clone)]
-// pub enum TunnelDirection {
-//     /// Client -> Server
-//     Client2Server,
-//     /// Client <- Server
-//     Server2Client,
-// }
-
-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<S>
-where
-    S: Read + Write + AsyncRead + AsyncWrite + Send,
-{
-    Stream(StreamDecryptedReader<TcpReadHalf<S>>),
-    Aead(AeadDecryptedReader<TcpReadHalf<S>>),
+/// Shadowsocks' Connection
+///
+/// The only feature: Supports timeout
+pub struct Connection<S> {
+    // Actual connection socket
+    stream: S,
+    // Timer instance
+    // Read and Write operations shares the same timer
+    timer: Option<Delay>,
+    // User defined server timeout
+    timeout: Option<Duration>,
 }
 
-macro_rules! ref_half_do {
-    ($self:expr,$t:ident,$m:ident$(,$p:expr)*) => {
-        match *$self {
-            $t::Stream(ref d) => d.$m($($p),*),
-            $t::Aead(ref d) => d.$m($($p),*),
+impl<S> Connection<S> {
+    pub fn new(stream: S, timeout: Option<Duration>) -> Connection<S> {
+        Connection {
+            stream,
+            timer: None,
+            timeout,
         }
     }
-}
 
-macro_rules! mut_half_do {
-    ($self:expr,$t:ident,$m:ident$(,$p:expr)*) => {
-        match *$self {
-            $t::Stream(ref mut  d) => d.$m($($p),*),
-            $t::Aead(ref mut d) => d.$m($($p),*),
-        }
-    }
-}
-
-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<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<S> AsyncRead for DecryptedHalf<S> where S: Read + Write + AsyncRead + AsyncWrite + Send {}
-
-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)
-    }
-
-    fn consume(&mut self, amt: usize) {
-        mut_half_do!(self, DecryptedHalf, consume, amt)
-    }
-}
-
-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<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<S>
-where
-    S: Read + Write + AsyncRead + AsyncWrite + Send,
-{
-    Stream(StreamEncryptedWriter<TcpWriteHalf<S>>),
-    Aead(AeadEncryptedWriter<TcpWriteHalf<S>>),
-}
-
-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)
-    }
-
-    fn flush(&mut self) -> io::Result<()> {
-        mut_half_do!(self, EncryptedHalf, flush)
-    }
-
-    fn encrypt<B: BufMut>(&mut self, data: &[u8], buf: &mut B) -> io::Result<()> {
-        mut_half_do!(self, EncryptedHalf, encrypt, data, buf)
-    }
-
-    fn buffer_size(&self, data: &[u8]) -> usize {
-        ref_half_do!(self, EncryptedHalf, buffer_size, data)
-    }
-}
-
-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<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)
-    }
-}
-
-/// Try to connect every IPs one by one
-enum TcpStreamConnect<I: Iterator<Item = SocketAddr>> {
-    Empty,
-    Connect {
-        last_err: Option<io::Error>,
-        addr_iter: I,
-        opt_stream_new: Option<(ConnectFuture, SocketAddr)>,
-    },
-}
-
-impl<I: Iterator<Item = SocketAddr>> TcpStreamConnect<I> {
-    fn new(iter: I) -> TcpStreamConnect<I> {
-        TcpStreamConnect::Connect {
-            last_err: None,
-            addr_iter: iter,
-            opt_stream_new: None,
-        }
-    }
-}
-
-impl<I: Iterator<Item = SocketAddr>> Future for TcpStreamConnect<I> {
-    type Error = io::Error;
-    type Item = TcpStream;
-
-    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
+    fn make_timeout_error() -> io::Error {
         use std::io::ErrorKind;
+        ErrorKind::TimedOut.into()
+    }
 
-        match *self {
-            TcpStreamConnect::Empty => unreachable!(),
-            TcpStreamConnect::Connect {
-                ref mut last_err,
-                ref mut addr_iter,
-                ref mut opt_stream_new,
-            } => {
-                loop {
-                    // 1. Poll before doing anything else
-                    if let Some((ref mut stream_new, ref addr)) = *opt_stream_new {
-                        match stream_new.poll() {
-                            Ok(Async::Ready(stream)) => return Ok(Async::Ready(stream)),
-                            Ok(Async::NotReady) => return Ok(Async::NotReady),
-                            Err(ref err) if err.kind() == ErrorKind::WouldBlock => return Ok(Async::NotReady),
-                            Err(err) => {
-                                error!("Failed to connect {}: {}", addr, err);
-
-                                *last_err = Some(err);
-                            }
-                        }
-                    }
-
-                    match addr_iter.next() {
-                        None => break,
-                        Some(addr) => {
-                            *opt_stream_new = Some((TcpStream::connect(&addr), addr));
-                        }
-                    }
+    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(()))
+    }
 
-        match mem::replace(self, TcpStreamConnect::Empty) {
-            TcpStreamConnect::Empty => unreachable!(),
-            TcpStreamConnect::Connect { last_err, .. } => match last_err {
-                None => {
-                    let err = io::Error::new(ErrorKind::Other, "connect TCP without any addresses");
-                    Err(err)
-                }
-                Some(err) => Err(err),
-            },
+    fn cancel_timeout(&mut self) {
+        let _ = self.timer.take();
+    }
+}
+
+impl<S> Connection<S>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+{
+    pub fn split(self) -> (ReadHalf<Connection<S>>, WriteHalf<Connection<S>>) {
+        use tokio::io::split;
+        split(self)
+    }
+}
+
+impl<S> Deref for Connection<S> {
+    type Target = S;
+
+    fn deref(&self) -> &Self::Target {
+        &self.stream
+    }
+}
+
+impl<S> DerefMut for Connection<S> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.stream
+    }
+}
+
+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
+            }
         }
     }
 }
 
-fn connect_proxy_server(
-    context: SharedContext,
-    svr_cfg: Arc<ServerConfig>,
-) -> impl Future<Item = TcpStream, Error = io::Error> + Send {
+impl<S> AsyncWrite for Connection<S>
+where
+    S: AsyncWrite + 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();
+                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)
+    }
+}
+
+pub type STcpStream = Connection<TcpStream>;
+
+/// Connect to proxy server with `ServerConfig`
+async fn connect_proxy_server(context: SharedContext, svr_cfg: Arc<ServerConfig>) -> io::Result<STcpStream> {
     let timeout = svr_cfg.timeout();
 
     let svr_addr = match context.config().config_type {
@@ -273,286 +183,90 @@ fn connect_proxy_server(
     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)
+            let stream = try_timeout(TcpStream::connect(addr), timeout).await?;
+            Ok(STcpStream::new(stream, timeout))
         }
+        #[cfg(feature = "trust-dns")]
         ServerAddr::DomainName(ref domain, port) => {
-            let fut = {
-                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)
-                })
-            };
-            boxed_future(fut)
+            use crate::relay::dns_resolver::resolve;
+            use log::error;
+
+            let vec_ipaddr = try_timeout(resolve(context, &domain[..], *port, false), timeout).await?;
+
+            assert!(!vec_ipaddr.is_empty());
+
+            let mut last_err: Option<io::Error> = None;
+            for addr in &vec_ipaddr {
+                match try_timeout(TcpStream::connect(addr), timeout).await {
+                    Ok(s) => return Ok(STcpStream::new(s, timeout)),
+                    Err(e) => {
+                        error!(
+                            "Failed to connect {}:{}, resolved address {}, try another (err: {})",
+                            domain, port, addr, e
+                        );
+                        last_err = Some(e);
+                    }
+                }
+            }
+
+            let err = last_err.unwrap();
+            error!(
+                "Failed to connect {}:{}, tried all addresses but still failed (last err: {})",
+                domain, port, err
+            );
+            Err(err)
+        }
+        #[cfg(not(feature = "trust-dns"))]
+        ServerAddr::DomainName(ref domain, port) => {
+            let stream = try_timeout(TcpStream::connect((domain.as_str(), *port)), timeout).await?;
+            Ok(STcpStream::new(stream, timeout))
         }
     }
 }
 
 /// Handshake logic for ShadowSocks Client
-pub fn proxy_server_handshake<S>(
-    remote_stream: S,
+pub async fn proxy_server_handshake(
+    remote_stream: STcpStream,
     svr_cfg: Arc<ServerConfig>,
-    relay_addr: Address,
-) -> impl Future<
-    Item = (
-        impl Future<Item = DecryptedHalf<S>, Error = io::Error> + Send,
-        impl Future<Item = EncryptedHalf<S>, Error = io::Error> + Send,
-    ),
-    Error = io::Error,
-> + 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| {
-            // Send relay address to remote
-            let mut buf = BytesMut::with_capacity(relay_addr.serialized_len());
-            relay_addr.write_to_buf(&mut buf);
+    relay_addr: &Address,
+) -> io::Result<CryptoStream<STcpStream>> {
+    let mut stream = CryptoStream::new(remote_stream, svr_cfg.clone());
 
-            trace!(
-                "Got encrypt stream and going to send addr: {:?}, buf: {:?}",
-                relay_addr,
-                buf
-            );
+    trace!("Got encrypt stream and going to send addr: {:?}", relay_addr);
 
-            try_timeout(enc_w.write_all(buf), timeout).map(|(w, _)| w)
-        });
+    // Send relay address to remote
+    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?;
 
-        Ok((r_fut, w_fut))
-    })
-}
-
-/// ShadowSocks Client-Server handshake protocol
-/// Exchange cipher IV and creates stream wrapper
-pub fn proxy_handshake<S>(
-    remote_stream: S,
-    svr_cfg: Arc<ServerConfig>,
-) -> impl Future<
-    Item = (
-        impl Future<Item = DecryptedHalf<S>, Error = io::Error> + Send,
-        impl Future<Item = EncryptedHalf<S>, Error = io::Error> + Send,
-    ),
-    Error = io::Error,
-> + 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();
-
-        let svr_cfg_cloned = svr_cfg.clone();
-
-        let enc = {
-            // Encrypt data to remote server
-
-            // Send initialize vector to remote and create encryptor
-
-            let method = svr_cfg.method();
-            let prev_buf = match method.category() {
-                CipherCategory::Stream => {
-                    let local_iv = method.gen_init_vec();
-                    trace!("Going to send initialize vector: {:?}", local_iv);
-                    local_iv
-                }
-                CipherCategory::Aead => {
-                    let local_salt = method.gen_salt();
-                    trace!("Going to send salt: {:?}", local_salt);
-                    local_salt
-                }
-            };
-
-            try_timeout(write_all(w, prev_buf), timeout).and_then(move |(w, prev_buf)| {
-                match svr_cfg.method().category() {
-                    CipherCategory::Stream => {
-                        let local_iv = prev_buf;
-                        let wtr = StreamEncryptedWriter::new(w, svr_cfg.method(), svr_cfg.key(), &local_iv);
-                        Ok(From::from(wtr))
-                    }
-                    CipherCategory::Aead => {
-                        let local_salt = prev_buf;
-                        let wtr = AeadEncryptedWriter::new(w, svr_cfg.method(), svr_cfg.key(), &local_salt[..]);
-                        Ok(From::from(wtr))
-                    }
-                }
-            })
-        };
-
-        let dec = {
-            let svr_cfg = svr_cfg_cloned;
-
-            // Decrypt data from remote server
-
-            let method = svr_cfg.method();
-            let prev_len = match method.category() {
-                CipherCategory::Stream => method.iv_size(),
-                CipherCategory::Aead => method.salt_size(),
-            };
-
-            try_timeout(read_exact(r, vec![0u8; prev_len]), timeout).and_then(move |(r, remote_iv)| {
-                match svr_cfg.method().category() {
-                    CipherCategory::Stream => {
-                        trace!("Got initialize vector {:?}", ByteStr::new(&remote_iv));
-                        let decrypt_stream = StreamDecryptedReader::new(r, svr_cfg.method(), svr_cfg.key(), &remote_iv);
-                        Ok(From::from(decrypt_stream))
-                    }
-                    CipherCategory::Aead => {
-                        trace!("Got salt {:?}", ByteStr::new(&remote_iv));
-                        let dr = AeadDecryptedReader::new(r, svr_cfg.method(), svr_cfg.key(), &remote_iv);
-                        Ok(From::from(dr))
-                    }
-                }
-            })
-        };
-
-        Ok((dec, enc))
-    })
+    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 fn tunnel<CF, CFI, SF, SFI>(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<()>
 where
-    CF: Future<Item = CFI, Error = io::Error> + Send + 'static,
-    SF: Future<Item = SFI, Error = io::Error> + Send + 'static,
+    CF: Future<Output = io::Result<CFI>> + Unpin + FusedFuture,
+    SF: Future<Output = io::Result<SFI>> + Unpin + FusedFuture,
 {
-    c2s.map(|_| ()).select(s2c.map(|_| ())).then(|r| match r {
-        Ok(..) => Ok(()),
-        Err((err, _)) => Err(err),
-    })
-
-    // let addr = Arc::new(addr);
-
-    // let cloned_addr = addr.clone();
-    // let c2s = c2s.then(move |res| {
-    //     match res {
-    //         Ok(..) => {
-    //             // Continue reading response from remote server
-    //             trace!("Relay {} client -> server is finished", cloned_addr);
-
-    //             Ok(TunnelDirection::Client2Server)
-    //         }
-    //         Err(err) => {
-    //             error!("Relay {} client -> server aborted: {}", cloned_addr, err);
-    //             Err(err)
-    //         }
-    //     }
-    // });
-
-    // let cloned_addr = addr.clone();
-    // let s2c = s2c.then(move |res| match res {
-    //     Ok(..) => {
-    //         trace!("Relay {} client <- server is finished", cloned_addr);
-
-    //         Ok(TunnelDirection::Server2Client)
-    //     }
-    //     Err(err) => {
-    //         error!("Relay {} client <- server aborted: {}", cloned_addr, err);
-    //         Err(err)
-    //     }
-    // });
-
-    // c2s.select(s2c)
-    //     .and_then(move |(dir, _)| {
-    //         match dir {
-    //             TunnelDirection::Server2Client => trace!("Relay {} client <- server is closed, abort connection", addr),
-    //             TunnelDirection::Client2Server => trace!("Relay {} server -> client is closed, abort connection", addr),
-    //         }
-
-    //         Ok(())
-    //     })
-    //     .map_err(|(err, _)| err)
-}
-
-/// Read until EOF, and ignore
-pub enum IgnoreUntilEnd<R: Read> {
-    Pending { r: R, amt: u64 },
-    Empty,
-}
-
-impl<R: Read> Future for IgnoreUntilEnd<R> {
-    type Error = io::Error;
-    type Item = u64;
-
-    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
-        match *self {
-            IgnoreUntilEnd::Empty => panic!("poll IgnoreUntilEnd after it is finished"),
-            IgnoreUntilEnd::Pending { ref mut r, ref mut amt } => {
-                let mut buf = [0u8; 4096];
-                loop {
-                    let n = try_nb!(r.read(&mut buf));
-                    *amt += n as u64;
-
-                    if n == 0 {
-                        break;
-                    }
-                }
-            }
-        }
-
-        match mem::replace(self, IgnoreUntilEnd::Empty) {
-            IgnoreUntilEnd::Pending { amt, .. } => Ok(amt.into()),
-            IgnoreUntilEnd::Empty => unreachable!(),
-        }
+    select! {
+        r1 = c2s => r1.map(|_| ()),
+        r2 = s2c => r2.map(|_| ()),
     }
 }
 
-/// Ignore all data from the reader
-pub fn ignore_until_end<R: Read>(r: R) -> IgnoreUntilEnd<R> {
-    IgnoreUntilEnd::Pending { r, amt: 0 }
-}
-
-pub enum TimeoutFuture<T, F>
+pub async fn ignore_until_end<R>(r: &mut R) -> io::Result<u64>
 where
-    T: 'static,
-    F: Future<Item = T, Error = io::Error> + Send + 'static,
+    R: AsyncRead + Unpin,
 {
-    Direct(F),
-    Wait(Timeout<F>),
-}
-
-impl<T, F> Future for TimeoutFuture<T, F>
-where
-    T: 'static,
-    F: Future<Item = T, Error = io::Error> + Send + 'static,
-{
-    type Error = io::Error;
-    type Item = T;
-
-    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
-        match *self {
-            TimeoutFuture::Direct(ref mut f) => f.poll(),
-            TimeoutFuture::Wait(ref mut f) => match f.poll() {
-                Ok(Async::NotReady) => Ok(Async::NotReady),
-                Ok(Async::Ready(x)) => Ok(Async::Ready(x)),
-                Err(err) => match err.into_inner() {
-                    Some(e) => Err(e),
-                    None => Err(io::Error::new(io::ErrorKind::TimedOut, "connection timed out")),
-                },
-            },
+    let mut buf = [0u8; BUFFER_SIZE];
+    let mut amt = 0u64;
+    loop {
+        let n = r.read(&mut buf).await?;
+        if n == 0 {
+            break;
         }
+        amt += n as u64;
     }
+    Ok(amt)
 }
-
-fn try_timeout<T, F>(fut: F, dur: Option<Duration>) -> impl Future<Item = T, Error = io::Error> + Send
-where
-    F: Future<Item = T, Error = io::Error> + Send + 'static,
-    T: 'static,
-{
-    use tokio::prelude::*;
-
-    match dur {
-        Some(dur) => TimeoutFuture::Wait(fut.timeout(dur)),
-        None => TimeoutFuture::Direct(fut),
-    }
-}
-
-// fn io_timeout<T, F>(fut: F, dur: Duration) -> impl Future<Item = T, Error = io::Error> + Send
-// where
-//     F: Future<Item = T, Error = io::Error> + Send + 'static,
-//     T: 'static,
-// {
-//     use tokio::prelude::*;
-
-//     TimeoutFuture::Wait(fut.timeout(dur))
-// }

src/relay/tcprelay/monitor.rs

@@ -1,67 +1,74 @@
 //! Server traffic monitor
 
 use std::{
-    io::{self, Read, Write},
+    io,
+    marker::Unpin,
     ops::{Deref, DerefMut},
+    pin::Pin,
+    task::{Context, Poll},
 };
 
-use tokio::{
-    io::{AsyncRead, AsyncWrite},
-    net::TcpStream,
-    prelude::Async,
-};
+use tokio::io::{AsyncRead, AsyncWrite};
 
 use super::context::SharedTcpServerContext;
 
-pub struct TcpMonStream {
-    stream: TcpStream,
+pub struct TcpMonStream<S> {
+    stream: S,
     context: SharedTcpServerContext,
 }
 
-impl TcpMonStream {
-    pub fn new(c: SharedTcpServerContext, s: TcpStream) -> TcpMonStream {
+impl<S> TcpMonStream<S> {
+    pub fn new(c: SharedTcpServerContext, s: S) -> TcpMonStream<S> {
         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)?;
+impl<S> AsyncRead for TcpMonStream<S>
+where
+    S: AsyncRead + Unpin,
+{
+    fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
+        let n = match Pin::new(&mut self.stream).poll_read(cx, buf)? {
+            Poll::Ready(n) => n,
+            Poll::Pending => return Poll::Pending,
+        };
         self.context.incr_rx(n);
-        Ok(n)
+        Poll::Ready(Ok(n))
     }
 }
 
-impl Write for TcpMonStream {
-    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        let n = self.stream.write(buf)?;
+impl<S> AsyncWrite for TcpMonStream<S>
+where
+    S: AsyncWrite + Unpin,
+{
+    fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
+        let n = match Pin::new(&mut self.stream).poll_write(cx, buf)? {
+            Poll::Ready(n) => n,
+            Poll::Pending => return Poll::Pending,
+        };
         self.context.incr_tx(n);
-        Ok(n)
+        Poll::Ready(Ok(n))
     }
 
-    fn flush(&mut self) -> io::Result<()> {
-        self.stream.flush()
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.stream).poll_flush(cx)
+    }
+
+    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.stream).poll_shutdown(cx)
     }
 }
 
-impl AsyncRead for TcpMonStream {}
+impl<S> Deref for TcpMonStream<S> {
+    type Target = S;
 
-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 {
+    fn deref(&self) -> &Self::Target {
         &self.stream
     }
 }
 
-impl DerefMut for TcpMonStream {
-    fn deref_mut(&mut self) -> &mut TcpStream {
+impl<S> DerefMut for TcpMonStream<S> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
         &mut self.stream
     }
 }

src/relay/tcprelay/server.rs

@@ -1,197 +1,33 @@
 //! Relay for TCP server that running on the server side
 
-use std::{
-    io::{self, ErrorKind},
-    net::SocketAddr,
-    sync::Arc,
-    time::Duration,
-};
+use std::{io, net::SocketAddr, sync::Arc};
 
-use crate::relay::{
-    boxed_future,
-    dns_resolver::resolve,
-    socks5::Address,
-    tcprelay::crypto_io::{DecryptedRead, EncryptedWrite},
-};
+use crate::relay::socks5::Address;
 
 use crate::context::SharedContext;
 
 use futures::{
-    self,
-    stream::{futures_unordered, Stream},
-    Future,
+    future::{self, Either},
+    stream::{FuturesUnordered, StreamExt},
 };
 use log::{debug, error, info, trace};
 use tokio::{
     self,
     net::{TcpListener, TcpStream},
 };
-use tokio_io::{
-    io::{ReadHalf, WriteHalf},
-    AsyncRead,
-};
 
 use super::{
     context::{SharedTcpServerContext, TcpServerContext},
     monitor::TcpMonStream,
-    proxy_handshake, try_timeout, tunnel, DecryptedHalf, EncryptedHalf, TcpStreamConnect,
+    CryptoStream, STcpStream,
 };
 
-/// Context for doing handshake with client
-pub struct TcpRelayClientHandshake {
-    s: TcpMonStream,
+#[allow(clippy::cognitive_complexity)]
+async fn handle_client(
     svr_context: SharedTcpServerContext,
-}
-
-impl TcpRelayClientHandshake {
-    #[inline]
-    fn error_handshake(peer_addr: SocketAddr) -> io::Error {
-        io::Error::new(
-            ErrorKind::Other,
-            format!(
-                "failed to decode Address, may be wrong method or key, peer: {}",
-                peer_addr
-            ),
-        )
-    }
-
-    /// Doing handshake with client
-    pub fn handshake(
-        self,
-    ) -> impl Future<
-        Item = TcpRelayClientPending<
-            impl Future<Item = EncryptedHalf<TcpMonStream>, Error = io::Error> + Send + 'static,
-        >,
-        Error = io::Error,
-    > + Send {
-        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_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 =
-                            Address::read_from(r).map_err(move |_| TcpRelayClientHandshake::error_handshake(peer_addr));
-                        try_timeout(fut, timeout)
-                    })
-                    .map(move |(r, addr)| TcpRelayClientPending {
-                        r,
-                        addr,
-                        w: w_fut,
-                        timeout,
-                        svr_context,
-                    })
-            })
-        })
-    }
-}
-
-/// Context for connecting remote
-pub struct TcpRelayClientPending<E>
-where
-    E: Future<Item = EncryptedHalf<TcpMonStream>, Error = io::Error> + Send + 'static,
-{
-    r: DecryptedHalf<TcpMonStream>,
-    addr: Address,
-    w: E,
-    timeout: Option<Duration>,
-    svr_context: SharedTcpServerContext,
-}
-
-/// Connect to the remote server
-#[inline]
-fn connect_remote(
-    context: SharedContext,
-    addr: Address,
-    timeout: Option<Duration>,
-) -> impl Future<Item = TcpStream, Error = io::Error> + Send {
-    debug!("Connecting to remote {}", addr);
-
-    match addr {
-        Address::SocketAddress(saddr) => {
-            if context.config().forbidden_ip.contains(&saddr.ip()) {
-                let err = io::Error::new(
-                    ErrorKind::Other,
-                    format!("{} is forbidden, failed to connect {}", saddr.ip(), saddr),
-                );
-                return boxed_future(futures::done(Err(err)));
-            }
-
-            let conn = TcpStream::connect(&saddr);
-            let fut = try_timeout(conn, timeout);
-            boxed_future(fut)
-        }
-        Address::DomainNameAddress(dname, port) => {
-            let fut = {
-                try_timeout(resolve(context, dname.as_str(), port, true), timeout).and_then(move |addrs| {
-                    let conn = TcpStreamConnect::new(addrs.into_iter());
-                    try_timeout(conn, timeout)
-                })
-            };
-            boxed_future(fut)
-        }
-    }
-}
-
-impl<E> TcpRelayClientPending<E>
-where
-    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<TcpMonStream>, Error = io::Error> + Send + 'static,
-        >,
-        Error = io::Error,
-    > + Send {
-        let client_pair = (self.r, self.w);
-        let timeout = self.timeout;
-        connect_remote(self.svr_context.context().clone(), self.addr, self.timeout).map(move |stream| {
-            TcpRelayClientConnected {
-                server: stream.split(),
-                client: client_pair,
-                timeout,
-            }
-        })
-    }
-}
-
-/// Context for extablishing tunnel
-pub struct TcpRelayClientConnected<E>
-where
-    E: Future<Item = EncryptedHalf<TcpMonStream>, Error = io::Error> + Send + 'static,
-{
-    server: (ReadHalf<TcpStream>, WriteHalf<TcpStream>),
-    client: (DecryptedHalf<TcpMonStream>, E),
-    timeout: Option<Duration>,
-}
-
-impl<E> TcpRelayClientConnected<E>
-where
-    E: Future<Item = EncryptedHalf<TcpMonStream>, Error = io::Error> + Send + 'static,
-{
-    /// Establish tunnel
-    #[inline]
-    pub fn tunnel(self) -> impl Future<Item = (), Error = io::Error> + Send {
-        let (svr_r, svr_w) = self.server;
-        let (r, w_fut) = self.client;
-        let timeout = self.timeout;
-
-        tunnel(
-            r.copy_timeout_opt(svr_w, self.timeout),
-            w_fut.and_then(move |w| w.copy_timeout_opt(svr_r, timeout)),
-        )
-    }
-}
-
-fn handle_client(svr_context: SharedTcpServerContext, socket: TcpStream) -> impl Future<Item = (), Error = ()> + Send {
-    let socket = TcpMonStream::new(svr_context.clone(), socket);
-
+    socket: TcpStream,
+    peer_addr: SocketAddr,
+) -> io::Result<()> {
     if let Err(err) = socket.set_keepalive(svr_context.svr_cfg().timeout()) {
         error!("Failed to set keep alive: {:?}", err);
     }
@@ -202,39 +38,167 @@ fn handle_client(svr_context: SharedTcpServerContext, socket: TcpStream) -> impl
         }
     }
 
-    futures::lazy(move || match socket.peer_addr() {
-        Ok(addr) => Ok((socket, addr)),
+    trace!(
+        "Got connection addr: {} with proxy server: {:?}",
+        peer_addr,
+        svr_context.svr_cfg()
+    );
+
+    let stream = TcpMonStream::new(
+        svr_context.clone(),
+        STcpStream::new(socket, svr_context.svr_cfg().timeout()),
+    );
+
+    // Do server-client handshake
+    // Perform encryption IV exchange
+    let mut stream = CryptoStream::new(stream, svr_context.svr_cfg().clone());
+
+    // Read remote Address
+    let remote_addr = match Address::read_from(&mut stream).await {
+        Ok(o) => o,
         Err(err) => {
-            error!("Failed to get peer_addr after accept: {}", err);
-            Err(())
+            error!(
+                "Failed to decode Address, may be wrong method or key, peer {}",
+                peer_addr
+            );
+            return Err(From::from(err));
         }
-    })
-    .and_then(move |(socket, addr)| {
-        trace!("Got connection, addr: {}", addr);
-        trace!("Picked proxy server: {:?}", svr_context.svr_cfg());
+    };
 
-        let client = TcpRelayClientHandshake { s: socket, svr_context };
+    debug!("Relay {} <-> {} establishing", peer_addr, remote_addr);
 
-        client
-            .handshake()
-            .and_then(TcpRelayClientPending::connect)
-            .and_then(TcpRelayClientConnected::tunnel)
-            .map_err(move |err| {
-                error!("Failed to handle client ({}): {}", addr, err);
-            })
-    })
+    let context = svr_context.context();
+
+    let mut remote_stream = match remote_addr {
+        Address::SocketAddress(ref saddr) => {
+            if context.config().forbidden_ip.contains(&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);
+            }
+
+            match TcpStream::connect(saddr).await {
+                Ok(s) => {
+                    debug!("Connected to remote {}", saddr);
+                    s
+                }
+                Err(err) => {
+                    error!("Failed to connect remote {}, {}", saddr, err);
+                    return Err(err);
+                }
+            }
+        }
+        #[cfg(feature = "trust-dns")]
+        Address::DomainNameAddress(ref dname, port) => {
+            use crate::relay::dns_resolver::resolve;
+
+            let addrs = match resolve(context.clone(), dname.as_str(), port, true).await {
+                Ok(r) => r,
+                Err(err) => {
+                    error!("Failed to resolve {}, {}", dname, err);
+                    return Err(err);
+                }
+            };
+
+            let mut last_err: Option<io::Error> = None;
+            let mut stream_opt = None;
+            for addr in &addrs {
+                match TcpStream::connect(addr).await {
+                    Ok(s) => stream_opt = Some(s),
+                    Err(err) => {
+                        error!(
+                            "Failed to connect remote {}:{} (resolved: {}), {}, try others",
+                            dname, port, addr, err
+                        );
+                        last_err = Some(err);
+                    }
+                }
+            }
+
+            match stream_opt {
+                Some(s) => {
+                    debug!("Connected to remote {}:{}", dname, port);
+                    s
+                }
+                None => {
+                    let err = last_err.unwrap();
+                    error!("Failed to connect remote {}:{}, {}", dname, port, err);
+                    return Err(io::Error::new(io::ErrorKind::Other, err));
+                }
+            }
+        }
+        #[cfg(not(feature = "trust-dns"))]
+        Address::DomainNameAddress(ref dname, port) => {
+            let s = match TcpStream::connect((dname.as_str(), port)).await {
+                Ok(s) => {
+                    debug!("Connected to remote {}:{}", dname, port);
+                    s
+                }
+                Err(err) => {
+                    error!("Failed to connect remote {}:{}, {}", dname, port, err);
+                    return Err(err);
+                }
+            };
+
+            // Still need to check forbidden IPs
+            let forbidden_ip = &context.config().forbidden_ip;
+            if !forbidden_ip.is_empty() {
+                let peer_addr = s.peer_addr()?;
+
+                if forbidden_ip.contains(&peer_addr.ip()) {
+                    error!("{} is forbidden, failed to connect {}", peer_addr.ip(), peer_addr);
+                    let err = io::Error::new(
+                        io::ErrorKind::Other,
+                        format!("{} is forbidden, failed to connect {}", peer_addr.ip(), peer_addr),
+                    );
+                    return Err(err);
+                }
+            }
+
+            s
+        }
+    };
+
+    debug!("Relay {} <-> {} established", peer_addr, remote_addr);
+
+    let (mut cr, mut cw) = stream.split();
+    let (mut sr, mut sw) = remote_stream.split();
+
+    use tokio::io::copy;
+
+    // CLIENT -> SERVER
+    let rhalf = copy(&mut cr, &mut sw);
+
+    // CLIENT <- SERVER
+    let whalf = copy(&mut sr, &mut cw);
+
+    match future::select(rhalf, whalf).await {
+        Either::Left((Ok(_), _)) => trace!("Relay {} -> {} closed", peer_addr, remote_addr),
+        Either::Left((Err(err), _)) => trace!("Relay {} -> {} closed with error {:?}", peer_addr, remote_addr, err),
+        Either::Right((Ok(_), _)) => trace!("Relay {} <- {} closed", peer_addr, remote_addr),
+        Either::Right((Err(err), _)) => trace!("Relay {} <- {} closed with error {:?}", peer_addr, remote_addr, err),
+    }
+
+    debug!("Relay {} <-> {} closing", peer_addr, remote_addr);
+
+    Ok(())
 }
 
 /// Runs the server
-pub fn run(context: SharedContext) -> impl Future<Item = (), Error = io::Error> + Send {
-    let mut vec_fut = Vec::with_capacity(context.config().server.len());
+pub async fn run(context: SharedContext) -> io::Result<()> {
+    let vec_fut = FuturesUnordered::new();
 
     for svr_cfg in &context.config().server {
-        let listener = {
+        let mut listener = {
             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));
+            let listener = TcpListener::bind(&addr)
+                .await
+                .unwrap_or_else(|err| panic!("Failed to listen, {}", err));
 
             info!("ShadowSocks TCP Listening on {}", addr);
             listener
@@ -244,44 +208,30 @@ pub fn run(context: SharedContext) -> impl Future<Item = (), Error = io::Error>
         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();
+        vec_fut.push(async move {
+            loop {
+                match listener.accept().await {
+                    Ok((socket, peer_addr)) => {
+                        let svr_context = svr_context.clone();
+                        tokio::spawn(async move {
+                            let _ = handle_client(svr_context, socket, peer_addr).await;
+                        });
+                    }
+                    Err(err) => {
+                        error!("Server run failed: {}", err);
+                        break;
+                    }
+                }
             }
-        }
-
-        let close_guard = CloseGuard(svr_context.clone());
-
-        let listening = listener
-            .incoming()
-            .for_each(move |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));
+        });
     }
 
-    futures_unordered(vec_fut).into_future().then(|res| match res {
-        Ok(..) => {
-            error!("One of TCP servers exited unexpectly without error");
+    match vec_fut.into_future().await.0 {
+        Some(()) => {
+            error!("One of TCP servers exited unexpectly");
             let err = io::Error::new(io::ErrorKind::Other, "server exited unexpectly");
             Err(err)
         }
-        Err((err, ..)) => {
-            error!("One of TCP servers exited unexpectly with error {}", err);
-            Err(err)
-        }
-    })
+        None => unreachable!(),
+    }
 }

src/relay/tcprelay/socks5_local.rs

@@ -2,28 +2,25 @@
 
 use std::{io, net::SocketAddr, sync::Arc};
 
-use futures::{self, stream::Stream, Future};
-
+use futures::future::{self, Either};
 use log::{debug, error, info, trace, warn};
 use tokio::{
     self,
-    net::{TcpListener, TcpStream},
-};
-use tokio_io::{
-    io::{flush, ReadHalf, WriteHalf},
-    AsyncRead,
+    net::{
+        tcp::{ReadHalf, WriteHalf},
+        TcpListener, TcpStream,
+    },
+    prelude::*,
 };
 
 use crate::{config::ServerConfig, context::SharedContext};
 
 use crate::relay::{
-    boxed_future,
     loadbalancing::server::{LoadBalancer, PingBalancer},
     socks5::{self, Address, HandshakeRequest, HandshakeResponse, TcpRequestHeader, TcpResponseHeader},
-    tcprelay::crypto_io::{DecryptedRead, EncryptedWrite},
 };
 
-use super::{ignore_until_end, try_timeout, tunnel};
+use super::ignore_until_end;
 
 #[derive(Debug, Clone)]
 struct UdpConfig {
@@ -31,62 +28,95 @@ struct UdpConfig {
     client_addr: SocketAddr,
 }
 
-fn handle_socks5_connect(
+async fn handle_socks5_connect<'a>(
     context: SharedContext,
-    (r, w): (ReadHalf<TcpStream>, WriteHalf<TcpStream>),
+    (mut r, mut w): (ReadHalf<'a>, WriteHalf<'a>),
     client_addr: SocketAddr,
-    addr: Address,
+    local_addr: SocketAddr,
+    addr: &Address,
     svr_cfg: Arc<ServerConfig>,
-) -> impl Future<Item = (), Error = io::Error> + Send {
-    let timeout = svr_cfg.timeout();
-    super::connect_proxy_server(context, svr_cfg.clone())
-        .then(move |res| {
-            let (header, r) = match res {
-                Ok(svr_s) => {
-                    trace!("Proxy server connected");
+) -> io::Result<()> {
+    let svr_s = match super::connect_proxy_server(context, svr_cfg.clone()).await {
+        Ok(svr_s) => {
+            trace!("Proxy server connected, {:?}", svr_cfg);
 
-                    // Tell the client that we are ready
-                    let header = TcpResponseHeader::new(socks5::Reply::Succeeded, Address::SocketAddress(client_addr));
-
-                    (header, Ok(svr_s))
-                }
-                Err(err) => {
-                    use crate::relay::socks5::Reply;
-                    use std::io::ErrorKind;
-
-                    error!("Failed to connect remote server, err: {}", err);
-
-                    let reply = match err.kind() {
-                        ErrorKind::ConnectionRefused => Reply::ConnectionRefused,
-                        ErrorKind::ConnectionAborted => Reply::HostUnreachable,
-                        _ => Reply::NetworkUnreachable,
-                    };
-
-                    let header = TcpResponseHeader::new(reply, Address::SocketAddress(client_addr));
-
-                    (header, Err(err))
-                }
-            };
+            // Tell the client that we are ready
+            let header = TcpResponseHeader::new(socks5::Reply::Succeeded, Address::SocketAddress(local_addr));
+            header.write_to(&mut w).await?;
+            w.flush().await?;
 
             trace!("Sent header: {:?}", header);
-            try_timeout(header.write_to(w).and_then(flush), timeout).and_then(|w| match r {
-                Ok(svr_s) => Ok((svr_s, w)),
-                Err(err) => Err(err),
-            })
-        })
-        .and_then(move |(svr_s, w)| {
-            let timeout = svr_cfg.timeout();
-            super::proxy_server_handshake(svr_s, svr_cfg, addr).and_then(move |(svr_r, svr_w)| {
-                let rhalf = svr_r.and_then(move |svr_r| svr_r.copy_timeout_opt(w, timeout));
-                let whalf = svr_w.and_then(move |svr_w| svr_w.copy_timeout_opt(r, timeout));
-                tunnel(whalf, rhalf)
-            })
-        })
+
+            svr_s
+        }
+        Err(err) => {
+            use crate::relay::socks5::Reply;
+            use std::io::ErrorKind;
+
+            error!("Failed to connect remote server, err: {}", err);
+
+            let reply = match err.kind() {
+                ErrorKind::ConnectionRefused => Reply::ConnectionRefused,
+                ErrorKind::ConnectionAborted => Reply::HostUnreachable,
+                _ => Reply::NetworkUnreachable,
+            };
+
+            let header = TcpResponseHeader::new(reply, Address::SocketAddress(local_addr));
+            header.write_to(&mut w).await?;
+            w.flush().await?;
+
+            return Err(err);
+        }
+    };
+
+    let mut svr_s = super::proxy_server_handshake(svr_s, svr_cfg.clone(), addr).await?;
+    let (mut svr_r, mut svr_w) = svr_s.split();
+
+    use tokio::io::copy;
+
+    let rhalf = copy(&mut r, &mut svr_w);
+    let whalf = copy(&mut svr_r, &mut w);
+
+    debug!(
+        "CONNECT relay established {} <-> {} ({})",
+        client_addr,
+        svr_cfg.addr(),
+        addr
+    );
+
+    match future::select(rhalf, whalf).await {
+        Either::Left((Ok(..), _)) => trace!("CONNECT relay {} -> {} ({}) closed", client_addr, svr_cfg.addr(), addr),
+        Either::Left((Err(err), _)) => trace!(
+            "CONNECT relay {} -> {} ({}) closed with error {:?}",
+            client_addr,
+            svr_cfg.addr(),
+            err,
+            addr,
+        ),
+        Either::Right((Ok(..), _)) => trace!("CONNECT relay {} <- {} ({}) closed", client_addr, svr_cfg.addr(), addr),
+        Either::Right((Err(err), _)) => trace!(
+            "CONNECT relay {} <- {} ({}) closed with error {:?}",
+            client_addr,
+            svr_cfg.addr(),
+            err,
+            addr,
+        ),
+    }
+
+    debug!(
+        "CONNECT relay {} <-> {} ({}) closing",
+        client_addr,
+        svr_cfg.addr(),
+        addr
+    );
+
+    Ok(())
 }
 
-fn handle_socks5_client(
+#[allow(clippy::cognitive_complexity)]
+async fn handle_socks5_client(
     context: SharedContext,
-    s: TcpStream,
+    mut s: TcpStream,
     conf: Arc<ServerConfig>,
     udp_conf: UdpConfig,
 ) -> io::Result<()> {
@@ -101,120 +131,108 @@ fn handle_socks5_client(
     }
 
     let client_addr = s.peer_addr()?;
-    let cloned_client_addr = client_addr;
-    let fut = futures::lazy(|| Ok(s.split()))
-        .and_then(|(r, w)| {
-            // Socks5 handshakes
-            HandshakeRequest::read_from(r).and_then(move |(r, req)| {
-                trace!("Socks5 {:?}", req);
+    let local_addr = s.local_addr()?;
 
-                let (resp, res) = if !req.methods.contains(&socks5::SOCKS5_AUTH_METHOD_NONE) {
-                    let resp = HandshakeResponse::new(socks5::SOCKS5_AUTH_METHOD_NOT_ACCEPTABLE);
-                    warn!("Currently shadowsocks-rust does not support authentication");
-                    (
-                        resp,
-                        Err(io::Error::new(
-                            io::ErrorKind::Other,
-                            "Currently shadowsocks-rust does not support authentication",
-                        )),
-                    )
-                } else {
-                    // Reply to client
-                    let resp = HandshakeResponse::new(socks5::SOCKS5_AUTH_METHOD_NONE);
-                    trace!("Reply handshake {:?}", resp);
-                    (resp, Ok(()))
-                };
+    let (mut r, mut w) = s.split();
 
-                resp.write_to(w).and_then(flush).and_then(move |w| match res {
-                    Ok(..) => Ok((r, w)),
-                    Err(err) => Err(err),
-                })
-            })
-        })
-        .and_then(move |(r, w)| {
-            // Fetch headers
-            TcpRequestHeader::read_from(r).then(move |res| match res {
-                Ok((r, h)) => boxed_future(futures::finished((r, w, h))),
-                Err(err) => {
-                    error!("Failed to get TcpRequestHeader: {}", err);
-                    let fut = TcpResponseHeader::new(err.reply, Address::SocketAddress(client_addr))
-                        .write_to(w)
-                        .then(|_| Err(From::from(err)));
-                    boxed_future(fut)
-                }
-            })
-        })
-        .and_then(move |(r, w, header)| {
-            trace!("Socks5 {:?}", header);
+    let handshake_req = HandshakeRequest::read_from(&mut r).await?;
 
-            let addr = header.address;
-            match header.command {
-                socks5::Command::TcpConnect => {
-                    let enable_tcp = context.config().mode.enable_tcp();
-                    if enable_tcp {
-                        debug!("CONNECT {}", addr);
-                        let fut = handle_socks5_connect(context, (r, w), cloned_client_addr, addr.clone(), conf)
-                            .map_err(move |err| {
-                                io::Error::new(err.kind(), format!("CONNECT {} failed with error \"{}\"", addr, err))
-                            });
-                        boxed_future(fut)
-                    } else {
-                        warn!("CONNECT is not enabled");
-                        let fut = TcpResponseHeader::new(socks5::Reply::CommandNotSupported, addr)
-                            .write_to(w)
-                            .map(|_| ());
-                        boxed_future(fut)
-                    }
-                }
-                socks5::Command::TcpBind => {
-                    warn!("BIND is not supported");
-                    let fut = TcpResponseHeader::new(socks5::Reply::CommandNotSupported, addr)
-                        .write_to(w)
-                        .map(|_| ());
-                    boxed_future(fut)
-                }
-                socks5::Command::UdpAssociate => {
-                    if udp_conf.enable_udp {
-                        debug!("UDP ASSOCIATE {}", addr);
-                        let fut = TcpResponseHeader::new(socks5::Reply::Succeeded, From::from(udp_conf.client_addr))
-                            .write_to(w)
-                            .and_then(flush)
-                            .and_then(|_| {
-                                // Hold the connection until it ends by its own
-                                ignore_until_end(r).map(|_| ())
-                            });
+    // Socks5 handshakes
+    trace!("Socks5 {:?}", handshake_req);
 
-                        boxed_future(fut)
-                    } else {
-                        warn!("UDP ASSOCIATE is not enabled");
-                        let fut = TcpResponseHeader::new(socks5::Reply::CommandNotSupported, addr)
-                            .write_to(w)
-                            .map(|_| ());
-                        boxed_future(fut)
-                    }
-                }
-            }
-        });
+    let (handshake_resp, res) = if !handshake_req.methods.contains(&socks5::SOCKS5_AUTH_METHOD_NONE) {
+        let resp = HandshakeResponse::new(socks5::SOCKS5_AUTH_METHOD_NOT_ACCEPTABLE);
+        warn!("Currently shadowsocks-rust does not support authentication");
+        (
+            resp,
+            Err(io::Error::new(
+                io::ErrorKind::Other,
+                "Currently shadowsocks-rust does not support authentication",
+            )),
+        )
+    } else {
+        // Reply to client
+        let resp = HandshakeResponse::new(socks5::SOCKS5_AUTH_METHOD_NONE);
+        trace!("Reply handshake {:?}", resp);
+        (resp, Ok(()))
+    };
 
-    tokio::spawn(fut.then(|res| match res {
-        Ok(..) => Ok(()),
+    handshake_resp.write_to(&mut w).await?;
+    w.flush().await?;
+
+    res?;
+
+    // Fetch headers
+    let header = match TcpRequestHeader::read_from(&mut r).await {
+        Ok(h) => h,
         Err(err) => {
-            if err.kind() != io::ErrorKind::BrokenPipe {
-                error!("Socks5 client: {}", err);
-            }
-            Err(())
+            error!("Failed to get TcpRequestHeader: {}", err);
+            let rh = TcpResponseHeader::new(err.reply, Address::SocketAddress(client_addr));
+            rh.write_to(&mut w).await?;
+            return Err(From::from(err));
         }
-    }));
+    };
 
-    Ok(())
+    trace!("Socks5 {:?}", header);
+
+    let addr = header.address;
+    match header.command {
+        socks5::Command::TcpConnect => {
+            let enable_tcp = context.config().mode.enable_tcp();
+            if enable_tcp {
+                debug!("CONNECT {}", addr);
+
+                match handle_socks5_connect(context, (r, w), client_addr, local_addr, &addr, conf).await {
+                    Ok(..) => Ok(()),
+                    Err(err) => Err(io::Error::new(
+                        err.kind(),
+                        format!("CONNECT {} failed with error \"{}\"", addr, err),
+                    )),
+                }
+            } else {
+                warn!("CONNECT is not enabled");
+                let rh = TcpResponseHeader::new(socks5::Reply::CommandNotSupported, addr);
+                rh.write_to(&mut w).await?;
+
+                Ok(())
+            }
+        }
+        socks5::Command::TcpBind => {
+            warn!("BIND is not supported");
+            let rh = TcpResponseHeader::new(socks5::Reply::CommandNotSupported, addr);
+            rh.write_to(&mut w).await?;
+
+            Ok(())
+        }
+        socks5::Command::UdpAssociate => {
+            if udp_conf.enable_udp {
+                debug!("UDP ASSOCIATE {}", addr);
+                let rh = TcpResponseHeader::new(socks5::Reply::Succeeded, From::from(udp_conf.client_addr));
+                rh.write_to(&mut w).await?;
+                w.flush().await?;
+
+                // Hold the connection until it ends by its own
+                ignore_until_end(&mut r).await?;
+
+                Ok(())
+            } else {
+                warn!("UDP ASSOCIATE is not enabled");
+                let rh = TcpResponseHeader::new(socks5::Reply::CommandNotSupported, addr);
+                rh.write_to(&mut w).await?;
+
+                Ok(())
+            }
+        }
+    }
 }
 
 /// Starts a TCP local server with Socks5 proxy protocol
-pub fn run(context: SharedContext) -> impl Future<Item = (), Error = io::Error> + Send {
+pub async fn run(context: SharedContext) -> io::Result<()> {
     let local_addr = *context.config().local.as_ref().expect("Missing local config");
 
-    let listener =
-        TcpListener::bind(&local_addr).unwrap_or_else(|err| panic!("Failed to listen on {}, {}", local_addr, err));
+    let mut listener = TcpListener::bind(&local_addr)
+        .await
+        .unwrap_or_else(|err| panic!("Failed to listen on {}, {}", local_addr, err));
 
     let actual_local_addr = listener.local_addr().expect("Could not determine port bound to");
 
@@ -226,12 +244,20 @@ pub fn run(context: SharedContext) -> impl Future<Item = (), Error = io::Error>
     };
 
     let mut servers = PingBalancer::new(context.clone());
-    listener.incoming().for_each(move |socket| {
+
+    loop {
+        let (socket, peer_addr) = listener.accept().await?;
         let server_cfg = servers.pick_server();
 
-        trace!("Got connection, addr: {}", socket.peer_addr()?);
+        trace!("Got connection, addr: {}", peer_addr);
         trace!("Picked proxy server: {:?}", server_cfg);
 
-        handle_socks5_client(context.clone(), socket, server_cfg, udp_conf.clone())
-    })
+        let context = context.clone();
+        let udp_conf = udp_conf.clone();
+        tokio::spawn(async move {
+            if let Err(err) = handle_socks5_client(context, socket, server_cfg, udp_conf).await {
+                error!("Socks5 client {}", err);
+            }
+        });
+    }
 }

src/relay/tcprelay/stream.rs

@@ -1,159 +1,145 @@
 //! Stream protocol implementation
 
 use std::{
-    cmp,
-    io::{self, BufRead, Read},
+    cmp, io,
+    marker::Unpin,
+    pin::Pin,
+    task::{Context, Poll},
 };
 
 use crate::crypto::{new_stream, BoxStreamCipher, CipherType, CryptoMode};
 use bytes::{BufMut, BytesMut};
-use tokio_io::{AsyncRead, AsyncWrite};
+use futures::ready;
+use tokio::prelude::*;
 
-use super::{DecryptedRead, EncryptedWrite, BUFFER_SIZE};
+use super::BUFFER_SIZE;
 
 const DUMMY_BUFFER: [u8; BUFFER_SIZE] = [0u8; BUFFER_SIZE];
 
 /// Reader wrapper that will decrypt data automatically
-pub struct DecryptedReader<R>
-where
-    R: AsyncRead,
-{
-    reader: R,
+pub struct DecryptedReader {
     buffer: BytesMut,
     cipher: BoxStreamCipher,
     pos: usize,
-    sent_final: bool,
+    got_final: bool,
+    incoming_buffer: Vec<u8>,
 }
 
-impl<R> DecryptedReader<R>
-where
-    R: AsyncRead,
-{
-    pub fn new(r: R, t: CipherType, key: &[u8], iv: &[u8]) -> DecryptedReader<R> {
+impl DecryptedReader {
+    pub fn new(t: CipherType, key: &[u8], iv: &[u8]) -> DecryptedReader {
         let cipher = new_stream(t, key, iv, CryptoMode::Decrypt);
         let buffer_size = cipher.buffer_size(&DUMMY_BUFFER);
         DecryptedReader {
-            reader: r,
             buffer: BytesMut::with_capacity(buffer_size),
             cipher,
             pos: 0,
-            sent_final: false,
+            got_final: false,
+            incoming_buffer: vec![0u8; BUFFER_SIZE],
         }
     }
 
-    pub fn get_ref(&self) -> &R {
-        &self.reader
-    }
-
-    /// Gets a mutable reference to the underlying reader.
-    ///
-    /// # Warning
-    ///
-    /// It is inadvisable to read directly from or write directly to the
-    /// underlying reader.
-    pub fn get_mut(&mut self) -> &mut R {
-        &mut self.reader
-    }
-
-    /// Unwraps this `DecryptedReader`, returning the underlying reader.
-    ///
-    /// The internal buffer is flushed before returning the reader. Any leftover
-    /// data in the read buffer is lost.
-    pub fn into_inner(self) -> R {
-        self.reader
-    }
-}
-
-impl<R> BufRead for DecryptedReader<R>
-where
-    R: AsyncRead,
-{
-    fn fill_buf(&mut self) -> io::Result<&[u8]> {
+    pub fn poll_read_decrypted<R>(
+        &mut self,
+        ctx: &mut Context<'_>,
+        r: &mut R,
+        dst: &mut [u8],
+    ) -> Poll<io::Result<usize>>
+    where
+        R: AsyncRead + Unpin,
+    {
         while self.pos >= self.buffer.len() {
-            if self.sent_final {
-                return Ok(&[]);
+            if self.got_final {
+                return Poll::Ready(Ok(0));
             }
 
-            let mut incoming = [0u8; BUFFER_SIZE];
+            let n = ready!(Pin::new(&mut *r).poll_read(ctx, &mut self.incoming_buffer))?;
+
+            // Reset pointers
             self.buffer.clear();
-            match self.reader.read(&mut incoming) {
-                Ok(0) => {
-                    // Ensure we have enough space
-                    let buffer_len = self.buffer_size(&[]);
-                    self.buffer.reserve(buffer_len);
-
-                    // EOF
-                    self.cipher.finalize(&mut self.buffer)?;
-                    self.sent_final = true;
-                }
-                Ok(l) => {
-                    let data = &incoming[..l];
-
-                    // Ensure we have enough space
-                    let buffer_len = self.buffer_size(data);
-                    self.buffer.reserve(buffer_len);
-
-                    self.cipher.update(data, &mut self.buffer)?;
-                }
-                Err(err) => {
-                    return Err(err);
-                }
-            }
-
             self.pos = 0;
+
+            if n == 0 {
+                // Finialize block
+                self.buffer.reserve(self.buffer_size(&[]));
+                self.cipher.finalize(&mut self.buffer)?;
+                self.got_final = true;
+            } else {
+                let data = &self.incoming_buffer[..n];
+                // Ensure we have enough space
+                let buffer_len = self.buffer_size(data);
+                self.buffer.reserve(buffer_len);
+                self.cipher.update(data, &mut self.buffer)?;
+            }
         }
 
-        Ok(&self.buffer[self.pos..])
+        let remaining_len = self.buffer.len() - self.pos;
+        let n = cmp::min(dst.len(), remaining_len);
+        (&mut dst[..n]).copy_from_slice(&self.buffer[self.pos..self.pos + n]);
+        self.pos += n;
+        Poll::Ready(Ok(n))
     }
 
-    fn consume(&mut self, amt: usize) {
-        self.pos = cmp::min(self.pos + amt, self.buffer.len());
-    }
-}
-
-impl<R> Read for DecryptedReader<R>
-where
-    R: AsyncRead,
-{
-    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        let nread = {
-            let mut available = self.fill_buf()?;
-            available.read(buf)?
-        };
-        self.consume(nread);
-        Ok(nread)
-    }
-}
-
-impl<R> DecryptedRead for DecryptedReader<R>
-where
-    R: AsyncRead,
-{
     fn buffer_size(&self, data: &[u8]) -> usize {
         self.cipher.buffer_size(data)
     }
 }
 
-impl<R> AsyncRead for DecryptedReader<R> where R: AsyncRead {}
-
-/// Writer wrapper that will encrypt data automatically
-pub struct EncryptedWriter<W>
-where
-    W: AsyncWrite,
-{
-    writer: W,
-    cipher: BoxStreamCipher,
+enum EncryptWriteStep {
+    Nothing,
+    Writing(BytesMut, usize),
 }
 
-impl<W> EncryptedWriter<W>
-where
-    W: AsyncWrite,
-{
+/// Writer wrapper that will encrypt data automatically
+pub struct EncryptedWriter {
+    cipher: BoxStreamCipher,
+    steps: EncryptWriteStep,
+}
+
+impl EncryptedWriter {
     /// Creates a new EncryptedWriter
-    pub fn new(w: W, t: CipherType, key: &[u8], iv: &[u8]) -> EncryptedWriter<W> {
+    pub fn new(t: CipherType, key: &[u8], iv: &[u8]) -> EncryptedWriter {
         EncryptedWriter {
-            writer: w,
             cipher: new_stream(t, key, iv, CryptoMode::Encrypt),
+            steps: EncryptWriteStep::Nothing,
+        }
+    }
+
+    pub fn poll_write_encrypted<W>(&mut self, ctx: &mut Context<'_>, w: &mut W, data: &[u8]) -> Poll<io::Result<usize>>
+    where
+        W: AsyncWrite + Unpin,
+    {
+        ready!(self.poll_write_all_encrypted(ctx, w, data))?;
+        Poll::Ready(Ok(data.len()))
+    }
+
+    pub fn poll_write_all_encrypted<W>(&mut self, ctx: &mut Context<'_>, w: &mut W, data: &[u8]) -> Poll<io::Result<()>>
+    where
+        W: AsyncWrite + Unpin,
+    {
+        // FIXME: How about finalize?
+
+        loop {
+            match self.steps {
+                EncryptWriteStep::Nothing => {
+                    let mut buf = BytesMut::with_capacity(self.buffer_size(data));
+                    self.cipher_update(data, &mut buf)?;
+
+                    self.steps = EncryptWriteStep::Writing(buf, 0);
+                }
+                EncryptWriteStep::Writing(ref mut buf, ref mut pos) => {
+                    while *pos < buf.len() {
+                        let n = ready!(Pin::new(&mut *w).poll_write(ctx, &buf[*pos..]))?;
+                        if n == 0 {
+                            use std::io::ErrorKind;
+                            return Poll::Ready(Err(ErrorKind::UnexpectedEof.into()));
+                        }
+                        *pos += n;
+                    }
+
+                    self.steps = EncryptWriteStep::Nothing;
+                    return Poll::Ready(Ok(()));
+                }
+            }
         }
     }
 
@@ -161,40 +147,10 @@ where
         self.cipher.update(data, buf).map_err(From::from)
     }
 
+    #[allow(dead_code)]
     fn cipher_finalize<B: BufMut>(&mut self, buf: &mut B) -> io::Result<()> {
         self.cipher.finalize(buf).map_err(From::from)
     }
-}
-
-impl<W> Drop for EncryptedWriter<W>
-where
-    W: AsyncWrite,
-{
-    fn drop(&mut self) {
-        let mut buf = Vec::new();
-        if let Ok(..) = self.cipher_finalize(&mut buf) {
-            if !buf.is_empty() {
-                let _ = self.write_raw(&buf);
-            }
-        }
-    }
-}
-
-impl<W> EncryptedWrite for EncryptedWriter<W>
-where
-    W: AsyncWrite,
-{
-    fn write_raw(&mut self, data: &[u8]) -> io::Result<usize> {
-        self.writer.write(data)
-    }
-
-    fn flush(&mut self) -> io::Result<()> {
-        self.writer.flush()
-    }
-
-    fn encrypt<B: BufMut>(&mut self, data: &[u8], buf: &mut B) -> io::Result<()> {
-        self.cipher_update(data, buf)
-    }
 
     fn buffer_size(&self, data: &[u8]) -> usize {
         self.cipher.buffer_size(data)

src/relay/tcprelay/tunnel_local.rs

@@ -0,0 +1,129 @@
+//! Local server that establish a TCP tunnel with server
+
+use std::io;
+use std::net::SocketAddr;
+use std::sync::Arc;
+
+use futures::future::{self, Either};
+use log::{debug, error, info, trace};
+use tokio::net::{TcpListener, TcpStream};
+
+use crate::config::ServerConfig;
+use crate::context::SharedContext;
+use crate::relay::loadbalancing::server::{LoadBalancer, PingBalancer};
+use crate::relay::socks5::Address;
+
+/// Established Client Tunnel
+///
+/// This method must be called after handshaking with client (for example, socks5 handshaking)
+async fn establish_client_tcp_tunnel<'a>(
+    context: SharedContext,
+    mut s: TcpStream,
+    client_addr: SocketAddr,
+    addr: &Address,
+    svr_cfg: Arc<ServerConfig>,
+) -> io::Result<()> {
+    let svr_s = match super::connect_proxy_server(context, svr_cfg.clone()).await {
+        Ok(svr_s) => {
+            trace!("Proxy server connected, {:?}", svr_cfg);
+            svr_s
+        }
+        Err(err) => {
+            // Just close the connection.
+            return Err(err);
+        }
+    };
+
+    let mut svr_s = super::proxy_server_handshake(svr_s, svr_cfg.clone(), addr).await?;
+    let (mut svr_r, mut svr_w) = svr_s.split();
+
+    let (mut r, mut w) = s.split();
+
+    use tokio::io::copy;
+
+    let rhalf = copy(&mut r, &mut svr_w);
+    let whalf = copy(&mut svr_r, &mut w);
+
+    debug!(
+        "TUNNEL relay established {} <-> {} ({})",
+        client_addr,
+        svr_cfg.addr(),
+        addr
+    );
+
+    match future::select(rhalf, whalf).await {
+        Either::Left((Ok(..), _)) => trace!("TUNNEL relay {} -> {} ({}) closed", client_addr, svr_cfg.addr(), addr),
+        Either::Left((Err(err), _)) => trace!(
+            "TUNNEL relay {} -> {} ({}) closed with error {:?}",
+            client_addr,
+            svr_cfg.addr(),
+            err,
+            addr,
+        ),
+        Either::Right((Ok(..), _)) => trace!("TUNNEL relay {} <- {} ({}) closed", client_addr, svr_cfg.addr(), addr),
+        Either::Right((Err(err), _)) => trace!(
+            "TUNNEL relay {} <- {} ({}) closed with error {:?}",
+            client_addr,
+            svr_cfg.addr(),
+            err,
+            addr
+        ),
+    }
+
+    debug!("TUNNEL relay {} <-> {} ({}) closing", client_addr, svr_cfg.addr(), addr);
+
+    Ok(())
+}
+
+async fn handle_tunnel_client(context: SharedContext, s: TcpStream, conf: Arc<ServerConfig>) -> io::Result<()> {
+    if let Err(err) = s.set_keepalive(conf.timeout()) {
+        error!("Failed to set keep alive: {:?}", err);
+    }
+
+    if context.config().no_delay {
+        if let Err(err) = s.set_nodelay(true) {
+            error!("Failed to set no delay: {:?}", err);
+        }
+    }
+
+    let client_addr = s.peer_addr()?;
+
+    // forward must not be None, it is already checked in local.rs
+    let target_addr = context.config().forward.as_ref().unwrap();
+
+    establish_client_tcp_tunnel(context.clone(), s, client_addr, target_addr, conf).await
+}
+
+pub async fn run(context: SharedContext) -> io::Result<()> {
+    assert!(
+        context.config().mode.enable_tcp(),
+        "You must enable TCP relay for tunneling"
+    );
+
+    let local_addr = *context.config().local.as_ref().expect("Missing local config");
+
+    let mut listener = TcpListener::bind(&local_addr)
+        .await
+        .unwrap_or_else(|err| panic!("Failed to listen on {}, {}", local_addr, err));
+
+    let actual_local_addr = listener.local_addr().expect("Could not determine port bound to");
+
+    info!("ShadowSocks TCP Tunnel Listening on {}", actual_local_addr);
+
+    let mut servers = PingBalancer::new(context.clone());
+
+    loop {
+        let (socket, peer_addr) = listener.accept().await?;
+        let server_cfg = servers.pick_server();
+
+        trace!("Got connection, addr: {}", peer_addr);
+        trace!("Picked proxy server: {:?}", server_cfg);
+
+        let context = context.clone();
+        tokio::spawn(async move {
+            if let Err(err) = handle_tunnel_client(context, socket, server_cfg).await {
+                error!("Socks5 client {}", err);
+            }
+        });
+    }
+}

src/relay/tcprelay/utils.rs

@@ -1,182 +1 @@
 //! Utility functions
-
-use std::{
-    io,
-    time::{Duration, Instant},
-};
-
-use futures::{Async, Future, Poll};
-use tokio::timer::Delay;
-use tokio_io::{
-    io::{copy, Copy},
-    try_nb, AsyncRead, AsyncWrite,
-};
-
-use super::BUFFER_SIZE;
-
-/// Copies all data from `r` to `w`, abort if timeout reaches
-pub fn copy_timeout<R, W>(r: R, w: W, dur: Duration) -> CopyTimeout<R, W>
-where
-    R: AsyncRead,
-    W: AsyncWrite,
-{
-    CopyTimeout::new(r, w, dur)
-}
-
-/// Copies all data from `r` to `w`, abort if timeout reaches
-pub struct CopyTimeout<R, W>
-where
-    R: AsyncRead,
-    W: AsyncWrite,
-{
-    r: Option<R>,
-    w: Option<W>,
-    timeout: Duration,
-    amt: u64,
-    timer: Option<Delay>,
-    buf: [u8; BUFFER_SIZE],
-    pos: usize,
-    cap: usize,
-}
-
-impl<R, W> CopyTimeout<R, W>
-where
-    R: AsyncRead,
-    W: AsyncWrite,
-{
-    fn new(r: R, w: W, timeout: Duration) -> CopyTimeout<R, W> {
-        CopyTimeout {
-            r: Some(r),
-            w: Some(w),
-            timeout,
-            amt: 0,
-            timer: None,
-            buf: [0u8; BUFFER_SIZE],
-            pos: 0,
-            cap: 0,
-        }
-    }
-
-    fn try_poll_timeout(&mut self) -> io::Result<()> {
-        match self.timer.as_mut() {
-            None => Ok(()),
-            Some(t) => match t.poll() {
-                Err(err) => panic!("Failed to poll on timer, err: {}", err),
-                Ok(Async::Ready(..)) => Err(io::Error::new(io::ErrorKind::TimedOut, "connection timed out")),
-                Ok(Async::NotReady) => Ok(()),
-            },
-        }
-    }
-
-    fn clear_timer(&mut self) {
-        let _ = self.timer.take();
-    }
-
-    fn read_or_set_timeout(&mut self) -> io::Result<usize> {
-        // First, return if timeout
-        self.try_poll_timeout()?;
-
-        // Then, unset the previous timeout
-        self.clear_timer();
-
-        match self.r.as_mut().unwrap().read(&mut self.buf) {
-            Ok(n) => Ok(n),
-            Err(e) => {
-                if e.kind() == io::ErrorKind::WouldBlock {
-                    self.timer = Some(Delay::new(Instant::now() + self.timeout));
-                }
-                Err(e)
-            }
-        }
-    }
-
-    fn write_or_set_timeout(&mut self, beg: usize, end: usize) -> io::Result<usize> {
-        // First, return if timeout
-        self.try_poll_timeout()?;
-
-        // Then, unset the previous timeout
-        self.clear_timer();
-
-        match self.w.as_mut().unwrap().write(&self.buf[beg..end]) {
-            Ok(n) => Ok(n),
-            Err(e) => {
-                if e.kind() == io::ErrorKind::WouldBlock {
-                    self.timer = Some(Delay::new(Instant::now() + self.timeout));
-                }
-                Err(e)
-            }
-        }
-    }
-}
-
-impl<R, W> Future for CopyTimeout<R, W>
-where
-    R: AsyncRead,
-    W: AsyncWrite,
-{
-    type Error = io::Error;
-    type Item = (u64, R, W);
-
-    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
-        loop {
-            if self.pos == self.cap {
-                let n = try_nb!(self.read_or_set_timeout());
-
-                if n == 0 {
-                    // If we've written all the data and we've seen EOF, flush out the
-                    // data and finish the transfer.
-                    // done with the entire transfer.
-                    try_nb!(self.w.as_mut().unwrap().flush());
-                    return Ok((self.amt, self.r.take().unwrap(), self.w.take().unwrap()).into());
-                }
-
-                self.pos = 0;
-                self.cap = n;
-
-                // Clear it before write
-                self.clear_timer();
-            }
-
-            // If our buffer has some data, let's write it out!
-            while self.pos < self.cap {
-                let (pos, cap) = (self.pos, self.cap);
-                let i = try_nb!(self.write_or_set_timeout(pos, cap));
-                self.pos += i;
-                self.amt += i as u64;
-            }
-
-            // Clear it before read
-            self.clear_timer();
-        }
-    }
-}
-
-/// Copies all data from `r` to `w` with optional timeout param
-pub fn copy_timeout_opt<R, W>(r: R, w: W, dur: Option<Duration>) -> CopyTimeoutOpt<R, W>
-where
-    R: AsyncRead,
-    W: AsyncWrite,
-{
-    match dur {
-        Some(d) => CopyTimeoutOpt::CopyTimeout(copy_timeout(r, w, d)),
-        None => CopyTimeoutOpt::Copy(copy(r, w)),
-    }
-}
-
-/// Copies all data from `R` to `W`
-pub enum CopyTimeoutOpt<R: AsyncRead, W: AsyncWrite> {
-    Copy(Copy<R, W>),
-    CopyTimeout(CopyTimeout<R, W>),
-}
-
-impl<R: AsyncRead, W: AsyncWrite> Future for CopyTimeoutOpt<R, W> {
-    type Error = io::Error;
-    type Item = (u64, R, W);
-
-    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
-        match *self {
-            CopyTimeoutOpt::CopyTimeout(ref mut c) => c.poll(),
-            CopyTimeoutOpt::Copy(ref mut c) => c.poll(),
-        }
-    }
-}

src/relay/udprelay/crypto_io.rs

@@ -21,81 +21,105 @@
 //! ```
 
 use std::io;
+use std::slice;
+
+use byte_string::ByteStr;
+use bytes::{BufMut, BytesMut};
+use log::trace;
 
 use crate::crypto::{self, CipherCategory, CipherType, CryptoMode};
 
 /// Encrypt payload into ShadowSocks UDP encrypted packet
-pub fn encrypt_payload(t: CipherType, key: &[u8], payload: &[u8]) -> io::Result<Vec<u8>> {
+pub fn encrypt_payload(t: CipherType, key: &[u8], payload: &[u8], dst: &mut BytesMut) -> io::Result<()> {
     match t.category() {
-        CipherCategory::Stream => encrypt_payload_stream(t, key, payload),
-        CipherCategory::Aead => encrypt_payload_aead(t, key, payload),
+        CipherCategory::Stream => encrypt_payload_stream(t, key, payload, dst),
+        CipherCategory::Aead => encrypt_payload_aead(t, key, payload, dst),
     }
 }
 
-fn encrypt_payload_stream(t: CipherType, key: &[u8], payload: &[u8]) -> io::Result<Vec<u8>> {
+fn encrypt_payload_stream(t: CipherType, key: &[u8], payload: &[u8], dst: &mut BytesMut) -> io::Result<()> {
     let iv = t.gen_init_vec();
     let mut cipher = crypto::new_stream(t, key, &iv, CryptoMode::Encrypt);
 
-    let mut send_payload = Vec::with_capacity(iv.len() + payload.len());
-    send_payload.extend_from_slice(&iv);
-    cipher.update(&payload[..], &mut send_payload)?;
-    cipher.finalize(&mut send_payload)?;
-    Ok(send_payload)
+    trace!("UDP packet generated stream iv {:?}", ByteStr::new(&iv));
+
+    dst.reserve(iv.len() + payload.len());
+
+    // First of all, IV
+    dst.put_slice(&iv);
+
+    // Encrypted data
+    cipher.update(&payload[..], dst)?;
+    cipher.finalize(dst)?;
+
+    Ok(())
 }
 
-fn encrypt_payload_aead(t: CipherType, key: &[u8], payload: &[u8]) -> io::Result<Vec<u8>> {
+fn encrypt_payload_aead(t: CipherType, key: &[u8], payload: &[u8], dst: &mut BytesMut) -> io::Result<()> {
     let salt = t.gen_salt();
     let tag_size = t.tag_size();
     let mut cipher = crypto::new_aead_encryptor(t, key, &salt);
 
-    let mut send_payload = Vec::with_capacity(salt.len() + payload.len() + tag_size);
-    send_payload.extend_from_slice(&salt);
-    let start_pos = send_payload.len();
-    send_payload.resize(start_pos + payload.len() + tag_size, 0);
+    trace!("UDP packet generated AEAD salt {:?}", ByteStr::new(&salt));
 
-    cipher.encrypt(payload, &mut send_payload[start_pos..]);
+    dst.reserve(salt.len() + payload.len() + tag_size);
 
-    Ok(send_payload)
+    // First of all, salt
+    dst.put_slice(&salt);
+
+    // Encrypted data
+    unsafe {
+        let remaining = dst.bytes_mut();
+        let b = slice::from_raw_parts_mut(remaining.as_mut_ptr() as *mut u8, remaining.len());
+
+        cipher.encrypt(payload, b);
+        dst.advance_mut(payload.len() + tag_size);
+    }
+
+    Ok(())
 }
 
 /// Decrypt payload from ShadowSocks UDP encrypted packet
-pub fn decrypt_payload(t: CipherType, key: &[u8], payload: &[u8]) -> io::Result<Vec<u8>> {
+pub fn decrypt_payload(t: CipherType, key: &[u8], payload: &[u8]) -> io::Result<Option<Vec<u8>>> {
     match t.category() {
         CipherCategory::Stream => decrypt_payload_stream(t, key, payload),
         CipherCategory::Aead => decrypt_payload_aead(t, key, payload),
     }
 }
 
-fn decrypt_payload_stream(t: CipherType, key: &[u8], payload: &[u8]) -> io::Result<Vec<u8>> {
+fn decrypt_payload_stream(t: CipherType, key: &[u8], payload: &[u8]) -> io::Result<Option<Vec<u8>>> {
     let iv_size = t.iv_size();
     if payload.len() < iv_size {
-        let err = io::Error::new(io::ErrorKind::Other, "udp packet too short");
-        return Err(err);
+        return Ok(None);
     }
 
     let iv = &payload[..iv_size];
     let data = &payload[iv_size..];
 
+    trace!("UDP packet got stream IV {:?}", ByteStr::new(iv));
+
     let mut cipher = crypto::new_stream(t, key, iv, CryptoMode::Decrypt);
 
     let mut recv_payload = Vec::with_capacity(data.len());
     cipher.update(data, &mut recv_payload)?;
     cipher.finalize(&mut recv_payload)?;
 
-    Ok(recv_payload)
+    Ok(Some(recv_payload))
 }
 
-fn decrypt_payload_aead(t: CipherType, key: &[u8], payload: &[u8]) -> io::Result<Vec<u8>> {
+fn decrypt_payload_aead(t: CipherType, key: &[u8], payload: &[u8]) -> io::Result<Option<Vec<u8>>> {
     let tag_size = t.tag_size();
     let salt_size = t.salt_size();
 
     if payload.len() < tag_size + salt_size {
-        let err = io::Error::new(io::ErrorKind::Other, "udp packet too short");
-        return Err(err);
+        return Ok(None);
     }
 
     let salt = &payload[..salt_size];
     let data = &payload[salt_size..];
+
+    trace!("UDP packet got AEAD salt {:?}", ByteStr::new(salt));
+
     let data_length = payload.len() - tag_size - salt_size;
 
     let mut cipher = crypto::new_aead_decryptor(t, key, salt);
@@ -103,5 +127,5 @@ fn decrypt_payload_aead(t: CipherType, key: &[u8], payload: &[u8]) -> io::Result
     let mut recv_payload = vec![0u8; data_length];
     cipher.decrypt(data, &mut recv_payload)?;
 
-    Ok(recv_payload)
+    Ok(Some(recv_payload))
 }

src/relay/udprelay/dns.rs

@@ -1,205 +0,0 @@
-//! UDP DNS relay
-
-use std::{
-    io::{self, Cursor},
-    net::{IpAddr, Ipv4Addr, SocketAddr},
-    sync::{Arc, Mutex},
-    time::Instant,
-};
-
-use dns_parser::Packet;
-use futures::{self, future::join_all, stream::futures_unordered, Future, Stream};
-use log::{error, info, trace};
-use tokio::{self, net::UdpSocket};
-
-use super::{
-    crypto_io::{decrypt_payload, encrypt_payload},
-    PacketStream, SendDgramRc, SharedUdpSocket,
-};
-use crate::{
-    config::{ServerAddr, ServerConfig},
-    context::SharedContext,
-    relay::{boxed_future, dns_resolver::resolve, socks5::Address},
-};
-
-/// Starts a UDP DNS server
-pub fn run(context: SharedContext) -> impl Future<Item = (), Error = io::Error> + Send {
-    let local_addr = *context.config().local.as_ref().unwrap();
-
-    futures::lazy(move || {
-        info!("ShadowSocks UDP DNS Listening on {}", local_addr);
-
-        UdpSocket::bind(&local_addr)
-    })
-    .and_then(move |l| listen(context, l))
-}
-
-fn listen(context: SharedContext, l: UdpSocket) -> impl Future<Item = (), Error = io::Error> + Send {
-    assert!(!context.config().server.is_empty());
-
-    let mut svr_fut = Vec::with_capacity(context.config().server.len());
-    for svr in &context.config().server {
-        match *svr.addr() {
-            ServerAddr::SocketAddr(ref addr) => {
-                svr_fut.push(boxed_future(futures::finished::<_, io::Error>(vec![*addr])));
-            }
-            ServerAddr::DomainName(ref dom, ref port) => {
-                svr_fut.push(boxed_future(resolve(context.clone(), &*dom, *port, false)));
-            }
-        }
-    }
-
-    let cloned_context = context.clone();
-    join_all(svr_fut)
-        .and_then(move |svr_addrs| {
-            let mut u = Vec::with_capacity(svr_addrs.len());
-            for (idx, svr_addr) in svr_addrs.into_iter().enumerate() {
-                let local_addr = SocketAddr::new(IpAddr::from(Ipv4Addr::new(0, 0, 0, 0)), 0);
-                let s = UdpSocket::bind(&local_addr)?;
-                let svr_cfg = Arc::new(cloned_context.config().server[idx].clone());
-                u.push((Arc::new(Mutex::new(s)), svr_cfg, svr_addr[0]));
-            }
-            Ok(u)
-        })
-        .and_then(move |vec_servers| {
-            let mut f = Vec::with_capacity(1 + vec_servers.len());
-            let l = Arc::new(Mutex::new(l));
-            f.push(boxed_future(handle_l2r(
-                context.clone(),
-                l.clone(),
-                vec_servers.clone(),
-            )));
-            for (svr, cfg, _) in vec_servers {
-                f.push(boxed_future(handle_r2l(context.clone(), l.clone(), svr, cfg)))
-            }
-            futures_unordered(f).into_future().then(|res| match res {
-                Ok(..) => {
-                    error!("One of DNS servers exited unexpectly without error");
-                    let err = io::Error::new(io::ErrorKind::Other, "server exited unexpectly");
-                    Err(err)
-                }
-                Err((err, ..)) => {
-                    error!("One of DNS servers exited unexpectly with error {}", err);
-                    Err(err)
-                }
-            })
-        })
-}
-
-fn handle_l2r(
-    context: SharedContext,
-    l: SharedUdpSocket,
-    server: Vec<(SharedUdpSocket, Arc<ServerConfig>, SocketAddr)>,
-) -> impl Future<Item = (), Error = io::Error> + Send {
-    assert!(!server.is_empty());
-
-    let mut server_idx: usize = 0;
-    let server = Arc::new(server);
-
-    PacketStream::new(l).for_each(move |(payload, src)| {
-        let server = server.clone();
-        let context = context.clone();
-        let context2 = context.clone();
-        let pkt_fut = futures::lazy(move || {
-            let pkt = Packet::parse(&payload[..]).map_err(|err| {
-                error!("Failed to parse DNS payload, err: {}", err);
-                io::Error::new(io::ErrorKind::Other, "parse DNS packet failed")
-            })?;
-
-            let (ref socket, ref svr_cfg, svr_addr) = server[server_idx % server.len()];
-            let (ni, _) = server_idx.overflowing_add(1);
-            server_idx = ni;
-
-            trace!("DNS {} -> {} packet: {:?}", src, svr_cfg.addr(), pkt);
-
-            let mut buf = Vec::new();
-            Address::SocketAddress(context.config().get_remote_dns()).write_to_buf(&mut buf);
-
-            buf.extend_from_slice(&payload);
-
-            let socket = socket.clone();
-            let id = pkt.header.id;
-            encrypt_payload(svr_cfg.method(), svr_cfg.key(), &buf)
-                .map(move |send_payload| (socket, svr_addr, send_payload, id))
-        })
-        .and_then(move |(socket, svr_addr, send_payload, id)| {
-            SendDgramRc::new(socket, send_payload, svr_addr).map(move |_| {
-                context2.dns_query_cache().insert(id, (src, Instant::now()));
-            })
-        });
-        tokio::spawn(pkt_fut.then(|res| match res {
-            Ok(..) => Ok(()),
-            Err(err) => {
-                error!("Failed to handle local -> remote packet, err: {}", err);
-                Err(())
-            }
-        }));
-
-        Ok(())
-    })
-}
-
-fn handle_r2l(
-    context: SharedContext,
-    l: SharedUdpSocket,
-    r: SharedUdpSocket,
-    svr_cfg: Arc<ServerConfig>,
-) -> impl Future<Item = (), Error = io::Error> + Send {
-    PacketStream::new(r).for_each(move |(payload, src)| {
-        let l = l.clone();
-        let svr_cfg = svr_cfg.clone();
-        let context = context.clone();
-        let pkt_fut = futures::lazy(move || decrypt_payload(svr_cfg.method(), svr_cfg.key(), &payload))
-            .and_then(move |payload| Address::read_from(Cursor::new(payload)).map_err(From::from))
-            .and_then(move |(cur, ..)| {
-                let pos = cur.position() as usize;
-                let payload = cur.into_inner();
-
-                let pkt = Packet::parse(&payload[pos..]).map_err(|err| {
-                    error!("Failed to parse DNS payload, err: {}", err);
-                    io::Error::new(io::ErrorKind::Other, "parse DNS packet failed")
-                })?;
-
-                let payload = payload[pos..].to_vec();
-
-                let mut cache = context.dns_query_cache();
-                match cache.remove(&pkt.header.id) {
-                    Some((cli_addr, start_time)) => {
-                        trace!(
-                            "DNS {} <- {} elapsed: {:?} packet: {:?}",
-                            cli_addr,
-                            src,
-                            Instant::now() - start_time,
-                            pkt,
-                        );
-
-                        Ok(Some((cli_addr, payload)))
-                    }
-                    None => {
-                        error!(
-                            "DNS received packet id={} opcode={:?} but found no local endpoint",
-                            pkt.header.id, pkt.header.opcode
-                        );
-                        Ok(None)
-                    }
-                }
-            })
-            .and_then(move |opt| match opt {
-                Some((cli_addr, payload)) => {
-                    let f = SendDgramRc::new(l, payload, cli_addr).map(|_| ());
-                    boxed_future(f)
-                }
-                None => boxed_future(futures::finished(())),
-            });
-
-        tokio::spawn(pkt_fut.then(|res| match res {
-            Ok(..) => Ok(()),
-            Err(err) => {
-                error!("Failed to handle local <- remote packet, err: {}", err);
-                Err(())
-            }
-        }));
-
-        Ok(())
-    })
-}

src/relay/udprelay/local.rs

@@ -3,178 +3,193 @@
 use std::{
     io::{self, Cursor, ErrorKind, Read},
     net::{IpAddr, Ipv4Addr, SocketAddr},
-    sync::{Arc, Mutex},
+    sync::Arc,
     time::Duration,
 };
 
-use futures::{self, Future, Stream};
-use log::{debug, error, info};
-use tokio::{self, net::UdpSocket, util::FutureExt};
+use bytes::BytesMut;
+use log::{debug, error, info, warn};
+use tokio::{self, net::UdpSocket, sync::mpsc};
 
 use crate::{
     config::{ServerAddr, ServerConfig},
     context::SharedContext,
     relay::{
-        boxed_future,
-        dns_resolver::resolve,
         loadbalancing::server::{LoadBalancer, RoundRobin},
         socks5::{Address, UdpAssociateHeader},
+        utils::try_timeout,
     },
 };
 
 use super::{
     crypto_io::{decrypt_payload, encrypt_payload},
-    PacketStream, SendDgramRc, MAXIMUM_UDP_PAYLOAD_SIZE,
+    MAXIMUM_UDP_PAYLOAD_SIZE,
 };
 
-/// Resolves server address to SocketAddr
-fn resolve_server_addr(
+#[allow(unused_variables)] // `context` is only used if trust-dns is enabled
+async fn udp_associate(
     context: SharedContext,
     svr_cfg: Arc<ServerConfig>,
-) -> impl Future<Item = SocketAddr, Error = io::Error> + Send {
-    match *svr_cfg.addr() {
-        // Return directly if it is a SocketAddr
-        ServerAddr::SocketAddr(ref addr) => boxed_future(futures::finished(*addr)),
-        // Resolve domain name to SocketAddr
+    pkt: Vec<u8>,
+    src: SocketAddr,
+) -> io::Result<Vec<u8>> {
+    const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5);
+
+    // PKT = UdpAssociateHeader + PAYLOAD
+    let mut cur = Cursor::new(pkt);
+
+    let header = UdpAssociateHeader::read_from(&mut cur).await?;
+
+    if header.frag != 0 {
+        error!("Received UDP associate with frag != 0, which is not supported by ShadowSocks");
+        let err = io::Error::new(ErrorKind::Other, "unsupported UDP fragmentation");
+        return Err(err);
+    }
+
+    let addr = header.address;
+
+    // The remaining is PAYLOAD
+    let mut payload = Vec::new();
+    Read::read_to_end(&mut cur, &mut payload)?;
+
+    // Binds to 0.0.0.0:0 (let system choose a random port)
+    // FIXME: Create a UdpSocket for every UDP associate requests
+    let local_addr = SocketAddr::new(IpAddr::from(Ipv4Addr::new(0, 0, 0, 0)), 0);
+    let mut remote_udp = UdpSocket::bind(&local_addr).await?;
+
+    // CLIENT -> SERVER protocol: ADDRESS + PAYLOAD
+    let mut send_buf = Vec::new();
+    addr.write_to_buf(&mut send_buf);
+    send_buf.extend_from_slice(&payload);
+
+    debug!(
+        "UDP ASSOCIATE {} -> {}, payload length {} bytes",
+        src,
+        addr,
+        payload.len()
+    );
+
+    let timeout = svr_cfg.udp_timeout().unwrap_or(DEFAULT_TIMEOUT);
+
+    // Write to remote socket CLIENT -> SERVER
+    let mut encrypt_buf = BytesMut::new();
+    encrypt_payload(svr_cfg.method(), svr_cfg.key(), &send_buf, &mut encrypt_buf)?;
+
+    let send_n = match *svr_cfg.addr() {
+        ServerAddr::SocketAddr(ref addr) => try_timeout(remote_udp.send_to(&encrypt_buf, addr), Some(timeout)).await?,
+        #[cfg(feature = "trust-dns")]
         ServerAddr::DomainName(ref dname, port) => {
-            let fut = resolve(context, dname, port, false).map(move |vec_ipaddr| {
-                assert!(!vec_ipaddr.is_empty());
-                vec_ipaddr[0]
-            });
-            boxed_future(fut)
+            use crate::relay::dns_resolver::resolve;
+
+            let vec_ipaddr = resolve(context, dname, port, false).await?;
+            assert!(!vec_ipaddr.is_empty());
+
+            try_timeout(remote_udp.send_to(&encrypt_buf, &vec_ipaddr[0]), Some(timeout)).await?
         }
+        #[cfg(not(feature = "trust-dns"))]
+        ServerAddr::DomainName(ref dname, port) => {
+            try_timeout(remote_udp.send_to(&encrypt_buf, (dname.as_str(), port)), Some(timeout)).await?
+        }
+    };
+
+    if send_n != encrypt_buf.len() {
+        warn!(
+            "Sent packet length {}, but expected length {}",
+            send_n,
+            encrypt_buf.len()
+        );
+    }
+
+    // Waiting for response from server SERVER -> CLIENT
+    // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
+    let mut recv_buf = [0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
+    let recv_n = try_timeout(remote_udp.recv(&mut recv_buf), Some(timeout)).await?;
+
+    let decrypt_buf = match decrypt_payload(svr_cfg.method(), svr_cfg.key(), &recv_buf[..recv_n])? {
+        None => {
+            error!("UDP packet too short, received length {}", recv_n);
+            let err = io::Error::new(io::ErrorKind::InvalidData, "packet too short");
+            return Err(err);
+        }
+        Some(b) => b,
+    };
+
+    // SERVER -> CLIENT protocol: ADDRESS + PAYLOAD
+    let mut cur = Cursor::new(decrypt_buf);
+
+    // FIXME: Address is ignored. Maybe useful in the future if we uses one common UdpSocket for communicate with remote server
+    let _ = Address::read_from(&mut cur).await?;
+
+    let payload_len = cur.get_ref().len() - cur.position() as usize;
+    debug!(
+        "UDP ASSOCIATE {} <- {}, payload length {} bytes",
+        src, addr, payload_len
+    );
+
+    let mut data = Vec::new();
+    UdpAssociateHeader::new(0, Address::SocketAddress(src)).write_to_buf(&mut data);
+
+    // Copy payload directly
+    Read::read_to_end(&mut cur, &mut data)?;
+
+    Ok(data)
+}
+
+async fn listen(context: SharedContext, l: UdpSocket) -> io::Result<()> {
+    let mut balancer = RoundRobin::new(context.config());
+
+    let (mut r, mut w) = l.split();
+
+    let mut pkt_buf = [0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
+
+    // FIXME: Channel size 1024?
+    let (tx, mut rx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);
+    tokio::spawn(async move {
+        while let Some((src, pkt)) = rx.recv().await {
+            if let Err(err) = w.send_to(&pkt, &src).await {
+                error!("UDP packet send failed, err: {:?}", err);
+                break;
+            }
+        }
+
+        // FIXME: How to stop the outer listener Future?
+    });
+
+    loop {
+        let (recv_len, src) = r.recv_from(&mut pkt_buf).await?;
+
+        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
+        // Copy bytes, because udp_associate runs in another tokio Task
+        let pkt = pkt_buf[..recv_len].to_vec();
+
+        let svr_cfg = balancer.pick_server();
+
+        let context = context.clone();
+        let svr_cfg = svr_cfg.clone();
+        let mut tx = tx.clone();
+        let assoc = async move {
+            match udp_associate(context, svr_cfg, pkt, src).await {
+                Ok(pkt) => {
+                    if let Err(..) = tx.send((src, pkt)).await {
+                        error!("UDP packet channel closed");
+                    }
+                }
+                Err(err) => {
+                    error!("Error occurs in UDP relay: {}", err);
+                }
+            }
+        };
+
+        tokio::spawn(assoc);
     }
 }
 
-fn listen(context: SharedContext, l: UdpSocket) -> impl Future<Item = (), Error = io::Error> + Send {
-    let socket = Arc::new(Mutex::new(l));
-    let mut balancer = RoundRobin::new(context.config());
-
-    PacketStream::new(socket.clone()).for_each(move |(pkt, src)| {
-        let svr_cfg = balancer.pick_server();
-        let svr_cfg_cloned = svr_cfg.clone();
-        let svr_cfg_cloned_cloned = svr_cfg.clone();
-        let socket = socket.clone();
-        let context = context.clone();
-        let timeout = *svr_cfg.udp_timeout();
-
-        const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5);
-
-        let rel = futures::lazy(|| UdpAssociateHeader::read_from(Cursor::new(pkt)))
-            .map_err(From::from)
-            .and_then(|(cur, header)| {
-                if header.frag != 0 {
-                    error!("Received UDP associate with frag != 0, which is not supported by ShadowSocks");
-                    let err = io::Error::new(ErrorKind::Other, "unsupported UDP fragmentation");
-                    Err(err)
-                } else {
-                    Ok((cur, header.address))
-                }
-            })
-            .and_then(|(mut cur, addr)| {
-                let svr_cfg = svr_cfg_cloned_cloned;
-
-                let mut payload = Vec::new();
-                cur.read_to_end(&mut payload).unwrap();
-
-                resolve_server_addr(context, svr_cfg)
-                    .and_then(|remote_addr| {
-                        let local_addr = SocketAddr::new(IpAddr::from(Ipv4Addr::new(0, 0, 0, 0)), 0);
-                        UdpSocket::bind(&local_addr).map(|remote_udp| (remote_udp, remote_addr))
-                    })
-                    .map(|(remote_udp, remote_addr)| (remote_udp, remote_addr, payload, addr))
-            })
-            .and_then(move |(remote_udp, remote_addr, payload, addr)| {
-                let mut buf = Vec::new();
-                addr.write_to_buf(&mut buf);
-                buf.extend_from_slice(&payload);
-                encrypt_payload(svr_cfg.method(), svr_cfg.key(), &buf)
-                    .map(|payload| (remote_udp, remote_addr, payload, addr))
-            })
-            .and_then(move |(remote_udp, remote_addr, payload, addr)| {
-                debug!(
-                    "UDP ASSOCIATE {} -> {}, payload length {} bytes",
-                    src,
-                    addr,
-                    payload.len()
-                );
-                let to = timeout.unwrap_or(DEFAULT_TIMEOUT);
-                let caddr = addr.clone();
-                remote_udp
-                    .send_dgram(payload, &remote_addr)
-                    .timeout(to)
-                    .map_err(move |err| match err.into_inner() {
-                        Some(e) => e,
-                        None => {
-                            error!(
-                                "Udp associate sending datagram {} -> {} timed out in {:?}",
-                                src, caddr, to
-                            );
-                            io::Error::new(io::ErrorKind::TimedOut, "udp send timed out")
-                        }
-                    })
-                    .map(|(remote_udp, _)| (remote_udp, addr))
-            })
-            .and_then(move |(remote_udp, addr)| {
-                let buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
-                let to = timeout.unwrap_or(DEFAULT_TIMEOUT);
-                let caddr = addr.clone();
-                remote_udp
-                    .recv_dgram(buf)
-                    .timeout(to)
-                    .map_err(move |err| match err.into_inner() {
-                        Some(e) => e,
-                        None => {
-                            error!(
-                                "Udp associate waiting datagram {} <- {} timed out in {:?}",
-                                src, caddr, to
-                            );
-                            io::Error::new(io::ErrorKind::TimedOut, "udp recv timed out")
-                        }
-                    })
-                    .and_then(move |(_remote_udp, buf, n, _from)| {
-                        let svr_cfg = svr_cfg_cloned;
-                        decrypt_payload(svr_cfg.method(), svr_cfg.key(), &buf[..n])
-                    })
-                    .map(|payload| (payload, addr))
-            })
-            .and_then(move |(payload, addr)| {
-                Address::read_from(Cursor::new(payload))
-                    .map_err(From::from)
-                    .map(|(cur, ..)| (cur, addr))
-            })
-            .and_then(move |(mut cur, addr)| {
-                let payload_len = cur.get_ref().len() - cur.position() as usize;
-                debug!(
-                    "UDP ASSOCIATE {} <- {}, payload length {} bytes",
-                    src, addr, payload_len
-                );
-
-                let mut data = Vec::new();
-                UdpAssociateHeader::new(0, Address::SocketAddress(src)).write_to_buf(&mut data);
-
-                cur.read_to_end(&mut data).unwrap();
-
-                SendDgramRc::new(socket, data, src)
-            })
-            .map(|_| ());
-
-        tokio::spawn(rel.map_err(|err| {
-            error!("Error occurs in UDP relay: {}", err);
-        }));
-
-        Ok(())
-    })
-}
-
 /// Starts a UDP local server
-pub fn run(context: SharedContext) -> impl Future<Item = (), Error = io::Error> + Send {
+pub async fn run(context: SharedContext) -> io::Result<()> {
     let local_addr = *context.config().local.as_ref().unwrap();
 
-    futures::lazy(move || {
-        info!("ShadowSocks UDP Listening on {}", local_addr);
+    let listener = UdpSocket::bind(&local_addr).await?;
+    info!("ShadowSocks UDP listening on {}", local_addr);
 
-        UdpSocket::bind(&local_addr)
-    })
-    .and_then(move |l| listen(context, l))
+    listen(context, listener).await
 }

src/relay/udprelay/mod.rs

@@ -37,17 +37,6 @@
 //! | Fixed |   Variable   |
 //! +-------+--------------+
 
-use std::{
-    io, mem,
-    net::SocketAddr,
-    sync::{Arc, Mutex},
-};
-
-use tokio::net::UdpSocket;
-
-use futures::{try_ready, Async, Future, Poll, Stream};
-
-pub mod dns;
 pub mod local;
 pub mod server;
 
@@ -59,74 +48,3 @@ mod crypto_io;
 /// Socks5 UDP ASSOCIATE packet size. The only thing I can find is
 /// [here](http://support.microsoft.com/kb/822061/)*
 pub const MAXIMUM_UDP_PAYLOAD_SIZE: usize = 65536;
-
-type SharedUdpSocket = Arc<Mutex<UdpSocket>>;
-
-/// UDP `recv_from` stream
-pub struct PacketStream {
-    udp: SharedUdpSocket,
-    buf: [u8; MAXIMUM_UDP_PAYLOAD_SIZE],
-}
-
-impl PacketStream {
-    /// Creates a new `PacketStream`
-    pub fn new(udp: SharedUdpSocket) -> PacketStream {
-        PacketStream {
-            udp,
-            buf: [0u8; MAXIMUM_UDP_PAYLOAD_SIZE],
-        }
-    }
-}
-
-impl Stream for PacketStream {
-    type Error = io::Error;
-    type Item = (Vec<u8>, SocketAddr);
-
-    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
-        let (n, addr) = try_ready!(self.udp.lock().unwrap().poll_recv_from(&mut self.buf));
-        Ok(Async::Ready(Some((self.buf[..n].to_vec(), addr))))
-    }
-}
-
-enum SendDgramStat<B: AsRef<[u8]>> {
-    Pending {
-        udp: SharedUdpSocket,
-        buf: B,
-        addr: SocketAddr,
-    },
-    Empty,
-}
-
-/// Send datagram with `Rc<UdpSocket>`
-pub struct SendDgramRc<B: AsRef<[u8]>> {
-    stat: SendDgramStat<B>,
-}
-
-impl<B: AsRef<[u8]>> SendDgramRc<B> {
-    pub fn new(udp: SharedUdpSocket, buf: B, addr: SocketAddr) -> SendDgramRc<B> {
-        SendDgramRc {
-            stat: SendDgramStat::Pending { udp, buf, addr },
-        }
-    }
-}
-
-impl<B: AsRef<[u8]>> Future for SendDgramRc<B> {
-    type Error = io::Error;
-    type Item = (SharedUdpSocket, usize, B);
-
-    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
-        let n = match self.stat {
-            SendDgramStat::Pending {
-                ref udp,
-                ref buf,
-                ref addr,
-            } => try_ready!(udp.lock().unwrap().poll_send_to(buf.as_ref(), addr)),
-            SendDgramStat::Empty => unreachable!(),
-        };
-
-        match mem::replace(&mut self.stat, SendDgramStat::Empty) {
-            SendDgramStat::Pending { udp, buf, .. } => Ok(Async::Ready((udp, n, buf))),
-            SendDgramStat::Empty => unreachable!(),
-        }
-    }
-}

src/relay/udprelay/server.rs

@@ -1,163 +1,183 @@
 //! UDP relay proxy server
 
 use std::{
-    io::{self, Cursor, ErrorKind},
+    io::{self, Cursor},
     net::{IpAddr, Ipv4Addr, SocketAddr},
-    sync::{Arc, Mutex},
+    sync::Arc,
     time::Duration,
 };
 
-use futures::{self, stream::futures_unordered, Future, Stream};
+use bytes::BytesMut;
+use futures::{self, stream::FuturesUnordered, StreamExt};
 use log::{debug, error, info};
-use tokio::{self, net::UdpSocket, util::FutureExt};
+use tokio::{self, net::UdpSocket, sync::mpsc};
 
 use crate::{
     config::ServerConfig,
     context::SharedContext,
-    relay::{boxed_future, dns_resolver::resolve, socks5::Address},
+    relay::{socks5::Address, utils::try_timeout},
 };
 
 use super::{
     crypto_io::{decrypt_payload, encrypt_payload},
-    PacketStream, SendDgramRc, MAXIMUM_UDP_PAYLOAD_SIZE,
+    MAXIMUM_UDP_PAYLOAD_SIZE,
 };
 
-fn resolve_remote_addr(
+#[allow(unused_variables)] // `context` is only used if trust-dns is enabled
+async fn udp_associate(
     context: SharedContext,
-    addr: Address,
-) -> impl Future<Item = SocketAddr, Error = io::Error> + Send {
-    match addr {
-        Address::SocketAddress(s) => {
-            if context.config().forbidden_ip.contains(&s.ip()) {
-                let err = io::Error::new(
-                    ErrorKind::Other,
-                    format!("{} is forbidden, failed to connect {}", s.ip(), s),
-                );
-                return boxed_future(futures::done(Err(err)));
+    svr_cfg: Arc<ServerConfig>,
+    decrypted_pkt: Vec<u8>,
+    src: SocketAddr,
+) -> io::Result<Vec<u8>> {
+    const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5);
+
+    // CLIENT -> SERVER protocol: ADDRESS + PAYLOAD
+    let mut cur = Cursor::new(decrypted_pkt);
+
+    let addr = Address::read_from(&mut cur).await?;
+
+    // Take out internal buffer for optimizing one byte copy
+    let header_len = cur.position() as usize;
+    let decrypted_pkt = cur.into_inner();
+    let body = &decrypted_pkt[header_len..];
+
+    debug!("UDP ASSOCIATE {} -> {}, payload length {} bytes", src, addr, body.len());
+
+    // FIXME: Create one UdpSocket for one associate
+    let local_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0);
+    let mut remote_udp = UdpSocket::bind(&local_addr).await?;
+
+    let timeout = svr_cfg.udp_timeout().unwrap_or(DEFAULT_TIMEOUT);
+
+    // Writes body to remote
+    let send_len = match addr {
+        Address::SocketAddress(ref remote_addr) => {
+            try_timeout(remote_udp.send_to(&body, remote_addr), Some(timeout)).await?
+        }
+        #[cfg(feature = "trust-dns")]
+        Address::DomainNameAddress(ref dname, port) => {
+            use crate::relay::dns_resolver::resolve;
+
+            let vec_ipaddr = resolve(context, dname, port, false).await?;
+            assert!(!vec_ipaddr.is_empty());
+
+            try_timeout(remote_udp.send_to(&body, &vec_ipaddr[0]), Some(timeout)).await?
+        }
+        #[cfg(not(feature = "trust-dns"))]
+        Address::DomainNameAddress(ref dname, port) => {
+            try_timeout(remote_udp.send_to(&body, (dname.as_str(), port)), Some(timeout)).await?
+        }
+    };
+    assert_eq!(body.len(), send_len);
+
+    // Waiting for response from server SERVER -> CLIENT
+    // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
+    let mut remote_buf = [0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
+    let remote_recv_len = try_timeout(remote_udp.recv(&mut remote_buf), Some(timeout)).await?;
+
+    // Making response packet, SERVER -> CLIENT: ADDRESS + PAYLOAD
+    let mut send_buf = Vec::new();
+    addr.write_to_buf(&mut send_buf);
+    send_buf.extend_from_slice(&remote_buf[..remote_recv_len]);
+
+    debug!(
+        "UDP ASSOCIATE {} <- {}, payload length {} bytes",
+        src,
+        addr,
+        send_buf.len()
+    );
+
+    Ok(send_buf)
+}
+
+async fn listen(context: SharedContext, svr_cfg: Arc<ServerConfig>) -> io::Result<()> {
+    let listen_addr = *svr_cfg.addr().listen_addr();
+    info!("ShadowSocks UDP listening on {}", listen_addr);
+
+    let listener = UdpSocket::bind(&listen_addr).await?;
+    let (mut r, mut w) = listener.split();
+
+    let svr_cfg_cloned = svr_cfg.clone();
+
+    // FIXME: Channel size 1024?
+    let (tx, mut rx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);
+    tokio::spawn(async move {
+        let svr_cfg = svr_cfg_cloned;
+
+        while let Some((src, pkt)) = rx.recv().await {
+            // Encrypts
+            let mut response_pkt = BytesMut::new();
+            if let Err(err) = encrypt_payload(svr_cfg.method(), svr_cfg.key(), &pkt, &mut response_pkt) {
+                error!("UDP packet encrypt failed, err: {:?}", err);
+                continue;
             }
 
-            boxed_future(futures::finished(s))
-        }
-        Address::DomainNameAddress(dname, port) => {
-            let fut = resolve(context, &dname, port, true).map(move |vec_ipaddr| {
-                assert!(!vec_ipaddr.is_empty());
-                vec_ipaddr[0]
-            });
-            boxed_future(fut)
+            if let Err(err) = w.send_to(&response_pkt, &src).await {
+                error!("UDP packet send failed, err: {:?}", err);
+                break;
+            }
         }
+
+        // FIXME: How to stop the outer listener Future?
+    });
+
+    let mut pkt_buf = [0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
+
+    loop {
+        let (recv_len, src) = r.recv_from(&mut pkt_buf).await?;
+
+        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
+        let pkt = &pkt_buf[..recv_len];
+
+        // First of all, decrypt payload CLIENT -> SERVER
+        let decrypted_pkt = match decrypt_payload(svr_cfg.method(), svr_cfg.key(), pkt) {
+            Ok(Some(pkt)) => pkt,
+            Ok(None) => {
+                error!("Failed to decrypt pkt in UDP relay, packet too short");
+                continue;
+            }
+            Err(err) => {
+                error!("Failed to decrypt pkt in UDP relay: {}", err);
+                continue;
+            }
+        };
+
+        let context = context.clone();
+        let svr_cfg = svr_cfg.clone();
+        let mut tx = tx.clone();
+        tokio::spawn(async move {
+            match udp_associate(context, svr_cfg, decrypted_pkt, src).await {
+                Ok(pkt) => {
+                    if let Err(..) = tx.send((src, pkt)).await {
+                        error!("UDP packet channel closed");
+                    }
+                }
+                Err(err) => {
+                    error!("Error occurs in UDP relay: {}", err);
+                }
+            }
+        });
     }
 }
 
-fn listen(context: SharedContext, svr_cfg: Arc<ServerConfig>) -> impl Future<Item = (), Error = io::Error> + Send {
-    let listen_addr = *svr_cfg.addr().listen_addr();
-    info!("ShadowSocks UDP listening on {}", listen_addr);
-    futures::lazy(move || UdpSocket::bind(&listen_addr)).and_then(move |socket| {
-        let socket = Arc::new(Mutex::new(socket));
-        PacketStream::new(socket.clone()).for_each(move |(pkt, src)| {
-            let svr_cfg = svr_cfg.clone();
-            let svr_cfg_cloned = svr_cfg.clone();
-            let socket = socket.clone();
-            let context = context.clone();
-            let timeout = svr_cfg.timeout();
-            let rel = futures::lazy(move || decrypt_payload(svr_cfg.method(), svr_cfg.key(), &pkt))
-                .and_then(move |payload| {
-                    // Read Address in the front (ShadowSocks protocol)
-                    Address::read_from(Cursor::new(payload))
-                        .map_err(From::from)
-                        .and_then(move |(r, addr)| {
-                            let header_len = r.position() as usize;
-                            let mut payload = r.into_inner();
-                            payload.drain(..header_len);
-                            let body = payload;
-
-                            debug!("UDP ASSOCIATE {} -> {}, payload length {} bytes", src, addr, body.len());
-                            Ok((addr, body))
-                        })
-                        .and_then(|(addr, body)| {
-                            let local_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0);
-                            UdpSocket::bind(&local_addr).map(|remote_udp| (remote_udp, addr, body))
-                        })
-                        .and_then(|(remote_udp, addr, body)| {
-                            resolve_remote_addr(context, addr.clone())
-                                .and_then(|addr| remote_udp.send_dgram(body, &addr))
-                                .map(|(remote_udp, _)| (remote_udp, addr))
-                        })
-                })
-                .and_then(move |(remote_udp, addr)| {
-                    let buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];
-                    let to = timeout.unwrap_or(Duration::from_secs(5));
-                    let caddr = addr.clone();
-                    remote_udp
-                        .recv_dgram(buf)
-                        .timeout(to)
-                        .map_err(move |err| match err.into_inner() {
-                            Some(e) => e,
-                            None => {
-                                error!(
-                                    "Udp associate waiting datagram {} -> {} timed out in {:?}",
-                                    src, caddr, to
-                                );
-                                io::Error::new(io::ErrorKind::TimedOut, "udp recv timed out")
-                            }
-                        })
-                        .and_then(|(_remote_udp, buf, n, _from)| {
-                            let svr_cfg = svr_cfg_cloned;
-
-                            let mut send_buf = Vec::new();
-                            addr.write_to_buf(&mut send_buf);
-                            send_buf.extend_from_slice(&buf[..n]);
-                            encrypt_payload(svr_cfg.method(), svr_cfg.key(), &send_buf).map(|buf| (buf, addr))
-                        })
-                })
-                .and_then(move |(buf, addr)| {
-                    debug!("UDP ASSOCIATE {} <- {}, payload length {} bytes", src, addr, buf.len());
-
-                    let to = timeout.unwrap_or(Duration::from_secs(5));
-                    let caddr = addr.clone();
-                    SendDgramRc::new(socket, buf, src)
-                        .timeout(to)
-                        .map_err(move |err| match err.into_inner() {
-                            Some(e) => e,
-                            None => {
-                                error!(
-                                    "Udp associate sending datagram {} <- {} timed out in {:?}",
-                                    src, caddr, to
-                                );
-                                io::Error::new(io::ErrorKind::TimedOut, "udp send timed out")
-                            }
-                        })
-                })
-                .map(|_| ());
-
-            tokio::spawn(rel.map_err(|err| {
-                error!("Udp relay error: {}", err);
-            }));
-
-            Ok(())
-        })
-    })
-}
-
 /// Starts a UDP relay server
-pub fn run(context: SharedContext) -> impl Future<Item = (), Error = io::Error> + Send {
-    let mut vec_fut = Vec::new();
+pub async fn run(context: SharedContext) -> io::Result<()> {
+    let vec_fut = FuturesUnordered::new();
 
     for svr in &context.config().server {
         let svr_cfg = Arc::new(svr.clone());
 
         let svr_fut = listen(context.clone(), svr_cfg);
-        vec_fut.push(boxed_future(svr_fut));
+        vec_fut.push(svr_fut);
     }
 
-    futures_unordered(vec_fut).into_future().then(|res| match res {
-        Ok(..) => {
-            error!("One of UDP servers exited unexpectly without error");
+    match vec_fut.into_future().await.0 {
+        Some(res) => {
+            error!("One of TCP servers exited unexpectly, result: {:?}", res);
             let err = io::Error::new(io::ErrorKind::Other, "server exited unexpectly");
             Err(err)
         }
-        Err((err, ..)) => {
-            error!("One of UDP servers exited unexpectly with error {}", err);
-            Err(err)
-        }
-    })
+        None => unreachable!(),
+    }
 }

src/relay/utils.rs

@@ -1,47 +1,16 @@
-use std::io;
+use std::future::Future;
+use std::{io, time::Duration};
 
-use futures::{try_ready, Async, Future, Poll};
-use tokio_io::{
-    io::{write_all, WriteAll},
-    AsyncWrite,
-};
+use tokio::time;
 
-/// Write all bytes without returning the internal bytes buffer
-pub struct WriteBytes<W, B>
+pub async fn try_timeout<T, E, F>(fut: F, timeout: Option<Duration>) -> io::Result<T>
 where
-    W: AsyncWrite,
-    B: AsRef<[u8]>,
+    F: Future<Output = Result<T, E>>,
+    io::Error: From<E>,
 {
-    inner: WriteAll<W, B>,
-}
-
-impl<W, B> WriteBytes<W, B>
-where
-    W: AsyncWrite,
-    B: AsRef<[u8]>,
-{
-    fn new(writer: W, bytes: B) -> WriteBytes<W, B> {
-        WriteBytes {
-            inner: write_all(writer, bytes),
-        }
+    match timeout {
+        Some(t) => time::timeout(t, fut).await?,
+        None => fut.await,
     }
-}
-
-impl<W, B> Future for WriteBytes<W, B>
-where
-    W: AsyncWrite,
-    B: AsRef<[u8]>,
-{
-    type Error = io::Error;
-    type Item = W;
-
-    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
-        let (w, _) = try_ready!(self.inner.poll());
-        Ok(Async::Ready(w))
-    }
-}
-
-/// Write all bytes without returning the internal bytes buffer
-pub fn write_bytes<W: AsyncWrite, B: AsRef<[u8]>>(w: W, b: B) -> WriteBytes<W, B> {
-    WriteBytes::new(w, b)
+    .map_err(From::from)
 }

tests/dns.rs

@@ -1,84 +0,0 @@
-use std::{
-    collections::HashSet,
-    net::{SocketAddr, UdpSocket},
-    thread,
-    time::Duration,
-};
-
-use dns_parser::{Builder, Packet, QueryClass, QueryType};
-use log::trace;
-use shadowsocks::{
-    config::{Config, ConfigType},
-    run_dns, run_server,
-};
-use tokio::runtime::current_thread::Runtime;
-
-const CONFIG: &str = r#"{
-        "server": "127.0.0.1",
-        "server_port": 8988,
-        "local_port": 5030,
-        "local_address": "127.0.0.1",
-        "password": "abc",
-        "timeout": 20,
-        "method": "aes-256-gcm",
-        "mode": "tcp_and_udp",
-        "remote_dns": "114.114.114.114:53"
-    }"#;
-
-#[test]
-fn dns_relay() {
-    let _ = env_logger::try_init();
-
-    let server_cfg = Config::load_from_str(CONFIG, ConfigType::Server).unwrap();
-    let dns_cfg = Config::load_from_str(CONFIG, ConfigType::Local).unwrap();
-
-    trace!("Server: {:?}", server_cfg);
-    trace!("Client: {:?}", dns_cfg);
-
-    thread::spawn(move || {
-        let mut runtime = Runtime::new().expect("Failed to create Runtime");
-        runtime.block_on(run_server(server_cfg)).unwrap();
-    });
-
-    thread::spawn(move || {
-        let mut runtime = Runtime::new().expect("Failed to create Runtime");
-        runtime.block_on(run_dns(dns_cfg)).unwrap();
-    });
-
-    thread::sleep(Duration::from_secs(1));
-
-    let dns_addr = "127.0.0.1:5030".parse::<SocketAddr>().unwrap();
-    let local_addr = "0.0.0.0:0".parse::<SocketAddr>().unwrap();
-    let local = UdpSocket::bind(&local_addr).unwrap();
-    local.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
-    local.set_write_timeout(Some(Duration::from_secs(5))).unwrap();
-
-    const LOOP_ROUND: usize = 10;
-
-    let mut sent_id = HashSet::<u16, _>::with_capacity(LOOP_ROUND);
-    for _ in 0..LOOP_ROUND {
-        let mut id = rand::random::<u16>();
-        while let Some(..) = sent_id.get(&id) {
-            id = rand::random::<u16>();
-        }
-
-        sent_id.insert(id);
-
-        let mut builder = Builder::new_query(id, false);
-        builder.add_question("www.example.com", false, QueryType::A, QueryClass::IN);
-
-        let payload = builder.build().unwrap();
-
-        local.send_to(&payload, &dns_addr).unwrap();
-
-        trace!("DNS SENT {:?}", payload);
-
-        let mut buf = [0u8; 65535];
-        let (len, _) = local.recv_from(&mut buf).unwrap();
-
-        let packet = Packet::parse(&buf[..len]).unwrap();
-
-        trace!("DNS GOT {:?}", packet);
-        assert_eq!(packet.header.id, id);
-    }
-}

tests/socks5.rs

@@ -1,12 +1,7 @@
-use std::{
-    net::{SocketAddr, ToSocketAddrs},
-    thread,
-    time::Duration,
-};
+use std::net::{SocketAddr, ToSocketAddrs};
 
-use futures::Future;
-use tokio::runtime::current_thread::Runtime;
-use tokio_io::io::{flush, read_to_end, write_all};
+use tokio::prelude::*;
+use tokio::time::{self, Duration};
 
 use shadowsocks::{
     config::{Config, ConfigType, Mode, ServerConfig},
@@ -52,27 +47,19 @@ impl Socks5TestServer {
         &self.local_addr
     }
 
-    pub fn run(&self) {
+    pub async fn run(&self) {
         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");
-        });
+        tokio::spawn(run_server(svr_cfg));
 
         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);
-            runtime.block_on(fut).expect("Failed to run Local");
-        });
+        tokio::spawn(run_local(client_cfg));
 
-        thread::sleep(Duration::from_secs(1));
+        time::delay_for(Duration::from_secs(1)).await;
     }
 }
 
-#[test]
-fn socks5_relay_stream() {
+#[tokio::test]
+async fn socks5_relay_stream() {
     let _ = env_logger::try_init();
 
     const SERVER_ADDR: &str = "127.0.0.1:8100";
@@ -82,29 +69,27 @@ fn socks5_relay_stream() {
     const METHOD: CipherType = CipherType::Aes256Cfb;
 
     let svr = Socks5TestServer::new(SERVER_ADDR, LOCAL_ADDR, PASSWORD, METHOD, false);
-    svr.run();
+    svr.run().await;
 
-    let c = Socks5Client::connect(
+    let mut c = Socks5Client::connect(
         Address::DomainNameAddress("www.example.com".to_owned(), 80),
-        *svr.client_addr(),
-    );
+        svr.client_addr(),
+    )
+    .await
+    .unwrap();
 
-    let fut = c.and_then(|c| {
-        let req = b"GET / HTTP/1.0\r\nHost: www.example.com\r\nAccept: */*\r\n\r\n";
-        write_all(c, req.to_vec())
-            .and_then(|(c, _)| flush(c))
-            .and_then(|c| read_to_end(c, Vec::new()))
-            .map(|(_, buf)| {
-                println!("Got reply from server: {}", String::from_utf8(buf).unwrap());
-            })
-    });
+    let req = b"GET / HTTP/1.0\r\nHost: www.example.com\r\nAccept: */*\r\n\r\n";
+    c.write_all(req).await.unwrap();
+    c.flush().await.unwrap();
 
-    let mut runtime = Runtime::new().expect("Failed to create Runtime");
-    runtime.block_on(fut).unwrap();
+    let mut buf = Vec::new();
+    c.read_to_end(&mut buf).await.unwrap();
+
+    println!("Got reply from server: {}", String::from_utf8(buf).unwrap());
 }
 
-#[test]
-fn socks5_relay_aead() {
+#[tokio::test]
+async fn socks5_relay_aead() {
     let _ = env_logger::try_init();
 
     const SERVER_ADDR: &str = "127.0.0.1:8110";
@@ -114,22 +99,21 @@ fn socks5_relay_aead() {
     const METHOD: CipherType = CipherType::Aes256Gcm;
 
     let svr = Socks5TestServer::new(SERVER_ADDR, LOCAL_ADDR, PASSWORD, METHOD, false);
-    svr.run();
+    svr.run().await;
 
-    let c = Socks5Client::connect(
+    let mut c = Socks5Client::connect(
         Address::DomainNameAddress("www.example.com".to_owned(), 80),
-        *svr.client_addr(),
-    );
-    let fut = c.and_then(|c| {
-        let req = b"GET / HTTP/1.0\r\nHost: www.example.com\r\nAccept: */*\r\n\r\n";
-        write_all(c, req.to_vec())
-            .and_then(|(c, _)| flush(c))
-            .and_then(|c| read_to_end(c, Vec::new()))
-            .map(|(_, buf)| {
-                println!("Got reply from server: {}", String::from_utf8(buf).unwrap());
-            })
-    });
+        svr.client_addr(),
+    )
+    .await
+    .unwrap();
 
-    let mut runtime = Runtime::new().expect("Failed to create Runtime");
-    runtime.block_on(fut).unwrap();
+    let req = b"GET / HTTP/1.0\r\nHost: www.example.com\r\nAccept: */*\r\n\r\n";
+    c.write_all(req).await.unwrap();
+    c.flush().await.unwrap();
+
+    let mut buf = Vec::new();
+    c.read_to_end(&mut buf).await.unwrap();
+
+    println!("Got reply from server: {}", String::from_utf8(buf).unwrap());
 }

tests/tunnel.rs

@@ -0,0 +1,57 @@
+use env_logger;
+use tokio;
+use tokio::net::TcpStream;
+use tokio::prelude::*;
+use tokio::time::{self, Duration};
+
+use shadowsocks::config::{Config, ConfigType};
+use shadowsocks::relay::socks5::Address;
+use shadowsocks::{run_local, run_server};
+
+#[tokio::test]
+async fn tcp_tunnel() {
+    let _ = env_logger::try_init();
+
+    let mut local_config = Config::load_from_str(
+        r#"{
+            "local_port": 9110,
+            "local_address": "127.0.0.1",
+            "server": "127.0.0.1",
+            "server_port": 9120,
+            "password": "password",
+            "method": "aes-256-gcm"
+        }"#,
+        ConfigType::Local,
+    )
+    .unwrap();
+
+    local_config.forward = Some("www.example.com:80".parse::<Address>().unwrap());
+
+    let server_config = Config::load_from_str(
+        r#"{
+            "server": "127.0.0.1",
+            "server_port": 9120,
+            "password": "password",
+            "method": "aes-256-gcm"
+        }"#,
+        ConfigType::Server,
+    )
+    .unwrap();
+
+    tokio::spawn(run_local(local_config));
+    tokio::spawn(run_server(server_config));
+
+    time::delay_for(Duration::from_secs(1)).await;
+
+    // Connect it directly, because it is now established a TCP tunnel with www.example.com
+    let mut stream = TcpStream::connect("127.0.0.1:9110").await.unwrap();
+
+    let req = b"GET / HTTP/1.0\r\nHost: www.example.com\r\nAccept: */*\r\n\r\n";
+    stream.write_all(req).await.unwrap();
+    stream.flush().await.unwrap();
+
+    let mut buf = Vec::new();
+    stream.read_to_end(&mut buf).await.unwrap();
+
+    println!("Got reply from server: {}", String::from_utf8(buf).unwrap());
+}

tests/udp.rs

@@ -1,17 +1,14 @@
 #![cfg_attr(clippy, allow(blacklisted_name))]
 
 use std::{
-    io::Cursor,
+    io::{self, Cursor},
     net::SocketAddr,
-    sync::{Arc, Barrier},
-    thread,
-    time::Duration,
 };
 
 use bytes::{BufMut, BytesMut};
-use futures::Future;
-use tokio::runtime::current_thread::Runtime;
-use tokio_io::io::read_to_end;
+use log::debug;
+use tokio::prelude::*;
+use tokio::time::{self, Duration};
 
 use shadowsocks::{
     config::{Config, ConfigType, Mode, ServerConfig},
@@ -59,88 +56,67 @@ fn get_client_addr() -> SocketAddr {
     LOCAL_ADDR.parse().unwrap()
 }
 
-fn start_server(bar: Arc<Barrier>) {
-    thread::spawn(move || {
-        let mut runtime = Runtime::new().expect("Failed to create Runtime");
+fn start_server() {
+    tokio::spawn(run_server(get_svr_config()));
+}
 
-        let fut = run_server(get_svr_config());
-        bar.wait();
-        runtime.block_on(fut).expect("Failed to run Server");
+fn start_local() {
+    tokio::spawn(run_local(get_cli_config()));
+}
+
+fn start_udp_echo_server() {
+    use tokio::net::UdpSocket;
+
+    tokio::spawn(async {
+        let mut l = UdpSocket::bind(UDP_ECHO_SERVER_ADDR).await.unwrap();
+
+        debug!("UDP echo server started {}", UDP_ECHO_SERVER_ADDR);
+
+        let mut buf = vec![0u8; 65536];
+        let (amt, src) = l.recv_from(&mut buf).await.unwrap();
+
+        debug!("UDP echo received {} bytes from {}", amt, src);
+
+        l.send_to(&buf[..amt], &src).await.unwrap();
+
+        debug!("UDP echo sent {} bytes to {}", amt, src);
     });
 }
 
-fn start_local(bar: Arc<Barrier>) {
-    thread::spawn(move || {
-        let mut runtime = Runtime::new().expect("Failed to create Runtime");
+fn start_udp_request_holder(addr: Address) {
+    tokio::spawn(async move {
+        let (mut c, addr) = Socks5Client::udp_associate(addr, &get_client_addr()).await?;
+        assert_eq!(addr, Address::SocketAddress(LOCAL_ADDR.parse().unwrap()));
 
-        let fut = run_local(get_cli_config());
-        bar.wait();
-        runtime.block_on(fut).expect("Failed to run Local");
+        debug!("TCP sent UDP associate {} request", addr);
+
+        // Holds it forever
+        let mut buf = Vec::new();
+        c.read_to_end(&mut buf).await?;
+
+        io::Result::Ok(())
     });
 }
 
-fn start_udp_echo_server(bar: Arc<Barrier>) {
-    use std::net::UdpSocket;
-
-    thread::spawn(move || {
-        let l = UdpSocket::bind(UDP_ECHO_SERVER_ADDR).unwrap();
-
-        bar.wait();
-
-        let mut buf = [0u8; 65536];
-        let (amt, src) = l.recv_from(&mut buf).unwrap();
-
-        l.send_to(&buf[..amt], &src).unwrap();
-    });
-}
-
-fn start_udp_request_holder(bar: Arc<Barrier>, addr: Address) {
-    thread::spawn(move || {
-        let mut runtime = Runtime::new().expect("Failed to create Runtime");
-
-        let c = Socks5Client::udp_associate(addr, get_client_addr());
-        let fut = c.and_then(|(c, addr)| {
-            assert_eq!(addr, Address::SocketAddress(LOCAL_ADDR.parse().unwrap()));
-
-            // Holds it forever
-            read_to_end(c, Vec::new()).map(|_| ())
-        });
-
-        bar.wait();
-
-        runtime.block_on(fut).expect("Failed to run UDP socks5 client");
-    });
-}
-
-#[test]
-fn udp_relay() {
-    use std::net::UdpSocket;
+#[tokio::test]
+async fn udp_relay() {
+    use tokio::net::UdpSocket;
 
     let _ = env_logger::try_init();
 
     let remote_addr = Address::SocketAddress(UDP_ECHO_SERVER_ADDR.parse().unwrap());
 
-    let bar = Arc::new(Barrier::new(4));
+    start_server();
+    start_local();
 
-    start_server(bar.clone());
-    start_local(bar.clone());
-
-    start_udp_echo_server(bar.clone());
-
-    bar.wait();
+    start_udp_echo_server();
 
     // Wait until all server starts
-    thread::sleep(Duration::from_secs(1));
+    time::delay_for(Duration::from_secs(1)).await;
 
-    let bar = Arc::new(Barrier::new(2));
+    start_udp_request_holder(remote_addr.clone());
 
-    start_udp_request_holder(bar.clone(), remote_addr.clone());
-
-    bar.wait();
-
-    let l = UdpSocket::bind(UDP_LOCAL_ADDR).unwrap();
-    l.set_read_timeout(Some(Duration::from_secs(5))).unwrap();
-    l.set_write_timeout(Some(Duration::from_secs(5))).unwrap();
+    let mut l = UdpSocket::bind(UDP_LOCAL_ADDR).await.unwrap();
 
     let header = UdpAssociateHeader::new(0, remote_addr);
     let mut buf = BytesMut::with_capacity(header.serialized_len());
@@ -152,14 +128,17 @@ fn udp_relay() {
     buf.put_slice(payload);
 
     let local_addr = LOCAL_ADDR.parse::<SocketAddr>().unwrap();
-    l.send_to(&buf[..], &local_addr).unwrap();
+    l.send_to(&buf[..], &local_addr).await.unwrap();
 
-    let mut buf = [0u8; 65536];
-    let (amt, _) = l.recv_from(&mut buf).unwrap();
+    let mut buf = vec![0u8; 65536];
+    let (amt, _) = time::timeout(Duration::from_secs(5), l.recv_from(&mut buf))
+        .await
+        .unwrap()
+        .unwrap();
     println!("Received buf size={} {:?}", amt, &buf[..amt]);
 
-    let cur = Cursor::new(buf[..amt].to_vec());
-    let (cur, header) = UdpAssociateHeader::read_from(cur).wait().expect("Invalid UDP header");
+    let mut cur = Cursor::new(buf[..amt].to_vec());
+    let header = UdpAssociateHeader::read_from(&mut cur).await.unwrap();
     println!("{:?}", header);
     let header_len = cur.position() as usize;
     let buf = cur.into_inner();