bufreader

src/flowgate/server.rs

@@ -1,13 +1,10 @@
 use std::{
-    error::Error, 
-    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6}, 
-    str::FromStr, 
-    sync::Arc
+    error::Error, net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6}, str::FromStr, sync::Arc
 };
 
 use ignore_result::Ignore;
 use tokio::{
-    io::{AsyncReadExt, AsyncWriteExt}, 
+    io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}, 
     net::{TcpListener, TcpStream}
 };
 
@@ -146,20 +143,11 @@ impl FlowgateServer {
     ) -> Option<Connection> {
         let mut addr = addr;
 
+        let mut stream = BufReader::new(stream);
+
         match &self.config.incoming_ip_forwarding {
             IpForwarding::Simple => {
-                let mut header = Vec::new();
-
-                {
-                    let mut buf = [0; 1];
-
-                    while let Ok(1) = stream.read(&mut buf).await {
-                        let byte = buf[0];
-                        if byte == b'\n' { break }
-                        header.push(byte);
-                    }
-                }
-
+                let header = read_until(&mut stream, b"\n").await?;
                 addr = SocketAddr::from_str(&String::from_utf8(header).ok()?).ok()?;
             },
             IpForwarding::Modern => {
@@ -183,30 +171,10 @@ impl FlowgateServer {
                     }, _ => { return None },
                 };
             },
-            _ => { }
+            _ => {}
         }
 
-        let mut head = Vec::new();
-
-        {
-            let mut buf = [0; 1];
-            let mut counter = 0;
-
-            while let Ok(1) = stream.read(&mut buf).await {
-                let byte = buf[0];
-                head.push(byte);
-
-                counter = match (counter, byte) {
-                    (0, b'\r') => 1,
-                    (1, b'\n') => 2,
-                    (2, b'\r') => 3,
-                    (3, b'\n') => break,
-                    _ => 0,
-                };
-            }
-
-            head.truncate(head.len() - 4);
-        }
+        let mut head = read_until(&mut stream, b"\r\n\r\n").await?;
 
         if head.is_empty() { return None; }
         
@@ -353,6 +321,7 @@ impl FlowgateServer {
         } else if is_chunked {
             loop {
                 let mut length = Vec::new();
+
                 {
                     let mut buf = [0; 1];
                     let mut counter = 0;
@@ -371,6 +340,7 @@ impl FlowgateServer {
 
                     length.truncate(length.len() - 2);
                 }
+
                 let length = String::from_utf8(length).ok()?;
                 let length = usize::from_str_radix(length.as_str(), 16).ok()?;
                 let mut data = vec![0u8; length+2];
@@ -490,3 +460,28 @@ impl FlowgateServer {
         Some(conn)
     }
 }
+
+async fn read_until(stream: &mut (impl AsyncBufReadExt + Unpin), delimiter: &[u8]) -> Option<Vec<u8>> {
+    let mut data = Vec::new();
+    let mut counter = 0usize;
+
+    let mut buf = vec![0];
+
+    while let Ok(1) = stream.read(&mut buf).await {
+        let char = buf[0];
+        data.push(char);
+
+        if char == delimiter[counter] {
+            counter += 1;
+            if counter == delimiter.len() {
+                break
+            }
+        } else {
+            counter = 0;
+        }
+    }
+
+    data.truncate(data.len()-delimiter.len());
+
+    Some(data)
+}