Skip to content

Commit

Permalink
Add support for chunked transfer encoding
Browse files Browse the repository at this point in the history
  • Loading branch information
bgreni committed Oct 28, 2024
1 parent cf8f81d commit 1783ed2
Show file tree
Hide file tree
Showing 4 changed files with 41 additions and 8 deletions.
5 changes: 1 addition & 4 deletions client.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
7 changes: 4 additions & 3 deletions lightbug_http/client.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -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^)
Expand Down Expand Up @@ -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)
32 changes: 31 additions & 1 deletion lightbug_http/http/response.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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()
Expand All @@ -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__())

Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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)

Expand Down
5 changes: 5 additions & 0 deletions lightbug_http/utils.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 1783ed2

Please sign in to comment.