Use pico for http basics

lightbug_http/header.mojo

@@ -1,8 +1,15 @@
 from collections import Dict, Optional
+from memory import UnsafePointer, Span
 from lightbug_http.io.bytes import Bytes, ByteReader, ByteWriter, is_newline, is_space
 from lightbug_http.strings import BytesConstant
 from lightbug_http._logger import logger
 from lightbug_http.strings import rChar, nChar, lineBreak, to_string
+from lightbug_http.pico import (
+    PhrHeader,
+    phr_parse_request,
+    phr_parse_response,
+    phr_parse_headers,
+)
 
 
 struct HeaderKey:
@@ -84,31 +91,158 @@ struct Headers(Writable, Stringable, Copyable, Movable):
             return 0
 
     fn parse_raw(mut self, mut r: ByteReader) raises -> (String, String, String, List[String]):
+        """Parse HTTP headers using picohttpparser for request/response.
+
+        This method delegates to parse_raw_request or parse_raw_response based on
+        whether the first token looks like an HTTP method or HTTP version.
+
+        Returns:
+            For requests: (method, path, protocol, cookies)
+            For responses: (protocol, status_code, status_text, cookies)
+        """
+        # Peek at first few bytes to determine if this is a request or response
         var first_byte = r.peek()
         if not first_byte:
-            raise Error("Headers.parse_raw: Failed to read first byte from response header")
+            raise Error("Headers.parse_raw: Failed to read first byte from header")
+
+        # Create buffer from ByteReader's remaining data
+        var buf_span = r._inner[r.read_pos:]
+        var buf_ptr = UnsafePointer[UInt8].alloc(len(buf_span))
+        for i in range(len(buf_span)):
+            buf_ptr[i] = buf_span[i]
+
+        # Check if starts with "HTTP/" (response) or method name (request)
+        var is_response = (
+            len(buf_span) >= 5
+            and buf_span[0] == ord('H')
+            and buf_span[1] == ord('T')
+            and buf_span[2] == ord('T')
+            and buf_span[3] == ord('P')
+            and buf_span[4] == ord('/')
+        )
+
+        var bytes_consumed: Int
+        var result: (String, String, String, List[String])
+        if is_response:
+            var parse_result = self._parse_raw_response(buf_ptr, len(buf_span))
+            bytes_consumed = parse_result[0]
+            result = parse_result[1]
+        else:
+            var parse_result = self._parse_raw_request(buf_ptr, len(buf_span))
+            bytes_consumed = parse_result[0]
+            result = parse_result[1]
+
+        buf_ptr.free()
+
+        # Advance ByteReader position to start of body (after headers end)
+        r.read_pos += bytes_consumed
+
+        return result
+
+    fn _parse_raw_request(
+        mut self,
+        buf_ptr: UnsafePointer[UInt8],
+        buf_len: Int
+    ) raises -> (Int, (String, String, String, List[String])):
+        """Parse HTTP request using picohttpparser."""
+        var method = String()
+        var method_len = 0
+        var path = String()
+        var path_len = 0
+        var minor_version = -1
+
+        # Allocate headers array (max 100 headers)
+        var max_headers = 100
+        var headers = UnsafePointer[PhrHeader].alloc(max_headers)
+        for i in range(max_headers):
+            headers[i] = PhrHeader()
+
+        var num_headers = max_headers
+        var ret = phr_parse_request(
+            buf_ptr, buf_len,
+            method, method_len,
+            path, path_len,
+            minor_version,
+            headers,
+            num_headers,
+            0  # last_len (0 for first parse)
+        )
+
+        if ret < 0:
+            headers.free()
+            if ret == -1:
+                raise Error("Headers.parse_raw: Invalid HTTP request")
+            else:  # ret == -2
+                raise Error("Headers.parse_raw: Incomplete HTTP request")
 
-        var first = r.read_word()
-        r.increment()
-        var second = r.read_word()
-        r.increment()
-        var third = r.read_line()
+        # Extract headers and cookies
         var cookies = List[String]()
+        for i in range(num_headers):
+            var key = headers[i].name.lower()
+            var value = headers[i].value
+
+            if key == HeaderKey.SET_COOKIE or key == HeaderKey.COOKIE:
+                cookies.append(value)
+            else:
+                self._inner[key] = value
+
+        # Build protocol string
+        var protocol = "HTTP/1." + String(minor_version)
+
+        headers.free()
+        return (ret, (method, path, protocol, cookies^))
+
+    fn _parse_raw_response(
+        mut self,
+        buf_ptr: UnsafePointer[UInt8],
+        buf_len: Int
+    ) raises -> (Int, (String, String, String, List[String])):
+        """Parse HTTP response using picohttpparser."""
+        var minor_version = -1
+        var status = 0
+        var msg = String()
+        var msg_len = 0
+
+        # Allocate headers array (max 100 headers)
+        var max_headers = 100
+        var headers = UnsafePointer[PhrHeader].alloc(max_headers)
+        for i in range(max_headers):
+            headers[i] = PhrHeader()
+
+        var num_headers = max_headers
+        var ret = phr_parse_response(
+            buf_ptr, buf_len,
+            minor_version,
+            status,
+            msg, msg_len,
+            headers,
+            num_headers,
+            0  # last_len (0 for first parse)
+        )
+
+        if ret < 0:
+            headers.free()
+            if ret == -1:
+                raise Error("Headers.parse_raw: Invalid HTTP response")
+            else:  # ret == -2
+                raise Error("Headers.parse_raw: Incomplete HTTP response")
+
+        # Extract headers and cookies
+        var cookies = List[String]()
+        for i in range(num_headers):
+            var key = headers[i].name.lower()
+            var value = headers[i].value
+
+            if key == HeaderKey.SET_COOKIE:
+                cookies.append(value)
+            else:
+                self._inner[key] = value
+
+        # Build protocol string
+        var protocol = "HTTP/1." + String(minor_version)
 
-        while not is_newline(r.peek()):
-            var key = r.read_until(BytesConstant.colon)
-            r.increment()
-            if is_space(r.peek()):
-                r.increment()
-            # TODO (bgreni): Handle possible trailing whitespace
-            var value = r.read_line()
-            var k = String(key).lower()
-            if k == HeaderKey.SET_COOKIE:
-                cookies.append(String(value))
-                continue
-
-            self._inner[k] = String(value)
-        return (String(first), String(second), String(third), cookies^)
+        headers.free()
+        return (ret, (protocol, String(status), msg, cookies^))
 
     fn write_to[T: Writer, //](self, mut writer: T):
         for header in self._inner.items():

lightbug_http/http/response.mojo

@@ -1,8 +1,12 @@
 from collections import Optional
+from memory import UnsafePointer
 from lightbug_http.external.small_time.small_time import now
 from lightbug_http.uri import URI
 from lightbug_http.io.bytes import Bytes, bytes, byte, ByteReader, ByteWriter
 from lightbug_http.connection import TCPConnection, default_buffer_size
+from lightbug_http.header import Headers, HeaderKey
+from lightbug_http.cookie import ResponseCookieJar
+from lightbug_http._logger import logger
 from lightbug_http.strings import (
     strHttp11,
     strHttp,
@@ -13,6 +17,7 @@ from lightbug_http.strings import (
     lineBreak,
     to_string,
 )
+from lightbug_http.pico import PhrChunkedDecoder, phr_decode_chunked
 
 
 struct StatusCode:
@@ -93,23 +98,31 @@ struct HTTPResponse(Writable, Stringable, Encodable, Sized, Movable):
 
         var transfer_encoding = response.headers.get(HeaderKey.TRANSFER_ENCODING)
         if transfer_encoding and transfer_encoding.value() == "chunked":
+            # Use pico's chunked decoder for proper RFC-compliant parsing
+            var decoder = PhrChunkedDecoder()
+            decoder.consume_trailer = True  # Consume trailing headers
+
             var b = reader.read_bytes().to_bytes()
             var buff = Bytes(capacity=default_buffer_size)
+
             try:
+                # Read chunks from connection
                 while conn.read(buff) > 0:
                     b += buff.copy()
 
+                    # Check if we've reached the end of chunked data (0\r\n\r\n)
                     if (
-                        buff[-5] == byte("0")
+                        len(buff) >= 5
+                        and buff[-5] == byte("0")
                         and buff[-4] == byte("\r")
                         and buff[-3] == byte("\n")
                         and buff[-2] == byte("\r")
                         and buff[-1] == byte("\n")
                     ):
                         break
 
-                    # buff.clear()
-                response.read_chunks(b)
+                # Decode chunks using pico
+                response._decode_chunks_pico(decoder, b)
                 return response^
             except e:
                 logger.error(e)
@@ -224,16 +237,36 @@ struct HTTPResponse(Writable, Stringable, Encodable, Sized, Movable):
         self.body_raw = r.read_bytes(self.content_length()).to_bytes()
         self.set_content_length(len(self.body_raw))
 
-    fn read_chunks(mut self, chunks: Span[Byte]) raises:
-        var reader = ByteReader(chunks)
-        while True:
-            var size = atol(String(reader.read_line()), 16)
-            if size == 0:
-                break
-            var data = reader.read_bytes(size).to_bytes()
-            reader.skip_carriage_return()
-            self.set_content_length(self.content_length() + len(data))
-            self.body_raw += data^
+    fn _decode_chunks_pico(mut self, mut decoder: PhrChunkedDecoder, var chunks: Bytes) raises:
+        """Decode chunked transfer encoding using picohttpparser.
+
+        Args:
+            decoder: The chunked decoder state machine.
+            chunks: The raw chunked data to decode.
+        """
+        # Convert Bytes to UnsafePointer for pico API
+        var buf_ptr = UnsafePointer[UInt8].alloc(len(chunks))
+        for i in range(len(chunks)):
+            buf_ptr[i] = chunks[i]
+
+        var bufsz = len(chunks)
+        var result = phr_decode_chunked(decoder, buf_ptr, bufsz)
+        var ret = result[0]
+        var decoded_size = result[1]
+
+        if ret == -1:
+            buf_ptr.free()
+            raise Error("HTTPResponse._decode_chunks_pico: Invalid chunked encoding")
+        # ret == -2 means incomplete, but we'll proceed with what we have
+        # ret >= 0 means complete, with ret bytes of trailing data
+
+        # Copy decoded data to body
+        self.body_raw = Bytes(capacity=decoded_size)
+        for i in range(decoded_size):
+            self.body_raw.append(buf_ptr[i])
+
+        self.set_content_length(len(self.body_raw))
+        buf_ptr.free()
 
     fn write_to[T: Writer](self, mut writer: T):
         writer.write(self.protocol, whitespace, self.status_code, whitespace, self.status_text, lineBreak)

lightbug_http/io/bytes.mojo

@@ -239,7 +239,11 @@ struct ByteReader[origin: Origin](Sized):
         self.read_pos += count
         return self._inner[start : start + count]
 
+
     fn read_until(mut self, char: Byte) -> ByteView[origin]:
+        """Read bytes until a specific character is found.
+
+        """
         var start = self.read_pos
         for i in range(start, len(self._inner)):
             if self._inner[i] == char:
@@ -250,9 +254,15 @@ struct ByteReader[origin: Origin](Sized):
 
     @always_inline
     fn read_word(mut self) -> ByteView[origin]:
+        """Read bytes until whitespace is found.
+
+        """
         return self.read_until(BytesConstant.whitespace)
 
     fn read_line(mut self) -> ByteView[origin]:
+        """Read bytes until newline (CRLF or LF) is found.
+
+        """
         var start = self.read_pos
         for i in range(start, len(self._inner)):
             if is_newline(self._inner[i]):