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