From 275abf3b39c5836b82105675c20439b3fde373aa Mon Sep 17 00:00:00 2001
From: yggverse <yggverse@project>
Date: Sun, 23 Feb 2025 05:23:21 +0200
Subject: [PATCH] fix header bytes read

---
 src/main.rs | 62 ++++++++++++++++++++++++++++++++---------------------
 1 file changed, 37 insertions(+), 25 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 51396b8..fc7202c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -59,14 +59,37 @@ fn handle(
     //println!("[{}] [info] [{peer}] New connection", now());
     match connection {
         Ok(mut stream) => {
-            let mut input = vec![0; titanite::HEADER_MAX_LEN];
-            match stream.read(&mut input) {
-                Ok(0) => println!("[{}] [warning] [{peer}] Peer closed connection", now()),
-                Ok(s) => match Request::from_bytes(&input[..s]) {
-                    Ok(request) => match request {
-                        Request::Gemini(this) => gemini(this, argument, peer, &mut stream),
-                        Request::Titan(this) => titan(this, argument, peer, &mut stream),
-                    },
+            let mut header = Vec::with_capacity(HEADER_MAX_LEN);
+            loop {
+                let mut buffer = vec![0];
+                match stream.read(&mut buffer) {
+                    Ok(0) => println!("[{}] [warning] [{peer}] Peer closed connection", now()),
+                    Ok(_) => {
+                        header.push(buffer[0]);
+                        if header.len() > HEADER_MAX_LEN || buffer[0] == b'\n' {
+                            return match Request::from_bytes(&header) {
+                                Ok(request) => match request {
+                                    Request::Gemini(this) => {
+                                        gemini(this, &argument, &peer, &mut stream)
+                                    }
+                                    Request::Titan(this) => {
+                                        titan(this, &argument, &peer, &mut stream)
+                                    }
+                                },
+                                Err(e) => send(
+                                    &response::failure::temporary::General {
+                                        message: Some("Internal server error".to_string()),
+                                    }
+                                    .into_bytes(),
+                                    &mut stream,
+                                    |result| match result {
+                                        Ok(()) => println!("[{}] [warning] [{peer}] {e}", now()),
+                                        Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
+                                    },
+                                ),
+                            };
+                        }
+                    }
                     Err(e) => send(
                         &response::failure::temporary::General {
                             message: Some("Internal server error".to_string()),
@@ -78,18 +101,7 @@ fn handle(
                             Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                         },
                     ),
-                },
-                Err(e) => send(
-                    &response::failure::temporary::General {
-                        message: Some("Internal server error".to_string()),
-                    }
-                    .into_bytes(),
-                    &mut stream,
-                    |result| match result {
-                        Ok(()) => println!("[{}] [warning] [{peer}] {e}", now()),
-                        Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
-                    },
-                ),
+                }
             }
         }
         Err(e) => println!("[{}] [warning] [{peer}] Handshake issue: {e}", now()),
@@ -98,8 +110,8 @@ fn handle(
 
 fn gemini(
     gemini: titanite::request::Gemini,
-    argument: Arc<Argument>,
-    peer: SocketAddr,
+    argument: &Argument,
+    peer: &SocketAddr,
     stream: &mut TlsStream<TcpStream>,
 ) {
     use titanite::*;
@@ -119,7 +131,7 @@ fn gemini(
                     let l = item.file.read_at(&mut data, read as u64).unwrap();
                     stream.write_all(&data[..l]).unwrap();
                     read += l;
-                    println!("[{}] [info] [{peer}] Chunk sent: {l} ({read} total)", now());
+                    // println!("[{}] [info] [{peer}] Chunk sent: {l} ({read} total)", now());
                     // EOF
                     if l == 0 {
                         stream.flush().unwrap();
@@ -147,8 +159,8 @@ fn gemini(
 
 fn titan(
     titan: titanite::request::Titan,
-    argument: Arc<Argument>,
-    peer: SocketAddr,
+    argument: &Argument,
+    peer: &SocketAddr,
     stream: &mut TlsStream<TcpStream>,
 ) {
     use titanite::*;