Fix socket

benchmark/bench.mojo

@@ -2,7 +2,7 @@ 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.io.bytes import ByteReader, ByteWriter
 from lightbug_http.http import HTTPRequest, HTTPResponse, encode
 from lightbug_http.uri import URI
 
@@ -11,9 +11,7 @@ alias headers = "GET /index.html HTTP/1.1\r\nHost: example.com\r\nUser-Agent: Mo
 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
+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
 
 
 fn main():
@@ -26,24 +24,12 @@ fn run_benchmark():
         config.verbose_timing = True
         config.tabular_view = True
         var m = Bench(config)
-        m.bench_function[lightbug_benchmark_header_encode](
-            BenchId("HeaderEncode")
-        )
-        m.bench_function[lightbug_benchmark_header_parse](
-            BenchId("HeaderParse")
-        )
-        m.bench_function[lightbug_benchmark_request_encode](
-            BenchId("RequestEncode")
-        )
-        m.bench_function[lightbug_benchmark_request_parse](
-            BenchId("RequestParse")
-        )
-        m.bench_function[lightbug_benchmark_response_encode](
-            BenchId("ResponseEncode")
-        )
-        m.bench_function[lightbug_benchmark_response_parse](
-            BenchId("ResponseParse")
-        )
+        m.bench_function[lightbug_benchmark_header_encode](BenchId("HeaderEncode"))
+        m.bench_function[lightbug_benchmark_header_parse](BenchId("HeaderParse"))
+        m.bench_function[lightbug_benchmark_request_encode](BenchId("RequestEncode"))
+        m.bench_function[lightbug_benchmark_request_parse](BenchId("RequestParse"))
+        m.bench_function[lightbug_benchmark_response_encode](BenchId("ResponseEncode"))
+        m.bench_function[lightbug_benchmark_response_parse](BenchId("ResponseParse"))
         m.dump_report()
     except:
         print("failed to start benchmark")
@@ -102,12 +88,12 @@ fn lightbug_benchmark_request_encode(mut b: Bencher):
     fn request_encode() raises:
         var uri = URI.parse("http://127.0.0.1:8080/some-path")
         var req = HTTPRequest(
-        uri=uri,
-        headers=headers_struct,
-        body=body_bytes,
+            uri=uri,
+            headers=headers_struct,
+            body=body_bytes,
         )
         _ = encode(req^)
-        
+
     try:
         b.iter[request_encode]()
     except e:
@@ -134,8 +120,7 @@ fn lightbug_benchmark_header_parse(mut b: Bencher):
             var header = Headers()
             var reader = ByteReader(headers.as_bytes())
             _ = header.parse_raw(reader)
-        except:
-            print("failed")
+        except e:
+            print("failed", e)
 
     b.iter[header_parse]()
-

lightbug_http/_logger.mojo

@@ -0,0 +1,113 @@
+from sys.param_env import env_get_string
+
+
+struct LogLevel:
+    alias FATAL = 0
+    alias ERROR = 1
+    alias WARN = 2
+    alias INFO = 3
+    alias DEBUG = 4
+
+
+fn get_log_level() -> Int:
+    """Returns the log level based on the parameter environment variable `LOG_LEVEL`.
+
+    Returns:
+        The log level.
+    """
+    alias level = env_get_string["LB_LOG_LEVEL", "INFO"]()
+    if level == "INFO":
+        return LogLevel.INFO
+    elif level == "WARN":
+        return LogLevel.WARN
+    elif level == "ERROR":
+        return LogLevel.ERROR
+    elif level == "DEBUG":
+        return LogLevel.DEBUG
+    elif level == "FATAL":
+        return LogLevel.FATAL
+    else:
+        return LogLevel.INFO
+
+
+alias LOG_LEVEL = get_log_level()
+"""Logger level determined by the `LB_LOG_LEVEL` param environment variable.
+
+When building or running the application, you can set `LB_LOG_LEVEL` by providing the the following option:
+
+```bash
+mojo build ... -D LB_LOG_LEVEL=DEBUG
+# or
+mojo ... -D LB_LOG_LEVEL=DEBUG
+```
+"""
+
+
+@value
+struct Logger[level: Int]:
+    alias STDOUT = 1
+    alias STDERR = 2
+
+    fn _log_message[event_level: Int](self, message: String):
+        @parameter
+        if level >= event_level:
+
+            @parameter
+            if event_level < LogLevel.WARN:
+                # Write to stderr if FATAL or ERROR
+                print(message, file=Self.STDERR)
+            else:
+                print(message)
+
+    fn info[*Ts: Writable](self, *messages: *Ts):
+        var msg = String.write("\033[36mINFO\033[0m  - ")
+
+        @parameter
+        fn write_message[T: Writable](message: T):
+            msg.write(message, " ")
+
+        messages.each[write_message]()
+        self._log_message[LogLevel.INFO](msg)
+
+    fn warn[*Ts: Writable](self, *messages: *Ts):
+        var msg = String.write("\033[33mWARN\033[0m  - ")
+
+        @parameter
+        fn write_message[T: Writable](message: T):
+            msg.write(message, " ")
+
+        messages.each[write_message]()
+        self._log_message[LogLevel.WARN](msg)
+
+    fn error[*Ts: Writable](self, *messages: *Ts):
+        var msg = String.write("\033[31mERROR\033[0m - ")
+
+        @parameter
+        fn write_message[T: Writable](message: T):
+            msg.write(message, " ")
+
+        messages.each[write_message]()
+        self._log_message[LogLevel.ERROR](msg)
+
+    fn debug[*Ts: Writable](self, *messages: *Ts):
+        var msg = String.write("\033[34mDEBUG\033[0m - ")
+
+        @parameter
+        fn write_message[T: Writable](message: T):
+            msg.write(message, " ")
+
+        messages.each[write_message]()
+        self._log_message[LogLevel.DEBUG](msg)
+
+    fn fatal[*Ts: Writable](self, *messages: *Ts):
+        var msg = String.write("\033[35mFATAL\033[0m - ")
+
+        @parameter
+        fn write_message[T: Writable](message: T):
+            msg.write(message, " ")
+
+        messages.each[write_message]()
+        self._log_message[LogLevel.FATAL](msg)
+
+
+alias logger = Logger[LOG_LEVEL]()

lightbug_http/client.mojo

@@ -1,23 +1,14 @@
 from collections import Dict
+from utils import StringSlice
 from memory import UnsafePointer
-from lightbug_http.libc import (
-    c_int,
-    AF_INET,
-    SOCK_STREAM,
-    socket,
-    connect,
-    send,
-    recv,
-    close,
-)
-from lightbug_http.strings import to_string
 from lightbug_http.net import default_buffer_size
 from lightbug_http.http import HTTPRequest, HTTPResponse, encode
 from lightbug_http.header import Headers, HeaderKey
 from lightbug_http.net import create_connection, TCPConnection
-from lightbug_http.io.bytes import Bytes
-from lightbug_http.utils import ByteReader, logger
-from lightbug_http.pool_manager import PoolManager
+from lightbug_http.io.bytes import Bytes, ByteReader
+from lightbug_http._logger import logger
+from lightbug_http.pool_manager import PoolManager, PoolKey
+from lightbug_http.uri import URI, Scheme
 
 
 struct Client:
@@ -41,7 +32,7 @@ struct Client:
         self.allow_redirects = allow_redirects
         self._connections = PoolManager[TCPConnection](cached_connections)
 
-    fn do(mut self, owned req: HTTPRequest) raises -> HTTPResponse:
+    fn do(mut self, owned request: HTTPRequest) raises -> HTTPResponse:
         """The `do` method is responsible for sending an HTTP request to a server and receiving the corresponding response.
 
         It performs the following steps:
@@ -54,81 +45,77 @@ struct Client:
         Note: The code assumes that the `HTTPRequest` object passed as an argument has a valid URI with a host and port specified.
 
         Args:
-            req: An `HTTPRequest` object representing the request to be sent.
+            request: An `HTTPRequest` object representing the request to be sent.
 
         Returns:
             The received response.
 
         Raises:
             Error: If there is a failure in sending or receiving the message.
         """
-        if req.uri.host == "":
-            raise Error("Client.do: Request failed because the host field is empty.")
-        var is_tls = False
+        if request.uri.host == "":
+            raise Error("Client.do: Host must not be empty.")
 
-        if req.uri.is_https():
+        var is_tls = False
+        var scheme = Scheme.HTTP
+        if request.uri.is_https():
             is_tls = True
+            scheme = Scheme.HTTPS
 
-        var host_str: String
-        var port: Int
-        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])
+        var port: UInt16
+        if request.uri.port:
+            port = request.uri.port.value()
         else:
-            host_str = req.uri.host
-            if is_tls:
+            if request.uri.scheme == Scheme.HTTP.value:
+                port = 80
+            elif request.uri.scheme == Scheme.HTTPS.value:
                 port = 443
             else:
-                port = 80
+                raise Error("Client.do: Invalid scheme received in the URI.")
 
+        var pool_key = PoolKey(request.uri.host, port, scheme)
         var cached_connection = False
         var conn: TCPConnection
         try:
-            conn = self._connections.take(host_str)
+            conn = self._connections.take(pool_key)
             cached_connection = True
         except e:
             if str(e) == "PoolManager.take: Key not found.":
-                conn = create_connection(host_str, port)
+                conn = create_connection(request.uri.host, port)
             else:
                 logger.error(e)
                 raise Error("Client.do: Failed to create a connection to host.")
 
         var bytes_sent: Int
         try:
-            bytes_sent = conn.write(encode(req))
+            bytes_sent = conn.write(encode(request))
         except e:
             # Maybe peer reset ungracefully, so try a fresh connection
             if str(e) == "SendError: Connection reset by peer.":
                 logger.debug("Client.do: Connection reset by peer. Trying a fresh connection.")
                 conn.teardown()
                 if cached_connection:
-                    return self.do(req^)
+                    return self.do(request^)
             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. (@thatstoasty)
         var new_buf = Bytes(capacity=default_buffer_size)
-
         try:
             _ = conn.read(new_buf)
         except e:
             if str(e) == "EOF":
                 conn.teardown()
                 if cached_connection:
-                    return self.do(req^)
+                    return self.do(request^)
                 raise Error("Client.do: No response received from the server.")
             else:
                 logger.error(e)
                 raise Error("Client.do: Failed to read response from peer.")
 
-        var res: HTTPResponse
+        var response: HTTPResponse
         try:
-            res = HTTPResponse.from_bytes(new_buf, conn)
+            response = HTTPResponse.from_bytes(new_buf, conn)
         except e:
             logger.error("Failed to parse a response...")
             try:
@@ -138,19 +125,19 @@ struct Client:
             raise e
 
         # Redirects should not keep the connection alive, as redirects can send the client to a different server.
-        if self.allow_redirects and res.is_redirect():
+        if self.allow_redirects and response.is_redirect():
             conn.teardown()
-            return self._handle_redirect(req^, res^)
+            return self._handle_redirect(request^, response^)
         # Server told the client to close the connection, we can assume the server closed their side after sending the response.
-        elif res.connection_close():
+        elif response.connection_close():
             conn.teardown()
         # Otherwise, persist the connection by giving it back to the pool manager.
         else:
-            self._connections.give(host_str, conn^)
-        return res
+            self._connections.give(pool_key, conn^)
+        return response
 
     fn _handle_redirect(
-        mut self, owned original_req: HTTPRequest, owned original_response: HTTPResponse
+        mut self, owned original_request: HTTPRequest, owned original_response: HTTPResponse
     ) raises -> HTTPResponse:
         var new_uri: URI
         var new_location: String
@@ -164,9 +151,9 @@ struct Client:
                 new_uri = URI.parse(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
+            original_request.headers[HeaderKey.HOST] = new_uri.host
         else:
-            new_uri = original_req.uri
+            new_uri = original_request.uri
             new_uri.path = new_location
-        original_req.uri = new_uri
-        return self.do(original_req^)
+        original_request.uri = new_uri
+        return self.do(original_request^)

lightbug_http/cookie/request_cookie_jar.mojo

@@ -3,7 +3,7 @@ from small_time import SmallTime, TimeZone
 from small_time.small_time import strptime
 from lightbug_http.strings import to_string, lineBreak
 from lightbug_http.header import HeaderKey, write_header
-from lightbug_http.utils import ByteReader, ByteWriter, is_newline, is_space
+from lightbug_http.io.bytes import ByteReader, ByteWriter, is_newline, is_space
 
 
 @value

lightbug_http/cookie/response_cookie_jar.mojo

@@ -1,7 +1,7 @@
 from collections import Optional, List, Dict, KeyElement
 from lightbug_http.strings import to_string
 from lightbug_http.header import HeaderKey, write_header
-from lightbug_http.utils import ByteWriter
+from lightbug_http.io.bytes import ByteWriter
 
 
 @value