Fix #427 websocket --admin-api-key connect issue

This is the proposed option number 2. It uses the websocket
to expect a message with the admin api key before it is being
used to transmit any other messages than 'ping'

Signed-off-by: Matthias Binzer <[email protected]>

aries_cloudagent/admin/server.py

@@ -172,6 +172,9 @@ async def make_application(self) -> web.Application:
             @web.middleware
             async def check_token(request, handler):
                 header_admin_api_key = request.headers.get("x-api-key")
+                # always allow websockets here
+                if request.path == "/ws":
+                    return await handler(request)
                 if not header_admin_api_key:
                     raise web.HTTPUnauthorized()
 
@@ -355,6 +358,17 @@ async def websocket_handler(self, request):
         queue = BasicMessageQueue()
         loop = asyncio.get_event_loop()
 
+        admin_api_key = self.context.settings.get("admin.admin_api_key")
+        admin_insecure_mode = self.context.settings.get("admin.admin_insecure_mode")
+        if admin_insecure_mode:
+            # open to send websocket messages without api key auth
+            queue.api_key_authenticated = True
+        else:
+            header_admin_api_key = request.headers.get("x-api-key")
+            if header_admin_api_key == admin_api_key:
+                # authenticated via http header
+                queue.api_key_authenticated = True
+
         try:
             self.websocket_queues[socket_id] = queue
             await queue.enqueue(
@@ -372,7 +386,7 @@ async def websocket_handler(self, request):
             )
 
             closed = False
-            receive = loop.create_task(ws.receive())
+            receive = loop.create_task(ws.receive_json())
             send = loop.create_task(queue.dequeue(timeout=5.0))
 
             while not closed:
@@ -384,9 +398,22 @@ async def websocket_handler(self, request):
                         closed = True
 
                     if receive.done():
-                        # ignored
                         if not closed:
-                            receive = loop.create_task(ws.receive())
+                            msg_received = None
+                            msg_api_key = None
+                            try:
+                                # this call can re-raise exeptions from inside the task
+                                msg_received = receive.result()
+                                msg_api_key = msg_received.get("x-api-key")
+                            except Exception as ex:
+                                LOGGER.error("Exception in websocket receiving task:")
+                                LOGGER.exception(ex)
+                            if admin_api_key:
+                                if admin_api_key == msg_api_key:
+                                    # authenticated via websocket message
+                                    queue.api_key_authenticated = True
+
+                            receive = loop.create_task(ws.receive_json())
 
                     if send.done():
                         try:
@@ -397,7 +424,10 @@ async def websocket_handler(self, request):
                         if msg is None:
                             # we send fake pings because the JS client
                             # can't detect real ones
-                            msg = {"topic": "ping"}
+                            msg = {
+                                "topic": "ping",
+                                "authenticated": queue.api_key_authenticated,
+                            }
                         if not closed:
                             if msg:
                                 await ws.send_json(msg)
@@ -441,4 +471,5 @@ async def send_webhook(self, topic: str, payload: dict):
                     )
 
         for queue in self.websocket_queues.values():
-            await queue.enqueue({"topic": topic, "payload": payload})
+            if queue.api_key_authenticated:
+                await queue.enqueue({"topic": topic, "payload": payload})

aries_cloudagent/admin/tests/test_admin_server.py

@@ -218,6 +218,35 @@ async def test_status_secure(self):
         result = await resp.json()
         assert isinstance(result, dict)
 
+    @unittest_run_loop
+    async def test_websocket_with_api_key_message(self):
+        async with self.client.ws_connect("/ws") as ws:
+            result = await ws.receive_json()
+            assert result["topic"] == "settings"
+
+            ping1 = await ws.receive_json()
+            assert ping1["topic"] == "ping"
+            assert ping1["authenticated"] == False
+
+            await ws.send_json({"dummy": ""})
+            ping2 = await ws.receive_json()
+            assert ping2["authenticated"] == False
+
+            await ws.send_json({"x-api-key": self.TEST_API_KEY})
+            ping3 = await ws.receive_json()
+            assert ping3["authenticated"] == True
+
+    @unittest_run_loop
+    async def test_websocket_with_api_key_header(self):
+        async with self.client.ws_connect(
+            "/ws", headers={"x-api-key": self.TEST_API_KEY}
+        ) as ws:
+            result = await ws.receive_json()
+            assert result["topic"] == "settings"
+
+            ping1 = await ws.receive_json()
+            assert ping1["authenticated"] == True
+
 
 class TestAdminServerWebhook(AioHTTPTestCase):
     async def setUpAsync(self):

aries_cloudagent/transport/queue/basic.py

@@ -14,6 +14,7 @@ def __init__(self):
         self.queue = self.make_queue()
         self.logger = logging.getLogger(__name__)
         self.stop_event = asyncio.Event()
+        self.api_key_authenticated = False
 
     def make_queue(self):
         """Create the queue instance."""