Merge pull request #247 from Lightbug-HQ/fix/incomplete-requests-handling

Fix/incomplete requests handling

Author: saviorand

lightbug_http/http/common_response.mojo

@@ -41,3 +41,11 @@ fn InternalError() -> HTTPResponse:
         headers=Headers(Header(HeaderKey.CONTENT_TYPE, "text/plain")),
         status_text="Internal Server Error",
     )
+
+fn BadRequest() -> HTTPResponse:
+    return HTTPResponse(
+        bytes("Bad Request"),
+        status_code=400,
+        headers=Headers(Header(HeaderKey.CONTENT_TYPE, "text/plain")),
+        status_text="Bad Request",
+    )

lightbug_http/http/request.mojo

@@ -128,8 +128,20 @@ struct HTTPRequest(Writable, Stringable):
         if content_length > max_body_size:
             raise Error("Request body too large")
 
-        self.body_raw = r.read_bytes(content_length).to_bytes()
-        self.set_content_length(content_length)
+        try:
+            self.body_raw = r.read_bytes(content_length).to_bytes()
+            self.set_content_length(len(self.body_raw))
+        except OutOfBoundsError:
+            logger.debug("Failed to read full request body as per content-length header. Proceeding with the available bytes.")
+            var available_bytes = len(r._inner) - r.read_pos
+            if available_bytes > 0:
+                self.body_raw = r.read_bytes(available_bytes).to_bytes()
+                self.set_content_length(len(self.body_raw))
+            else:
+                logger.debug("No body bytes available. Setting content-length to 0.")
+                self.body_raw = Bytes()
+                self.set_content_length(0)
+
 
     fn write_to[T: Writer, //](self, mut writer: T):
         path = self.uri.path if len(self.uri.path) > 1 else strSlash

lightbug_http/server.mojo

@@ -6,7 +6,7 @@ from lightbug_http._logger import logger
 from lightbug_http.connection import NoTLSListener, default_buffer_size, TCPConnection, ListenConfig
 from lightbug_http.socket import Socket
 from lightbug_http.http import HTTPRequest, encode
-from lightbug_http.http.common_response import InternalError
+from lightbug_http.http.common_response import InternalError, BadRequest
 from lightbug_http.uri import URI
 from lightbug_http.header import Headers
 from lightbug_http.service import HTTPService
@@ -159,47 +159,50 @@ struct Server(Movable):
                     if String(e) == "EOF":
                         return
                     else:
-                        logger.error(e)
-                        raise Error("Server.serve_connection: Failed to read request")
+                        logger.error("Server.serve_connection: Failed to read request. Expected EOF, got:", String(e))
+                        return
 
             var request: HTTPRequest
             try:
                 request = HTTPRequest.from_bytes(self.address(), max_request_body_size, request_buffer)
+                var response: HTTPResponse
+                var close_connection = (not self.tcp_keep_alive) or request.connection_close()
+                try:
+                    response = handler.func(request)
+                    if close_connection:
+                        response.set_connection_close()
+                    logger.debug(
+                        conn.socket._remote_address.ip,
+                        String(conn.socket._remote_address.port),
+                        request.method,
+                        request.uri.path,
+                        response.status_code,
+                    )
+                    try:
+                        _ = conn.write(encode(response^))
+                    except e:
+                        logger.error("Failed to write encoded response to the connection:", String(e))
+                        conn.teardown()
+                        break
+
+                    if close_connection:
+                        conn.teardown()
+                        break
+                except e:
+                    logger.error("Handler error:", String(e))
+                    if not conn.is_closed():
+                        try:
+                            _ = conn.write(encode(InternalError()))
+                        except e:
+                            raise Error("Failed to send InternalError response")
+                        finally:
+                            conn.teardown()
+                        return
             except e:
                 logger.error("Failed to parse HTTPRequest:", String(e))
-                raise Error("Server.serve_connection: Failed to parse request")
-
-            var response: HTTPResponse
-            var close_connection = (not self.tcp_keep_alive) or request.connection_close()
-            try:
-                response = handler.func(request)
-                if close_connection:
-                    response.set_connection_close()
-                logger.debug(
-                    conn.socket._remote_address.ip,
-                    String(conn.socket._remote_address.port),
-                    request.method,
-                    request.uri.path,
-                    response.status_code,
-                )
                 try:
-                    _ = conn.write(encode(response^))
+                    _ = conn.write(encode(BadRequest()))
                 except e:
-                    logger.error("Failed to write encoded response to the connection:", String(e))
-                    conn.teardown()
-                    break
-
-                if close_connection:
+                    logger.error("Failed to write BadRequest response to the connection:", String(e))
                     conn.teardown()
                     break
-            except e:
-                logger.error("Handler error:", String(e))
-
-                if not conn.is_closed():
-                    try:
-                        _ = conn.write(encode(InternalError()))
-                    except e:
-                        raise Error("Failed to send InternalError response")
-                    finally:
-                        conn.teardown()
-                    return

tests/integration/integration_client.py

@@ -1,4 +1,6 @@
 import requests
+import socket
+import time
 
 session = requests.Session()
 
@@ -23,3 +25,16 @@
 }
 response = session.get('http://127.0.0.1:8080/large-headers', headers=large_headers)
 assert response.status_code == 200
+
+print("\n~~~ Testing content-length mismatch (smaller) ~~~")
+def test_content_length_smaller():
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    s.connect(('127.0.0.1', 8080))
+    s.sendall(b'POST / HTTP/1.1\r\nHost: localhost\r\nContent-Length: 100\r\n\r\nOnly sending 20 bytes')
+    time.sleep(1)
+    s.close()
+
+test_content_length_smaller()
+time.sleep(1)
+
+print("\n~~~ All tests completed ~~~")

tests/integration/integration_test_server.mojo

@@ -5,7 +5,9 @@ from lightbug_http import *
 struct IntegrationTestService(HTTPService):
     fn func(mut self, req: HTTPRequest) raises -> HTTPResponse:
         var p = req.uri.path
-        if p == "/redirect":
+        if p == "/":
+            return OK("hello")
+        elif p == "/redirect":
             return HTTPResponse(
                 "get off my lawn".as_bytes(),
                 headers=Headers(Header(HeaderKey.LOCATION, "/rd-destination")),