From 6a1a0213316239271139e88e3bbc98e11914930a Mon Sep 17 00:00:00 2001
From: "Y. T. Chung" <zonyitoo@gmail.com>
Date: Sat, 18 Feb 2017 21:55:13 +0800
Subject: [PATCH] add aead test

---
 tests/socks5.rs | 127 ++++++++++++++++++++++++++++++++----------------
 1 file changed, 85 insertions(+), 42 deletions(-)

diff --git a/tests/socks5.rs b/tests/socks5.rs
index c540d2a5..93149d02 100644
--- a/tests/socks5.rs
+++ b/tests/socks5.rs
@@ -1,5 +1,3 @@
-#![allow(dead_code, unused_imports)]
-
 extern crate shadowsocks;
 extern crate tokio_core;
 extern crate futures;
@@ -8,10 +6,8 @@ extern crate log;
 extern crate env_logger;
 
 use std::thread;
-use std::net::SocketAddr;
-use std::sync::{Arc, Barrier};
+use std::net::{SocketAddr, ToSocketAddrs};
 use std::time::Duration;
-use std::io::Cursor;
 
 use tokio_core::reactor::Core;
 use tokio_core::io::{read_to_end, write_all, flush};
@@ -23,58 +19,105 @@ use shadowsocks::crypto::CipherType;
 use shadowsocks::relay::socks5::Address;
 use shadowsocks::{run_local, run_server};
 
-const SERVER_ADDR: &'static str = "127.0.0.1:8092";
-const LOCAL_ADDR: &'static str = "127.0.0.1:8290";
-
-const PASSWORD: &'static str = "test-password";
-const METHOD: CipherType = CipherType::Aes128Cfb;
-
-fn get_config() -> Config {
-    let mut cfg = Config::new();
-    cfg.local = Some(LOCAL_ADDR.parse().unwrap());
-    cfg.server = vec![ServerConfig::basic(SERVER_ADDR.parse().unwrap(), PASSWORD.to_owned(), METHOD)];
-    cfg
+pub struct Socks5TestServer {
+    local_addr: SocketAddr,
+    config: Config,
 }
 
-fn get_client_addr() -> SocketAddr {
-    LOCAL_ADDR.parse().unwrap()
-}
+impl Socks5TestServer {
+    pub fn new<S, L>(svr_addr: S,
+                     local_addr: L,
+                     pwd: &'static str,
+                     method: CipherType,
+                     enable_udp: bool)
+                     -> Socks5TestServer
+        where S: ToSocketAddrs,
+              L: ToSocketAddrs
+    {
+        let svr_addr = svr_addr.to_socket_addrs().unwrap().next().unwrap();
+        let local_addr = local_addr.to_socket_addrs().unwrap().next().unwrap();
 
-fn start_server(bar: Arc<Barrier>) {
-    thread::spawn(move || {
+        Socks5TestServer {
+            local_addr: local_addr.clone(),
+            config: {
+                let mut cfg = Config::new();
+                cfg.local = Some(local_addr);
+                cfg.server = vec![ServerConfig::basic(svr_addr, pwd.to_owned(), method)];
+                cfg.enable_udp = enable_udp;
+                cfg
+            },
+        }
+    }
+
+    pub fn client_addr(&self) -> &SocketAddr {
+        &self.local_addr
+    }
+
+    pub fn run(&self) {
         drop(env_logger::init());
-        bar.wait();
-        run_server(get_config()).unwrap();
-    });
-}
 
-fn start_local(bar: Arc<Barrier>) {
-    thread::spawn(move || {
-        drop(env_logger::init());
-        bar.wait();
-        run_local(get_config()).unwrap();
-    });
+        let svr_cfg = self.config.clone();
+        thread::spawn(move || {
+            drop(env_logger::init());
+            run_server(svr_cfg).unwrap();
+        });
+
+        let client_cfg = self.config.clone();
+        thread::spawn(move || {
+            drop(env_logger::init());
+            run_local(client_cfg).unwrap();
+        });
+
+        thread::sleep(Duration::from_secs(1));
+    }
 }
 
 #[test]
-fn socks5_relay() {
-    drop(env_logger::init());
+fn socks5_relay_stream() {
+    const SERVER_ADDR: &'static str = "127.0.0.1:8100";
+    const LOCAL_ADDR: &'static str = "127.0.0.1:8200";
 
-    let bar = Arc::new(Barrier::new(3));
+    const PASSWORD: &'static str = "test-password";
+    const METHOD: CipherType = CipherType::Aes256Cfb;
 
-    start_server(bar.clone());
-    start_local(bar.clone());
-
-    bar.wait();
-
-    // Wait until all server starts
-    thread::sleep(Duration::from_secs(1));
+    let svr = Socks5TestServer::new(SERVER_ADDR, LOCAL_ADDR, PASSWORD, METHOD, false);
+    svr.run();
 
     let mut lp = Core::new().unwrap();
     let handle = lp.handle();
 
     let c = Socks5Client::connect(Address::DomainNameAddress("www.example.com".to_owned(), 80),
-                                  get_client_addr(),
+                                  *svr.client_addr(),
+                                  handle);
+    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());
+            })
+    });
+
+    lp.run(fut).unwrap();
+}
+
+#[test]
+fn socks5_relay_aead() {
+    const SERVER_ADDR: &'static str = "127.0.0.1:8110";
+    const LOCAL_ADDR: &'static str = "127.0.0.1:8210";
+
+    const PASSWORD: &'static str = "test-password";
+    const METHOD: CipherType = CipherType::Aes256Gcm;
+
+    let svr = Socks5TestServer::new(SERVER_ADDR, LOCAL_ADDR, PASSWORD, METHOD, false);
+    svr.run();
+
+    let mut lp = Core::new().unwrap();
+    let handle = lp.handle();
+
+    let c = Socks5Client::connect(Address::DomainNameAddress("www.example.com".to_owned(), 80),
+                                  *svr.client_addr(),
                                   handle);
     let fut = c.and_then(|c| {
         let req = b"GET / HTTP/1.0\r\nHost: www.example.com\r\nAccept: */*\r\n\r\n";