Skip to content

Commit

Permalink
follow redirects
Browse files Browse the repository at this point in the history
Signed-off-by: Brian Grenier <[email protected]>
  • Loading branch information
bgreni committed Sep 21, 2024
1 parent 805372e commit 6a4ca35
Show file tree
Hide file tree
Showing 6 changed files with 81 additions and 22 deletions.
4 changes: 2 additions & 2 deletions lightbug_http/header.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@ struct HeaderKey:
alias CONTENT_LENGTH = "content-length"
alias CONTENT_ENCODING = "content-encoding"
alias DATE = "date"
alias LOCATION = "location"
alias HOST = "host"


@value
Expand Down Expand Up @@ -70,8 +72,6 @@ struct Headers(Formattable, Stringable):
self._inner[key.lower()] = value

fn content_length(self) -> Int:
if HeaderKey.CONTENT_LENGTH not in self:
return 0
try:
return int(self[HeaderKey.CONTENT_LENGTH])
except:
Expand Down
29 changes: 27 additions & 2 deletions lightbug_http/http.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,14 @@ fn encode(owned res: HTTPResponse) -> Bytes:
return res._encoded()


struct StatusCode:
alias OK = 200
alias MOVED_PERMANENTLY = 301
alias FOUND = 302
alias TEMPORARY_REDIRECT = 307
alias PERMANENT_REDIRECT = 308
alias NOT_FOUND = 404

@value
struct HTTPRequest(Formattable, Stringable):
var headers: Headers
Expand Down Expand Up @@ -103,6 +111,8 @@ struct HTTPRequest(Formattable, Stringable):
self.set_content_length(len(body))
if HeaderKey.CONNECTION not in self.headers:
self.set_connection_close()
if HeaderKey.HOST not in self.headers:
self.headers[HeaderKey.HOST] = uri.host

fn set_connection_close(inout self):
self.headers[HeaderKey.CONNECTION] = "close"
Expand All @@ -126,8 +136,15 @@ struct HTTPRequest(Formattable, Stringable):
fn format_to(self, inout writer: Formatter):
writer.write(
self.method,
whitespace,
self.uri.path if len(self.uri.path) > 1 else strSlash,
whitespace
)
path = self.uri.path if len(self.uri.path) > 1 else strSlash
if len(self.uri.query_string) > 0:
path += "?" + self.uri.query_string

writer.write(path)

writer.write(
whitespace,
self.protocol,
lineBreak,
Expand All @@ -147,6 +164,8 @@ struct HTTPRequest(Formattable, Stringable):
writer.write(self.method)
writer.write(whitespace)
var path = self.uri.path if len(self.uri.path) > 1 else strSlash
if len(self.uri.query_string) > 0:
path += "?" + self.uri.query_string
writer.write(path)
writer.write(whitespace)
writer.write(self.protocol)
Expand Down Expand Up @@ -236,6 +255,12 @@ struct HTTPResponse(Formattable, Stringable):
fn set_content_length(inout self, l: Int):
self.headers[HeaderKey.CONTENT_LENGTH] = str(l)

fn is_redirect(self) -> Bool:
return self.status_code == StatusCode.MOVED_PERMANENTLY
or self.status_code == StatusCode.FOUND
or self.status_code == StatusCode.TEMPORARY_REDIRECT
or self.status_code == StatusCode.PERMANENT_REDIRECT

@always_inline
fn read_body(inout self, inout r: ByteReader) raises -> None:
r.consume(self.body_raw)
Expand Down
34 changes: 26 additions & 8 deletions lightbug_http/sys/client.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -12,26 +12,23 @@ from lightbug_http.strings import to_string
from lightbug_http.client import Client
from lightbug_http.net import default_buffer_size
from lightbug_http.http import HTTPRequest, HTTPResponse, encode
from lightbug_http.header import Headers
from lightbug_http.header import Headers, HeaderKey
from lightbug_http.sys.net import create_connection
from lightbug_http.io.bytes import Bytes
from lightbug_http.utils import ByteReader


struct MojoClient(Client):
var fd: c_int
var host: StringLiteral
var port: Int
var name: String

fn __init__(inout self) raises:
self.fd = socket(AF_INET, SOCK_STREAM, 0)
self.host = "127.0.0.1"
self.port = 8888
self.name = "lightbug_http_client"

fn __init__(inout self, host: StringLiteral, port: Int) raises:
self.fd = socket(AF_INET, SOCK_STREAM, 0)
self.host = host
self.port = port
self.name = "lightbug_http_client"
Expand Down Expand Up @@ -87,8 +84,10 @@ struct MojoClient(Client):
port = 443
else:
port = 80
var conn = create_connection(self.fd, host_str, port)
var bytes_sent = conn.write(encode(req^))

# TODO: Actually handle persistent connections
var conn = create_connection(socket(AF_INET, SOCK_STREAM, 0), host_str, port)
var bytes_sent = conn.write(encode(req))
if bytes_sent == -1:
raise Error("Failed to send message")

Expand All @@ -97,11 +96,30 @@ struct MojoClient(Client):

if bytes_recv == 0:
conn.close()

try:
return HTTPResponse.from_bytes(new_buf^)
var res = HTTPResponse.from_bytes(new_buf^)
if res.is_redirect():
conn.close()
return self._handle_redirect(req^, res^)
return res
except e:
conn.close()
raise e

return HTTPResponse(Bytes())

fn _handle_redirect(self,
owned original_req: HTTPRequest,
owned original_response: HTTPResponse) raises -> HTTPResponse:
var new_uri: URI
var new_location = original_response.headers[HeaderKey.LOCATION]
if new_location.startswith("http"):
new_uri = URI.parse_raises(new_location)
else:
new_uri = original_req.uri
new_uri.path = new_location

original_req.uri = new_uri
return self.do(original_req^)


2 changes: 1 addition & 1 deletion lightbug_http/uri.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ struct URI:
host_and_port = remainder_uri
request_uri = strSlash
self.host = host_and_port

if is_https:
self.scheme = https
else:
Expand Down
20 changes: 16 additions & 4 deletions tests/test_client.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -12,20 +12,32 @@ from lightbug_http.io.bytes import bytes
def test_client():
var mojo_client = MojoClient()
test_mojo_client_lightbug_external_req(mojo_client)
test_mojo_client_redirect_external_req(mojo_client)

fn test_mojo_client_redirect_external_req(client: MojoClient) raises:
var req = HTTPRequest(
uri=URI.parse_raises("http://httpbin.org/status/302"),
headers=Headers(
Header("Connection", "close")),
method="GET",
)
try:
var res = client.do(req)
testing.assert_equal(res.status_code, 200)
# testing.assert_equal(res.headers.content_length(), 323)
except e:
print(e)

fn test_mojo_client_lightbug_external_req(client: MojoClient) raises:
var req = HTTPRequest(
uri=URI.parse("http://httpbin.org/status/200")[URI],
uri=URI.parse_raises("http://httpbin.org/status/200"),
headers=Headers(
Header("Connection", "keep-alive"),
Header("Host", "httpbin.org")),
Header("Connection", "close")),
method="GET",
)

try:
var res = client.do(req)
testing.assert_equal(res.status_code, 200)

except e:
print(e)
14 changes: 9 additions & 5 deletions tests/test_http.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ def test_http():


def test_encode_http_request():
var uri = URI(default_server_conn_string + "/foobar?baz")
var uri = URI.parse_raises(default_server_conn_string + "/foobar?baz")
var req = HTTPRequest(
uri,
body=String("Hello world!").as_bytes(),
Expand All @@ -23,12 +23,16 @@ def test_encode_http_request():

var as_str = str(req)
var req_encoded = to_string(encode(req^))


var expected = String(
"GET /foobar?baz HTTP/1.1\r\nconnection: keep-alive\r\ncontent-length:"
" 12\r\nhost: localhost:8080\r\n\r\nHello world!"
)

testing.assert_equal(
req_encoded,
(
"GET / HTTP/1.1\r\nconnection: keep-alive\r\ncontent-length:"
" 12\r\n\r\nHello world!"
),
expected
)
testing.assert_equal(req_encoded, as_str)

Expand Down

0 comments on commit 6a4ca35

Please sign in to comment.