From 1783ed22a273093c37e4e9bc145555367b0a7895 Mon Sep 17 00:00:00 2001
From: Brian Grenier <grenierb96@gmail.com>
Date: Mon, 28 Oct 2024 14:35:11 -0600
Subject: [PATCH] Add support for chunked transfer encoding

---
 client.mojo                      |  5 +----
 lightbug_http/client.mojo        |  7 ++++---
 lightbug_http/http/response.mojo | 32 +++++++++++++++++++++++++++++++-
 lightbug_http/utils.mojo         |  5 +++++
 4 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/client.mojo b/client.mojo
index e37f9b2..ebe3a4a 100644
--- a/client.mojo
+++ b/client.mojo
@@ -12,10 +12,7 @@ fn test_request(inout client: Client) raises -> None:
     # print status code
     print("Response:", response.status_code)
 
-    # print parsed headers (only some are parsed for now)
-    print("Content-Type:", response.headers["Content-Type"])
-    print("Content-Length", response.headers["Content-Length"])
-    print("Server:", to_string(response.headers["Server"]))
+    print(response.headers)
 
     print(
         "Is connection set to connection-close? ", response.connection_close()
diff --git a/lightbug_http/client.mojo b/lightbug_http/client.mojo
index e77466c..f6ded40 100644
--- a/lightbug_http/client.mojo
+++ b/lightbug_http/client.mojo
@@ -123,7 +123,7 @@ struct Client:
                 return self.do(req^)
             raise Error("No response received")
         try:
-            var res = HTTPResponse.from_bytes(new_buf^)
+            var res = HTTPResponse.from_bytes(new_buf^, conn)
             if res.is_redirect():
                 self._close_conn(host_str)
                 return self._handle_redirect(req^, res^)
@@ -151,5 +151,6 @@ struct Client:
         return self.do(original_req^)
 
     fn _close_conn(inout self, host: String) raises:
-        self._connections[host].close()
-        _ = self._connections.pop(host)
+        if host in self._connections:
+            self._connections[host].close()
+            _ = self._connections.pop(host)
diff --git a/lightbug_http/http/response.mojo b/lightbug_http/http/response.mojo
index 7990770..10bf49b 100644
--- a/lightbug_http/http/response.mojo
+++ b/lightbug_http/http/response.mojo
@@ -12,6 +12,9 @@ from lightbug_http.strings import (
     lineBreak,
     to_string,
 )
+from collections import Optional
+from utils import StringSlice
+from lightbug_http.net import SysConnection, default_buffer_size
 
 
 struct StatusCode:
@@ -33,7 +36,7 @@ struct HTTPResponse(Formattable, Stringable):
     var protocol: String
 
     @staticmethod
-    fn from_bytes(owned b: Bytes) raises -> HTTPResponse:
+    fn from_bytes(owned b: Bytes, conn: Optional[SysConnection] = None) raises -> HTTPResponse:
         var reader = ByteReader(b^)
 
         var headers = Headers()
@@ -43,6 +46,7 @@ struct HTTPResponse(Formattable, Stringable):
 
         try:
             protocol, status_code, status_text = headers.parse_raw(reader)
+            reader.skip_newlines()
         except e:
             raise Error("Failed to parse response headers: " + e.__str__())
 
@@ -54,6 +58,18 @@ struct HTTPResponse(Formattable, Stringable):
             status_text=status_text,
         )
 
+        if response.headers[HeaderKey.TRANSFER_ENCODING] == "chunked":
+            var b = Bytes()
+            reader.consume(b)
+
+            var buff = Bytes(capacity=default_buffer_size)
+            while conn.value().read(buff) > 0:
+                b += buff
+                buff.resize(0)
+
+            response.read_chunks(b^)
+            return response
+
         try:
             response.read_body(reader)
             return response
@@ -111,6 +127,7 @@ struct HTTPResponse(Formattable, Stringable):
         except:
             return 0
 
+    @always_inline
     fn is_redirect(self) -> Bool:
         return (
             self.status_code == StatusCode.MOVED_PERMANENTLY
@@ -124,6 +141,19 @@ struct HTTPResponse(Formattable, Stringable):
         r.consume(self.body_raw, self.content_length())
         self.set_content_length(len(self.body_raw))
 
+    fn read_chunks(inout self, owned chunks: Bytes) raises:
+        var reader = ByteReader(chunks^)
+
+        while True:
+            var size = atol(StringSlice(unsafe_from_utf8=reader.read_line()), 16)
+            if size == 0:
+                break
+            var data = Bytes()
+            reader.consume(data, size)
+            reader.skip_newlines()
+            self.set_content_length(self.content_length() + len(data))
+            self.body_raw += data
+
     fn format_to(self, inout writer: Formatter):
         writer.write(self.protocol, whitespace, self.status_code, whitespace, self.status_text, lineBreak)
 
diff --git a/lightbug_http/utils.mojo b/lightbug_http/utils.mojo
index 7055edc..e527e69 100644
--- a/lightbug_http/utils.mojo
+++ b/lightbug_http/utils.mojo
@@ -85,6 +85,11 @@ struct ByteReader:
         while is_space(self.peek()):
             self.increment()
 
+    @always_inline
+    fn skip_newlines(inout self):
+        while self.peek() == BytesConstant.rChar:
+            self.increment(2)
+
     @always_inline
     fn increment(inout self, v: Int = 1):
         self.read_pos += v