From 02d166907c0373a07287bec94fd5e381c61c0379 Mon Sep 17 00:00:00 2001
From: NotAShelf <raf@notashelf.dev>
Date: Fri, 2 May 2025 14:01:50 +0300
Subject: [PATCH] network: better header handling

Please work...
---
 src/network.rs | 62 ++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 58 insertions(+), 4 deletions(-)

diff --git a/src/network.rs b/src/network.rs
index 226ed00..efdb8c8 100644
--- a/src/network.rs
+++ b/src/network.rs
@@ -66,6 +66,8 @@ pub async fn handle_connection(
     let mut buffer = vec![0; 8192];
     let mut request_data = Vec::with_capacity(8192);
     let mut header_end_pos = 0;
+    let mut content_length: Option<usize> = None;
+    let mut chunked_encoding = false;
 
     // Read with timeout to prevent hanging resource load ops.
     let read_fut = async {
@@ -76,13 +78,65 @@ pub async fn handle_connection(
                     let new_data = &buffer[..n];
                     request_data.extend_from_slice(new_data);
 
-                    // Look for end of headers
+                    // Look for end of headers if we haven't found it yet
                     if header_end_pos == 0 {
                         if let Some(pos) = find_header_end(&request_data) {
                             header_end_pos = pos;
-                            // XXX: Breaking here appears to be malforming the request
-                            // and causing 404 errors.
-                            // So, continue reading the body if present but do not break.
+
+                            // Parse headers to determine body size
+                            let headers_str =
+                                String::from_utf8_lossy(&request_data[..header_end_pos]);
+
+                            // Check for Content-Length
+                            if let Some(cl_line) = headers_str
+                                .lines()
+                                .find(|l| l.to_lowercase().starts_with("content-length:"))
+                            {
+                                if let Some(len_str) = cl_line.split(':').nth(1) {
+                                    if let Ok(len) = len_str.trim().parse::<usize>() {
+                                        content_length = Some(len);
+                                    }
+                                }
+                            }
+
+                            // Check for chunked encoding
+                            if headers_str.lines().any(|l| {
+                                l.to_lowercase().starts_with("transfer-encoding:")
+                                    && l.to_lowercase().contains("chunked")
+                            }) {
+                                chunked_encoding = true;
+                            }
+
+                            // For requests with no body (GET, HEAD, etc.), we're done
+                            if (content_length == Some(0) || content_length.is_none())
+                                && !chunked_encoding
+                            {
+                                let method = headers_str
+                                    .lines()
+                                    .next()
+                                    .and_then(|l| l.split_whitespace().next())
+                                    .unwrap_or("")
+                                    .to_uppercase();
+
+                                if method == "GET"
+                                    || method == "HEAD"
+                                    || method == "OPTIONS"
+                                    || method == "TRACE"
+                                {
+                                    break;
+                                }
+                            }
+                        }
+                    } else if !chunked_encoding && content_length.is_some() {
+                        // We have Content-Length, check if we've read enough
+                        let body_len = request_data.len() - (header_end_pos + 4); // +4 for CRLF CRLF
+                        if body_len >= content_length.unwrap() {
+                            break;
+                        }
+                    } else if chunked_encoding {
+                        // For chunked encoding, look for the terminating chunk
+                        if request_data.windows(5).any(|w| w == b"0\r\n\r\n") {
+                            break;
                         }
                     }