Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: improve request header parsing #166

Merged
merged 1 commit into from
Sep 19, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 49 additions & 4 deletions python/src/uagents/asgi.py
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
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()
Loading