Merge pull request #262 from nowex35/main

add uri length check

Author: saviorand

benchmark/bench.mojo

@@ -5,6 +5,7 @@ from lightbug_http.header import Headers, Header
 from lightbug_http.io.bytes import ByteReader, ByteWriter
 from lightbug_http.http import HTTPRequest, HTTPResponse, encode
 from lightbug_http.uri import URI
+from lightbug_http.server import default_max_request_body_size, default_max_request_uri_length
 
 alias headers = "GET /index.html HTTP/1.1\r\nHost: example.com\r\nUser-Agent: Mozilla/5.0\r\nContent-Type: text/html\r\nContent-Length: 1234\r\nConnection: close\r\nTrailer: end-of-message\r\n\r\n"
 
@@ -73,7 +74,7 @@ fn lightbug_benchmark_request_parse(mut b: Bencher):
     @parameter
     fn request_parse():
         try:
-            _ = HTTPRequest.from_bytes("127.0.0.1/path", 4096, Request.as_bytes())
+            _ = HTTPRequest.from_bytes("127.0.0.1/path", default_max_request_body_size, default_max_request_uri_length, Request.as_bytes())
         except:
             pass
 

lightbug_http/cookie/request_cookie_jar.mojo

@@ -79,3 +79,13 @@ struct RequestCookieJar(Writable, Stringable):
 
     fn __str__(self) -> String:
         return to_string(self)
+
+    fn __eq__(self, other: RequestCookieJar) -> Bool:
+        if len(self._inner) != len(other._inner):
+            return False
+
+        for value in self._inner.items():
+            for other_value in other._inner.items():
+                if value.key != other_value.key or value.value != other_value.value:
+                    return False
+        return True

lightbug_http/header.mojo

@@ -116,3 +116,13 @@ struct Headers(Writable, Stringable):
 
     fn __str__(self) -> String:
         return String.write(self)
+
+    fn __eq__(self, other: Headers) -> Bool:
+        if len(self._inner) != len(other._inner):
+            return False
+
+        for value in self._inner.items():
+            for other_value in other._inner.items():
+                if value.key != other_value.key or value.value != other_value.value:
+                    return False
+        return True

lightbug_http/http/common_response.mojo

@@ -59,6 +59,14 @@ fn NotFound(path: String) -> HTTPResponse:
         status_text="Not Found",
     )
 
+fn URITooLong() -> HTTPResponse:
+    return HTTPResponse(
+        bytes("URI Too Long"),
+        headers=Headers(Header(HeaderKey.CONTENT_TYPE, "text/plain")),
+        status_code=414,
+        status_text="URI Too Long"
+    )
+
 
 fn InternalError() -> HTTPResponse:
     return HTTPResponse(

lightbug_http/http/request.mojo

@@ -44,7 +44,7 @@ struct HTTPRequest(Writable, Stringable, Encodable):
     var timeout: Duration
 
     @staticmethod
-    fn from_bytes(addr: String, max_body_size: Int, b: Span[Byte]) raises -> HTTPRequest:
+    fn from_bytes(addr: String, max_body_size: Int, max_uri_length: Int, b: Span[Byte]) raises -> HTTPRequest:
         var reader = ByteReader(b)
         var headers = Headers()
         var method: String
@@ -56,6 +56,9 @@ struct HTTPRequest(Writable, Stringable, Encodable):
         except e:
             raise Error("HTTPRequest.from_bytes: Failed to parse request headers: " + String(e))
 
+        if len(uri.as_bytes()) > max_uri_length:
+            raise Error("HTTPRequest.from_bytes: Request URI too long")
+
         var cookies = RequestCookieJar()
         try:
             cookies.parse_cookies(headers)
@@ -189,3 +192,21 @@ struct HTTPRequest(Writable, Stringable, Encodable):
 
     fn __str__(self) -> String:
         return String.write(self)
+    
+    fn __eq__(self, other: HTTPRequest) -> Bool:
+        return (
+            self.method == other.method
+            and self.protocol == other.protocol
+            and self.uri == other.uri
+            and self.headers == other.headers
+            and self.cookies == other.cookies
+            and self.body_raw.__str__() == other.body_raw.__str__()
+        )
+    
+    fn __isnot__(self, other: HTTPRequest) -> Bool:
+        return not self.__eq__(other)
+
+    fn __isnot__(self, other: None) -> Bool:
+        if self.get_body() != "" or self.uri.request_uri != "":
+            return True
+        return False

lightbug_http/server.mojo

@@ -5,7 +5,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, BadRequest
+from lightbug_http.http.common_response import InternalError, BadRequest, URITooLong
 from lightbug_http.uri import URI
 from lightbug_http.header import Headers
 from lightbug_http.service import HTTPService
@@ -14,6 +14,7 @@ from lightbug_http.error import ErrorHandler
 
 alias DefaultConcurrency: Int = 256 * 1024
 alias default_max_request_body_size = 4 * 1024 * 1024  # 4MB
+alias default_max_request_uri_length = 8192
 
 
 struct Server(Movable):
@@ -27,6 +28,7 @@ struct Server(Movable):
     var max_requests_per_connection: UInt
 
     var _max_request_body_size: UInt
+    var _max_request_uri_length: UInt
     var tcp_keep_alive: Bool
 
     fn __init__(
@@ -37,13 +39,15 @@ struct Server(Movable):
         max_concurrent_connections: UInt = 1000,
         max_requests_per_connection: UInt = 0,
         max_request_body_size: UInt = default_max_request_body_size,
+        max_request_uri_length: UInt = default_max_request_uri_length,
         tcp_keep_alive: Bool = False,
     ) raises:
         self.error_handler = error_handler
         self.name = name
         self._address = address
         self.max_requests_per_connection = max_requests_per_connection
-        self._max_request_body_size = default_max_request_body_size
+        self._max_request_body_size = max_request_body_size
+        self._max_request_uri_length = max_request_uri_length
         self.tcp_keep_alive = tcp_keep_alive
         if max_concurrent_connections == 0:
             self.max_concurrent_connections = DefaultConcurrency
@@ -57,6 +61,7 @@ struct Server(Movable):
         self.max_concurrent_connections = other.max_concurrent_connections
         self.max_requests_per_connection = other.max_requests_per_connection
         self._max_request_body_size = other._max_request_body_size
+        self._max_request_uri_length = other._max_request_uri_length
         self.tcp_keep_alive = other.tcp_keep_alive
 
     fn address(self) -> ref [self._address] String:
@@ -71,6 +76,12 @@ struct Server(Movable):
     fn set_max_request_body_size(mut self, size: UInt) -> None:
         self._max_request_body_size = size
 
+    fn max_request_uri_length(self) -> UInt:
+        return self._max_request_uri_length
+
+    fn set_max_request_uri_length(mut self, length: UInt) -> None:
+        self._max_request_uri_length = length
+
     fn get_concurrency(self) -> UInt:
         """Retrieve the concurrency level which is either
         the configured `max_concurrent_connections` or the `DefaultConcurrency`.
@@ -128,6 +139,10 @@ struct Server(Movable):
         var max_request_body_size = self.max_request_body_size()
         if max_request_body_size <= 0:
             max_request_body_size = default_max_request_body_size
+        
+        var max_request_uri_length = self.max_request_uri_length()
+        if max_request_uri_length <= 0:
+            max_request_uri_length = default_max_request_uri_length
 
         var req_number = 0
         while True:
@@ -163,7 +178,7 @@ struct Server(Movable):
 
             var request: HTTPRequest
             try:
-                request = HTTPRequest.from_bytes(self.address(), max_request_body_size, request_buffer)
+                request = HTTPRequest.from_bytes(self.address(), max_request_body_size, max_request_uri_length, request_buffer)
                 var response: HTTPResponse
                 var close_connection = (not self.tcp_keep_alive) or request.connection_close()
                 try:
@@ -200,7 +215,10 @@ struct Server(Movable):
             except e:
                 logger.error("Failed to parse HTTPRequest:", String(e))
                 try:
-                    _ = conn.write(encode(BadRequest()))
+                    if String(e) == "HTTPRequest.from_bytes: Request URI too long":
+                        _ = conn.write(encode(URITooLong()))
+                    else:
+                        _ = conn.write(encode(BadRequest()))
                 except e:
                     logger.error("Failed to write BadRequest response to the connection:", String(e))
                     conn.teardown()

lightbug_http/uri.mojo

@@ -240,6 +240,17 @@ struct URI(Writable, Stringable, Representable):
     fn __repr__(self) -> String:
         return String.write(self)
 
+    fn __eq__(self, other: URI) -> Bool:
+        return (
+            self.scheme == other.scheme
+            and self.host == other.host
+            and self.path == other.path
+            and self.query_string == other.query_string
+            and self._original_path == other._original_path
+            and self.full_uri == other.full_uri
+            and self.request_uri == other.request_uri
+        )
+
     fn write_to[T: Writer](self, mut writer: T):
         writer.write(
             "URI(",

tests/lightbug_http/http/test_request.mojo

@@ -3,32 +3,42 @@ from memory import Span
 from collections.string import StringSlice
 from lightbug_http.http import HTTPRequest, StatusCode
 from lightbug_http.strings import to_string
-
+from lightbug_http.server import default_max_request_body_size, default_max_request_uri_length
 
 def test_request_from_bytes():
     alias data = "GET /redirect HTTP/1.1\r\nHost: 127.0.0.1:8080\r\nUser-Agent: python-requests/2.32.3\r\nAccept-Encoding: gzip, deflate, br, zstd\r\nAccept: */*\r\nconnection: keep-alive\r\n\r\n"
-    var request = HTTPRequest.from_bytes("127.0.0.1", 4096, data.as_bytes())
-    testing.assert_equal(request.protocol, "HTTP/1.1")
-    testing.assert_equal(request.method, "GET")
-    testing.assert_equal(request.uri.request_uri, "/redirect")
-    testing.assert_equal(request.headers["Host"], "127.0.0.1:8080")
-    testing.assert_equal(request.headers["User-Agent"], "python-requests/2.32.3")
-
-    testing.assert_false(request.connection_close())
-    request.set_connection_close()
-    testing.assert_true(request.connection_close())
+    var request: HTTPRequest
+    try:
+        request = HTTPRequest.from_bytes("127.0.0.1", default_max_request_body_size, default_max_request_uri_length, data.as_bytes())
+        if request is not None:
+            testing.assert_equal(request.protocol, "HTTP/1.1")
+            testing.assert_equal(request.method, "GET")
+            testing.assert_equal(request.uri.request_uri, "/redirect")
+            testing.assert_equal(request.headers["Host"], "127.0.0.1:8080")
+            testing.assert_equal(request.headers["User-Agent"], "python-requests/2.32.3")
+
+            testing.assert_false(request.connection_close())
+            request.set_connection_close()
+            testing.assert_true(request.connection_close())
+    except e:
+        testing.assert_true(False, "Failed to parse HTTP request: " + String(e))
+    
 
 
 def test_read_body():
     alias data = "GET /redirect HTTP/1.1\r\nHost: 127.0.0.1:8080\r\nUser-Agent: python-requests/2.32.3\r\nAccept-Encoding: gzip, deflate, br, zstd\r\nAccept: */\r\nContent-Length: 17\r\nconnection: keep-alive\r\n\r\nThis is the body!"
-    var request = HTTPRequest.from_bytes("127.0.0.1", 4096, data.as_bytes())
-    testing.assert_equal(request.protocol, "HTTP/1.1")
-    testing.assert_equal(request.method, "GET")
-    testing.assert_equal(request.uri.request_uri, "/redirect")
-    testing.assert_equal(request.headers["Host"], "127.0.0.1:8080")
-    testing.assert_equal(request.headers["User-Agent"], "python-requests/2.32.3")
-
-    testing.assert_equal(String(request.get_body()), String("This is the body!"))
+    var request: HTTPRequest
+    try:
+        request = HTTPRequest.from_bytes("127.0.0.1", default_max_request_body_size, default_max_request_uri_length, data.as_bytes())
+        if request is not None:
+            testing.assert_equal(request.protocol, "HTTP/1.1")
+            testing.assert_equal(request.method, "GET")
+            testing.assert_equal(request.uri.request_uri, "/redirect")
+            testing.assert_equal(request.headers["Host"], "127.0.0.1:8080")
+            testing.assert_equal(request.headers["User-Agent"], "python-requests/2.32.3")
+            testing.assert_equal(String(request.get_body()), String("This is the body!"))
+    except e:
+        testing.assert_true(False, "Failed to parse HTTP request: " + String(e))
 
 
 def test_encode():