Skip to content

Commit

Permalink
fix: improve request header parsing (#166)
Browse files Browse the repository at this point in the history
  • Loading branch information
@jrriehl
jrriehl authored Sep 19, 2023
1 parent 56e86a7 commit b42d02d
Show file tree
Hide file tree
Showing 2 changed files with 111 additions and 8 deletions.
53 changes: 49 additions & 4 deletions python/src/uagents/asgi.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,9 @@ async def serve(self):
)
await self._server.serve()

async def __call__(self, scope, receive, send):
async def __call__(
self, scope, receive, send
): # pylint: disable=too-many-branches
"""
Handle an incoming ASGI message, dispatching the envelope to the appropriate handler,
and waiting for any queries to be resolved.
Expand All @@ -107,6 +109,43 @@ async def __call__(self, scope, receive, send):
return

headers = CaseInsensitiveDict(scope.get("headers", {}))

if b"content-type" not in headers:
# if connecting from browser, return a 200 OK
if b"user-agent" in headers:
await send(
{
"type": "http.response.start",
"status": 200,
"headers": [
[b"content-type", b"application/json"],
],
}
)
await send(
{
"type": "http.response.body",
"body": b'{"status": "OK - Agent is running"}',
}
)
else: # otherwise, return a 400 Bad Request
await send(
{
"type": "http.response.start",
"status": 400,
"headers": [
[b"content-type", b"application/json"],
],
}
)
await send(
{
"type": "http.response.body",
"body": b'{"error": "missing header: content-type"}',
}
)
return

if b"application/json" not in headers[b"content-type"]:
await send(
{
Expand All @@ -118,7 +157,10 @@ async def __call__(self, scope, receive, send):
}
)
await send(
{"type": "http.response.body", "body": b'{"error": "invalid format"}'}
{
"type": "http.response.body",
"body": b'{"error": "invalid content-type"}',
}
)
return

Expand All @@ -139,7 +181,10 @@ async def __call__(self, scope, receive, send):
}
)
await send(
{"type": "http.response.body", "body": b'{"error": "invalid format"}'}
{
"type": "http.response.body",
"body": b'{"error": "contents do not match envelope schema"}',
}
)
return

Expand All @@ -163,7 +208,7 @@ async def __call__(self, scope, receive, send):
await send(
{
"type": "http.response.body",
"body": b'{"error": "unable to verify payload"}',
"body": b'{"error": "signature verification failed"}',
}
)
return
Expand Down
66 changes: 62 additions & 4 deletions python/tests/test_server.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -297,7 +297,7 @@ async def test_message_fail_wrong_headers(self):
call(
{
"type": "http.response.body",
"body": b'{"error": "invalid format"}',
"body": b'{"error": "invalid content-type"}',
}
),
]
Expand Down Expand Up @@ -329,7 +329,7 @@ async def test_message_fail_bad_data(self):
call(
{
"type": "http.response.body",
"body": b'{"error": "invalid format"}',
"body": b'{"error": "contents do not match envelope schema"}',
}
),
]
Expand Down Expand Up @@ -370,7 +370,7 @@ async def test_message_fail_unsigned(self):
call(
{
"type": "http.response.body",
"body": b'{"error": "unable to verify payload"}',
"body": b'{"error": "signature verification failed"}',
}
),
]
Expand Down Expand Up @@ -412,7 +412,7 @@ async def test_message_fail_verify(self):
call(
{
"type": "http.response.body",
"body": b'{"error": "unable to verify payload"}',
"body": b'{"error": "signature verification failed"}',
}
),
]
Expand Down Expand Up @@ -460,6 +460,64 @@ async def test_message_fail_dispatch(self):
]
)

async def test_request_fail_missing_header(self):
mock_send = AsyncMock()
await self.agent._server(
scope={
"type": "http",
"path": "/submit",
"headers": {},
},
receive=None,
send=mock_send,
)
mock_send.assert_has_calls(
[
call(
{
"type": "http.response.start",
"status": 400,
"headers": [[b"content-type", b"application/json"]],
}
),
call(
{
"type": "http.response.body",
"body": b'{"error": "missing header: content-type"}',
}
),
]
)

async def test_request_from_browser(self):
mock_send = AsyncMock()
await self.agent._server(
scope={
"type": "http",
"path": "/submit",
"headers": {b"User-Agent": b"Mozilla/5.0"},
},
receive=None,
send=mock_send,
)
mock_send.assert_has_calls(
[
call(
{
"type": "http.response.start",
"status": 200,
"headers": [[b"content-type", b"application/json"]],
}
),
call(
{
"type": "http.response.body",
"body": b'{"status": "OK - Agent is running"}',
}
),
]
)


if __name__ == "__main__":
unittest.main()

0 comments on commit b42d02d

Please sign in to comment.