From 10f8e30b92ec9ba993741dda5f673a13af868c48 Mon Sep 17 00:00:00 2001
From: zonyitoo <zonyitoo@gmail.com>
Date: Wed, 28 Apr 2021 15:59:11 +0800
Subject: [PATCH] tests based on HTTP should read status line only

---
 tests/http.rs   |  7 +++++--
 tests/socks4.rs | 10 +++++-----
 tests/socks5.rs | 18 +++++++++---------
 tests/tunnel.rs | 12 +++++-------
 4 files changed, 24 insertions(+), 23 deletions(-)

diff --git a/tests/http.rs b/tests/http.rs
index e0386b37..70ae7fdb 100644
--- a/tests/http.rs
+++ b/tests/http.rs
@@ -3,7 +3,7 @@
 use std::time::Duration;
 
 use tokio::{
-    io::{AsyncReadExt, AsyncWriteExt},
+    io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader},
     net::TcpStream,
     time,
 };
@@ -59,6 +59,7 @@ async fn http_proxy() {
             .unwrap();
         c.flush().await.unwrap();
 
+        // Proxy should close connection actively because HTTP/1.0 use short connection by default
         let mut buf = Vec::new();
         c.read_to_end(&mut buf).await.unwrap();
 
@@ -77,8 +78,10 @@ async fn http_proxy() {
             .unwrap();
         c.flush().await.unwrap();
 
+        let mut r = BufReader::new(c);
+
         let mut buf = Vec::new();
-        c.read_to_end(&mut buf).await.unwrap();
+        r.read_until(b'\n', &mut buf).await.unwrap();
 
         assert!(buf.starts_with(b"HTTP/1.0 200 OK\r\n"));
     }
diff --git a/tests/socks4.rs b/tests/socks4.rs
index 898369b5..83cd408f 100644
--- a/tests/socks4.rs
+++ b/tests/socks4.rs
@@ -6,7 +6,7 @@ use std::{
 };
 
 use tokio::{
-    io::{AsyncReadExt, AsyncWriteExt},
+    io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
     time::{self, Duration},
 };
 
@@ -89,11 +89,11 @@ async fn socks4_relay_connect() {
     c.write_all(HTTP_REQUEST).await.unwrap();
     c.flush().await.unwrap();
 
-    let mut buf = Vec::new();
-    c.read_to_end(&mut buf).await.unwrap();
+    let mut r = BufReader::new(c);
 
-    println!("Got reply from server: {}", str::from_utf8(&buf).unwrap());
+    let mut buf = Vec::new();
+    r.read_until(b'\n', &mut buf).await.unwrap();
 
     let http_status = b"HTTP/1.0 200 OK\r\n";
-    buf.starts_with(http_status);
+    assert!(buf.starts_with(http_status));
 }
diff --git a/tests/socks5.rs b/tests/socks5.rs
index 2b963827..81f34e96 100644
--- a/tests/socks5.rs
+++ b/tests/socks5.rs
@@ -6,7 +6,7 @@ use std::{
 };
 
 use tokio::{
-    io::{AsyncReadExt, AsyncWriteExt},
+    io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
     time::{self, Duration},
 };
 
@@ -95,13 +95,13 @@ async fn socks5_relay_stream() {
     c.write_all(req).await.unwrap();
     c.flush().await.unwrap();
 
-    let mut buf = Vec::new();
-    c.read_to_end(&mut buf).await.unwrap();
+    let mut r = BufReader::new(c);
 
-    println!("Got reply from server: {}", str::from_utf8(&buf).unwrap());
+    let mut buf = Vec::new();
+    r.read_until(b'\n', &mut buf).await.unwrap();
 
     let http_status = b"HTTP/1.0 200 OK\r\n";
-    buf.starts_with(http_status);
+    assert!(buf.starts_with(http_status));
 }
 
 #[tokio::test]
@@ -128,11 +128,11 @@ async fn socks5_relay_aead() {
     c.write_all(req).await.unwrap();
     c.flush().await.unwrap();
 
-    let mut buf = Vec::new();
-    c.read_to_end(&mut buf).await.unwrap();
+    let mut r = BufReader::new(c);
 
-    println!("Got reply from server: {}", str::from_utf8(&buf).unwrap());
+    let mut buf = Vec::new();
+    r.read_until(b'\n', &mut buf).await.unwrap();
 
     let http_status = b"HTTP/1.0 200 OK\r\n";
-    buf.starts_with(http_status);
+    assert!(buf.starts_with(http_status));
 }
diff --git a/tests/tunnel.rs b/tests/tunnel.rs
index 8a57510b..8f0fe157 100644
--- a/tests/tunnel.rs
+++ b/tests/tunnel.rs
@@ -1,12 +1,10 @@
 #![cfg(all(feature = "local-tunnel", feature = "server"))]
 
-use std::str;
-
 use byte_string::ByteStr;
 use log::debug;
 use tokio::{
     self,
-    io::{AsyncReadExt, AsyncWriteExt},
+    io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
     net::{TcpStream, UdpSocket},
     time::{self, Duration},
 };
@@ -64,13 +62,13 @@ async fn tcp_tunnel() {
     stream.write_all(req).await.unwrap();
     stream.flush().await.unwrap();
 
-    let mut buf = Vec::new();
-    stream.read_to_end(&mut buf).await.unwrap();
+    let mut r = BufReader::new(stream);
 
-    println!("Got reply from server: {}", str::from_utf8(&buf).unwrap());
+    let mut buf = Vec::new();
+    r.read_until(b'\n', &mut buf).await.unwrap();
 
     let http_status = b"HTTP/1.0 200 OK\r\n";
-    buf.starts_with(http_status);
+    assert!(buf.starts_with(http_status));
 }
 
 #[tokio::test]