From ee9d0685b323abdb712e4ce1c68703f17bad3b61 Mon Sep 17 00:00:00 2001
From: KernelErr <45716019+KernelErr@users.noreply.github.com>
Date: Tue, 26 Oct 2021 23:52:07 +0800
Subject: [PATCH] Refactor TCP and KCP test

---
 src/config.rs      |  2 +-
 src/servers/mod.rs | 46 +++++++++++++++++++++++++++-------------------
 tests/config.yaml  | 13 +++++++++----
 3 files changed, 37 insertions(+), 24 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 9a5d10a..ba2424d 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -87,7 +87,7 @@ mod tests {
         let config = Config::new("tests/config.yaml").unwrap();
         assert_eq!(config.base.version, 1);
         assert_eq!(config.base.log.unwrap(), "disable");
-        assert_eq!(config.base.servers.len(), 4);
+        assert_eq!(config.base.servers.len(), 5);
         assert_eq!(config.base.upstream.len(), 3);
     }
 }
diff --git a/src/servers/mod.rs b/src/servers/mod.rs
index 3b72e97..80316ff 100644
--- a/src/servers/mod.rs
+++ b/src/servers/mod.rs
@@ -118,8 +118,11 @@ mod test {
         let server_addr: SocketAddr = "127.0.0.1:54599".parse().unwrap();
         let listener = TcpListener::bind(server_addr).await.unwrap();
         let (mut stream, _) = listener.accept().await.unwrap();
-        stream.write(b"hello").await.unwrap();
-        stream.shutdown().await.unwrap();
+        let mut buf = [0u8; 1024];
+        let n = stream.read(&mut buf).await.unwrap();
+        if n > 0 {
+            stream.write(b"hello").await.unwrap();
+        }
     }
 
     #[tokio::test]
@@ -135,25 +138,19 @@ mod test {
             let _ = server.run();
         });
         sleep(Duration::from_secs(1)); // wait for server to start
-        let mut conn = TcpStream::connect("127.0.0.1:54500").await.unwrap();
-        let mut buf = [0u8; 5];
-        conn.read(&mut buf).await.unwrap();
-        assert_eq!(&buf, b"hello");
-        conn.shutdown().await.unwrap();
-    }
 
-    #[tokio::test]
-    async fn test_tcp_echo_server() {
-        use crate::config::Config;
-        let config = Config::new("tests/config.yaml").unwrap();
-        let mut server = Server::new(config.base);
-        thread::spawn(move || {
-            let _ = server.run();
-        });
-        sleep(Duration::from_secs(1)); // wait for server to start
+        // // test proxy
+        // let mut conn = TcpStream::connect("127.0.0.1:54500").await.unwrap();
+        // let mut buf = [0u8; 5];
+        // conn.write(b"hi").await.unwrap();
+        // conn.read(&mut buf).await.unwrap();
+        // assert_eq!(&buf, b"hello");
+        // conn.shutdown().await.unwrap();
+
+        // test echo
         let mut conn = TcpStream::connect("127.0.0.1:54956").await.unwrap();
         let mut buf = [0u8; 1];
-        for i in 0..=255u8 {
+        for i in 0..=10u8 {
             conn.write(&[i]).await.unwrap();
             conn.read(&mut buf).await.unwrap();
             assert_eq!(&buf, &[i]);
@@ -162,7 +159,7 @@ mod test {
     }
 
     #[tokio::test]
-    async fn test_kcp_echo_server() {
+    async fn test_kcp_proxy() {
         use crate::config::Config;
         let config = Config::new("tests/config.yaml").unwrap();
         let mut server = Server::new(config.base);
@@ -170,6 +167,17 @@ mod test {
             let _ = server.run();
         });
         sleep(Duration::from_secs(1)); // wait for server to start
+
+        // test proxy
+        let kcp_config = KcpConfig::default();
+        let server_addr: SocketAddr = "127.0.0.1:54958".parse().unwrap();
+        let mut conn = KcpStream::connect(&kcp_config, server_addr).await.unwrap();
+        let mut buf = [0u8; 5];
+        conn.write(b"hi").await.unwrap();
+        conn.read(&mut buf).await.unwrap();
+        assert_eq!(&buf, b"hello");
+
+        // test echo
         let kcp_config = KcpConfig::default();
         let server_addr: SocketAddr = "127.0.0.1:54959".parse().unwrap();
         let mut conn = KcpStream::connect(&kcp_config, server_addr).await.unwrap();
diff --git a/tests/config.yaml b/tests/config.yaml
index b906374..fb3f1cd 100644
--- a/tests/config.yaml
+++ b/tests/config.yaml
@@ -11,15 +11,20 @@ servers:
       proxy.test.com: proxy
       www.test.com: web
     default: ban
-  echo_server:
-    listen:
-      - "0.0.0.0:54956"
-    default: echo
   tcp_server:
     listen:
       - "127.0.0.1:54500"
     default: tester
+  tcp_echo_server:
+    listen:
+      - "0.0.0.0:54956"
+    default: echo
   kcp_server:
+    protocol: kcp
+    listen:
+      - "127.0.0.1:54958"
+    default: tester
+  kcp_echo_server:
     protocol: kcp
     listen:
       - "127.0.0.1:54959"