Improve code safety and error handling + performance improvements

bench.mojo

@@ -1,22 +1,19 @@
+from memory import Span
 from benchmark import *
 from lightbug_http.io.bytes import bytes, Bytes
 from lightbug_http.header import Headers, Header
 from lightbug_http.utils import ByteReader, ByteWriter
 from lightbug_http.http import HTTPRequest, HTTPResponse, encode
 from lightbug_http.uri import URI
 
-alias headers = bytes(
-    """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"""
-)
-alias body = bytes(String("I am the body of an HTTP request") * 5)
-alias Request = bytes(
-    """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"""
-) + body
-alias Response = bytes(
-    "HTTP/1.1 200 OK\r\nserver: lightbug_http\r\ncontent-type:"
+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"
+
+alias body = "I am the body of an HTTP request" * 5
+alias body_bytes = bytes(body)
+alias Request = "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" + body
+alias Response = "HTTP/1.1 200 OK\r\nserver: lightbug_http\r\ncontent-type:"
     " application/octet-stream\r\nconnection: keep-alive\r\ncontent-length:"
-    " 13\r\ndate: 2024-06-02T13:41:50.766880+00:00\r\n\r\n"
-) + body
+    " 13\r\ndate: 2024-06-02T13:41:50.766880+00:00\r\n\r\n" + body
 
 
 fn main():
@@ -66,7 +63,7 @@ fn lightbug_benchmark_response_encode(mut b: Bencher):
     @always_inline
     @parameter
     fn response_encode():
-        var res = HTTPResponse(body, headers=headers_struct)
+        var res = HTTPResponse(body.as_bytes(), headers=headers_struct)
         _ = encode(res^)
 
     b.iter[response_encode]()
@@ -79,7 +76,7 @@ fn lightbug_benchmark_response_parse(mut b: Bencher):
     fn response_parse():
         var res = Response
         try:
-            _ = HTTPResponse.from_bytes(res^)
+            _ = HTTPResponse.from_bytes(res.as_bytes())
         except:
             pass
 
@@ -93,7 +90,7 @@ fn lightbug_benchmark_request_parse(mut b: Bencher):
     fn request_parse():
         var r = Request
         try:
-            _ = HTTPRequest.from_bytes("127.0.0.1/path", 4096, r^)
+            _ = HTTPRequest.from_bytes("127.0.0.1/path", 4096, r.as_bytes())
         except:
             pass
 
@@ -108,7 +105,7 @@ fn lightbug_benchmark_request_encode(mut b: Bencher):
         var req = HTTPRequest(
             URI.parse("http://127.0.0.1:8080/some-path")[URI],
             headers=headers_struct,
-            body=body,
+            body=body_bytes,
         )
         _ = encode(req^)
 
@@ -122,7 +119,7 @@ fn lightbug_benchmark_header_encode(mut b: Bencher):
     fn header_encode():
         var b = ByteWriter()
         var h = headers_struct
-        h.encode_to(b)
+        b.write(h)
 
     b.iter[header_encode]()
 
@@ -135,7 +132,7 @@ fn lightbug_benchmark_header_parse(mut b: Bencher):
         try:
             var b = headers
             var header = Headers()
-            var reader = ByteReader(b^)
+            var reader = ByteReader(b.as_bytes())
             _ = header.parse_raw(reader)
         except:
             print("failed")

integration_test_client.mojo

@@ -1,44 +1,83 @@
+from collections import Dict
 from lightbug_http import *
 from lightbug_http.client import Client
+from lightbug_http.utils import logger
 from testing import *
 
 fn u(s: String) raises -> URI:
     return URI.parse_raises("http://127.0.0.1:8080/" + s)
 
 struct IntegrationTest:
     var client: Client
+    var results: Dict[String, String]
 
     fn __init__(out self):
         self.client = Client()
+        self.results = Dict[String, String]()
+
+    fn mark_successful(mut self, name: String):
+        self.results[name] = "✅"
+
+    fn mark_failed(mut self, name: String):
+        self.results[name] = "❌"
 
-    fn test_redirect(mut self) raises:
-        print("Testing redirect...")
+    fn test_redirect(mut self):
+        alias name = "test_redirect"
+        logger.info("Testing redirect...")
         var h = Headers(Header(HeaderKey.CONNECTION, 'keep-alive'))
-        var res = self.client.do(HTTPRequest(u("redirect"), headers=h))
-        assert_equal(res.status_code, StatusCode.OK)
-        assert_equal(to_string(res.body_raw), "yay you made it")
-        assert_equal(res.headers[HeaderKey.CONNECTION], "keep-alive")
+        try:
+            var res = self.client.do(HTTPRequest(u("redirect"), headers=h))
+            assert_equal(res.status_code, StatusCode.OK)
+            assert_equal(to_string(res.body_raw), "yay you made it")
+            var conn = res.headers.get(HeaderKey.CONNECTION)
+            if conn:
+                assert_equal(conn.value(), "keep-alive")
+            self.mark_successful(name)
+        except e:
+            logger.error("IntegrationTest.test_redirect has run into an error.")
+            logger.error(e)
+            self.mark_failed(name)
+            return
 
-    fn test_close_connection(mut self) raises:
-        print("Testing close connection...")
+    fn test_close_connection(mut self):
+        alias name = "test_close_connection"
+        logger.info("Testing close connection...")
         var h = Headers(Header(HeaderKey.CONNECTION, 'close'))
-        var res = self.client.do(HTTPRequest(u("close-connection"), headers=h))
-        assert_equal(res.status_code, StatusCode.OK)
-        assert_equal(to_string(res.body_raw), "connection closed")
-        assert_equal(res.headers[HeaderKey.CONNECTION], "close")
+        try:
+            var res = self.client.do(HTTPRequest(u("close-connection"), headers=h))
+            assert_equal(res.status_code, StatusCode.OK)
+            assert_equal(to_string(res.body_raw), "connection closed")
+            assert_equal(res.headers[HeaderKey.CONNECTION], "close")
+            self.mark_successful(name)
+        except e:
+            logger.error("IntegrationTest.test_close_connection has run into an error.")
+            logger.error(e)
+            self.mark_failed(name)
+            return
 
-    fn test_server_error(mut self) raises:
-        print("Testing internal server error...")
-        var res = self.client.do(HTTPRequest(u("error")))
-        assert_equal(res.status_code, StatusCode.INTERNAL_ERROR)
-        assert_equal(res.status_text, "Internal Server Error")
+    fn test_server_error(mut self):
+        alias name = "test_server_error"
+        logger.info("Testing internal server error...")
+        try:
+            var res = self.client.do(HTTPRequest(u("error")))
+            assert_equal(res.status_code, StatusCode.INTERNAL_ERROR)
+            assert_equal(res.status_text, "Internal Server Error")
+            self.mark_successful(name)
+        except e:
+            logger.error("IntegrationTest.test_server_error has run into an error.")
+            logger.error(e)
+            self.mark_failed(name)
+            return
 
-
-    fn run_tests(mut self) raises:
+    fn run_tests(mut self):
+        logger.info("Running Client Integration Tests...")
         self.test_redirect()
         self.test_close_connection()
         self.test_server_error()
 
-fn main() raises:
+        for test in self.results.items():
+            print(test[].key + ":", test[].value)
+
+fn main():
     var test = IntegrationTest()
     test.run_tests()

integration_test_server.mojo

@@ -1,7 +1,8 @@
 from lightbug_http import *
 
+
 @value
-struct IntegerationTestService(HTTPService):
+struct IntegrationTestService(HTTPService):
     fn func(mut self, req: HTTPRequest) raises -> HTTPResponse:
         var p = req.uri.path
         if p == "/redirect":
@@ -21,8 +22,9 @@ struct IntegerationTestService(HTTPService):
 
         return NotFound("wrong")
 
+
 fn main() raises:
     var server = Server(tcp_keep_alive=True)
-    var service = IntegerationTestService()
+    var service = IntegrationTestService()
     server.listen_and_serve("127.0.0.1:8080", service)
-
+

lightbug_http/client.mojo

@@ -14,24 +14,18 @@ from lightbug_http.http import HTTPRequest, HTTPResponse, encode
 from lightbug_http.header import Headers, HeaderKey
 from lightbug_http.net import create_connection, SysConnection
 from lightbug_http.io.bytes import Bytes
-from lightbug_http.utils import ByteReader
+from lightbug_http.utils import ByteReader, logger
 from collections import Dict
 
 
 struct Client:
-    var host: StringLiteral
+    var host: String
     var port: Int
     var name: String
 
     var _connections: Dict[String, SysConnection]
 
-    fn __init__(out self):
-        self.host = "127.0.0.1"
-        self.port = 8888
-        self.name = "lightbug_http_client"
-        self._connections = Dict[String, SysConnection]()
-
-    fn __init__(out self, host: StringLiteral, port: Int):
+    fn __init__(out self, host: String = "127.0.0.1", port: Int = 8888):
         self.host = host
         self.port = port
         self.name = "lightbug_http_client"
@@ -46,8 +40,7 @@ struct Client:
                 pass
 
     fn do(mut self, owned req: HTTPRequest) raises -> HTTPResponse:
-        """
-        The `do` method is responsible for sending an HTTP request to a server and receiving the corresponding response.
+        """The `do` method is responsible for sending an HTTP request to a server and receiving the corresponding response.
 
         It performs the following steps:
         1. Creates a connection to the server specified in the request.
@@ -58,72 +51,78 @@ struct Client:
 
         Note: The code assumes that the `HTTPRequest` object passed as an argument has a valid URI with a host and port specified.
 
-        Parameters
-        ----------
-        req : HTTPRequest :
-            An `HTTPRequest` object representing the request to be sent.
+        Args:
+            req: An `HTTPRequest` object representing the request to be sent.
 
-        Returns
-        -------
-        HTTPResponse :
+        Returns:
             The received response.
 
-        Raises
-        ------
-        Error :
-            If there is a failure in sending or receiving the message.
+        Raises:
+            Error: If there is a failure in sending or receiving the message.
         """
-        var uri = req.uri
-        var host = uri.host
-
-        if host == "":
-            raise Error("URI is nil")
+        if req.uri.host == "":
+            raise Error("Client.do: Request failed because the host field is empty.")
         var is_tls = False
 
-        if uri.is_https():
+        if req.uri.is_https():
             is_tls = True
 
         var host_str: String
         var port: Int
-
-        if ":" in host:
-            var host_port = host.split(":")
+        if ":" in req.uri.host:
+            var host_port: List[String]
+            try:
+                host_port = req.uri.host.split(":")
+            except:
+                raise Error("Client.do: Failed to split host and port.")
             host_str = host_port[0]
             port = atol(host_port[1])
         else:
-            host_str = host
+            host_str = req.uri.host
             if is_tls:
                 port = 443
             else:
                 port = 80
 
         var conn: SysConnection
         var cached_connection = False
-        if host_str in self._connections:
+        try:
             conn = self._connections[host_str]
             cached_connection = True
-        else:
-            conn = create_connection(socket(AF_INET, SOCK_STREAM, 0), host_str, port)
-            self._connections[host_str] = conn
+        except:
+            # If connection is not cached, create a new one.
+            try:
+                conn = create_connection(socket(AF_INET, SOCK_STREAM, 0), host_str, port)
+                self._connections[host_str] = conn
+            except e:
+                logger.error(e)
+                raise Error("Client.do: Failed to create a connection to host.")
 
-        var bytes_sent = conn.write(encode(req))
-        if bytes_sent == -1:
+        var bytes_sent: Int
+        try:
+            bytes_sent = conn.write(encode(req))
+        except e:
             # Maybe peer reset ungracefully, so try a fresh connection
-            self._close_conn(host_str)
-            if cached_connection:
-                return self.do(req^)
-            raise Error("Failed to send message")
+            if str(e) == "SendError: Connection reset by peer.":
+                logger.debug("Client.do: Connection reset by peer. Trying a fresh connection.")
+                self._close_conn(host_str)
+                if cached_connection:
+                    return self.do(req^)
+            logger.error("Client.do: Failed to send message.")
+            raise e
 
+        # TODO: What if the response is too large for the buffer? We should read until the end of the response.
         var new_buf = Bytes(capacity=default_buffer_size)
         var bytes_recv = conn.read(new_buf)
 
         if bytes_recv == 0:
             self._close_conn(host_str)
             if cached_connection:
                 return self.do(req^)
-            raise Error("No response received")
+            raise Error("Client.do: No response received from the server.")
+
         try:
-            var res = HTTPResponse.from_bytes(new_buf^, conn)
+            var res = HTTPResponse.from_bytes(new_buf, conn)
             if res.is_redirect():
                 self._close_conn(host_str)
                 return self._handle_redirect(req^, res^)
@@ -140,9 +139,17 @@ struct Client:
         mut self, owned original_req: HTTPRequest, owned original_response: HTTPResponse
     ) raises -> HTTPResponse:
         var new_uri: URI
-        var new_location = original_response.headers[HeaderKey.LOCATION]
-        if new_location.startswith("http"):
-            new_uri = URI.parse_raises(new_location)
+        var new_location: String
+        try:
+            new_location = original_response.headers[HeaderKey.LOCATION]
+        except e:
+            raise Error("Client._handle_redirect: `Location` header was not received in the response.")
+
+        if new_location and new_location.startswith("http"):
+            try:
+                new_uri = URI.parse_raises(new_location)
+            except e:
+                raise Error("Client._handle_redirect: Failed to parse the new URI - " + str(e))
             original_req.headers[HeaderKey.HOST] = new_uri.host
         else:
             new_uri = original_req.uri

lightbug_http/cookie/request_cookie_jar.mojo

@@ -19,10 +19,11 @@ struct RequestCookieJar(Writable, Stringable):
             self._inner[cookie[].name] = cookie[].value
 
     fn parse_cookies(mut self, headers: Headers) raises:
-        var cookie_header = headers[HeaderKey.COOKIE]
+        var cookie_header = headers.get(HeaderKey.COOKIE)
         if not cookie_header:
             return None
-        var cookie_strings = cookie_header.split("; ")
+
+        var cookie_strings = cookie_header.value().split("; ")
 
         for chunk in cookie_strings:
             var key = String("")