From d8c9268c674eccd1b5d03cdb9a3f1a29ff6609a8 Mon Sep 17 00:00:00 2001
From: James Riehl <james.riehl@fetch.ai>
Date: Tue, 19 Sep 2023 11:16:19 +0100
Subject: [PATCH] fix: improve request header parsing

---
 python/src/uagents/asgi.py  | 53 ++++++++++++++++++++++++++---
 python/tests/test_server.py | 66 ++++++++++++++++++++++++++++++++++---
 2 files changed, 111 insertions(+), 8 deletions(-)

diff --git a/python/src/uagents/asgi.py b/python/src/uagents/asgi.py
index 641b3617..9867b79a 100644
--- a/python/src/uagents/asgi.py
+++ b/python/src/uagents/asgi.py
@@ -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.
@@ -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(
                 {
@@ -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
 
@@ -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
 
@@ -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
diff --git a/python/tests/test_server.py b/python/tests/test_server.py
index c904e7e1..a58fb063 100644
--- a/python/tests/test_server.py
+++ b/python/tests/test_server.py
@@ -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"}',
                     }
                 ),
             ]
@@ -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"}',
                     }
                 ),
             ]
@@ -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"}',
                     }
                 ),
             ]
@@ -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"}',
                     }
                 ),
             ]
@@ -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()