From 31e33bca287b10c0a44494d91e5571b829e0f51e Mon Sep 17 00:00:00 2001
From: yggverse <yggverse@project>
Date: Sun, 23 Feb 2025 23:27:21 +0200
Subject: [PATCH] update main handler, add comments

---
 src/main.rs | 94 +++++++++++++++++++++++++++++++++--------------------
 1 file changed, 59 insertions(+), 35 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 30d8ca3..9999ccd 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -59,45 +59,69 @@ fn handle(
     match connection {
         Ok(mut stream) => {
             // server should work with large files without memory overload,
-            // because of that incoming data read partially, using chunks;
-            // collect header bytes first to route the request
-            let mut header = Vec::with_capacity(HEADER_MAX_LEN);
+            // because of that, incoming data read partially, using chunks;
+            // collect header bytes first to route the request by its type.
+            let mut header_buffer = Vec::with_capacity(HEADER_MAX_LEN);
             loop {
-                let mut buffer = vec![0];
-                match stream.read(&mut buffer) {
+                let mut header_chunk = vec![0];
+                match stream.read(&mut header_chunk) {
                     Ok(0) => println!("[{}] [warning] [{peer}] Peer closed connection.", now()),
-                    Ok(_) => {
-                        header.push(buffer[0]);
-                        // header bytes collected
-                        if header.len() > HEADER_MAX_LEN || buffer[0] == b'\n' {
-                            // detect controller for the request by parse its header bytes
-                            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)
+                    Ok(l) => {
+                        // validate header buffer, break on its length reached protocol limits
+                        if header_buffer.len() + l > HEADER_MAX_LEN {
+                            return send(
+                                &response::failure::permanent::BadRequest {
+                                    message: Some("Bad request".to_string()),
+                                }
+                                .into_bytes(),
+                                &mut stream,
+                                |result| match result {
+                                    Ok(()) => {
+                                        println!("[{}] [warning] [{peer}] Bad request", now())
                                     }
+                                    Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
                                 },
-                                Err(e) => send(
-                                    &response::failure::permanent::BadRequest {
-                                        message: Some("Invalid request".to_string()),
+                            );
+                        }
+
+                        // take chunk bytes at this point
+                        header_buffer.extend(header_chunk);
+
+                        // ending header byte received
+                        if header_buffer.last().is_some_and(|&b| b == b'\n') {
+                            // header bytes contain valid Gemini **request**
+                            if let Ok(request) = request::Gemini::from_bytes(&header_buffer) {
+                                return gemini(request, &argument, &peer, &mut stream);
+                            }
+
+                            // header bytes contain valid Titan **meta**
+                            // * yet, no data has been received to parse the entire Titan request,
+                            //   the data will be handled separately upon success, in chunks.
+                            if let Ok(meta) = request::titan::Meta::from_bytes(&header_buffer) {
+                                return titan(meta, &argument, &peer, &mut stream);
+                            }
+
+                            // header bytes received but yet could not be parsed,
+                            // complete with request failure
+                            return send(
+                                &response::failure::permanent::BadRequest {
+                                    message: Some("Bad request".to_string()),
+                                }
+                                .into_bytes(),
+                                &mut stream,
+                                |result| match result {
+                                    Ok(()) => {
+                                        println!("[{}] [warning] [{peer}] Bad request", now())
                                     }
-                                    .into_bytes(),
-                                    &mut stream,
-                                    |result| match result {
-                                        Ok(()) => println!("[{}] [warning] [{peer}] {e}", now()),
-                                        Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
-                                    },
-                                ),
-                            };
+                                    Err(e) => println!("[{}] [error] [{peer}] {e}", now()),
+                                },
+                            );
                         }
                     }
                     Err(e) => {
                         return send(
                             &response::failure::permanent::BadRequest {
-                                message: Some("Invalid request".to_string()),
+                                message: Some("Bad request".to_string()),
                             }
                             .into_bytes(),
                             &mut stream,
@@ -174,16 +198,16 @@ fn gemini(
 }
 
 fn titan(
-    titan: titanite::request::Titan,
+    meta: titanite::request::titan::Meta,
     argument: &Argument,
     peer: &SocketAddr,
     stream: &mut TlsStream<TcpStream>,
 ) {
     use titanite::*;
-    println!("[{}] [info] [{peer}] Request: {}", now(), titan.meta.url);
+    println!("[{}] [info] [{peer}] Request: {}", now(), meta.url);
     // require content type for application,
     // even MIME value is optional by Titan specification
-    let mime = match titan.meta.mime {
+    let mime = match meta.mime {
         Some(mime) => mime,
         None => {
             const MESSAGE: &str = "Content type is required";
@@ -236,7 +260,7 @@ fn titan(
                     }
 
                     // validate client-side limits (from header)
-                    if total > titan.meta.size {
+                    if total > meta.size {
                         if let Err(e) = tmp.delete() {
                             println!("[{}] [error] [{peer}] {e}", now());
                         }
@@ -278,12 +302,12 @@ fn titan(
                             }
 
                             // just to make sure
-                            if total > titan.meta.size {
+                            if total > meta.size {
                                 panic!()
                             }
 
                             // all data received
-                            if titan.meta.size == total {
+                            if meta.size == total {
                                 return match tmp.commit() {
                                     Ok(pmt) => send(
                                         &response::redirect::Permanent {