Fix: Logs access was not authenticated correctly

Author: hoh

src/aleph/vm/orchestrator/views/operator.py

@@ -164,10 +164,8 @@ def get_signed_operation(request: web.Request) -> SignedOperation:
         raise web.HTTPBadRequest(reason="Invalid X-SignedOperation format") from error
 
 
-async def authenticate_jwk(request: web.Request) -> str:
-    signed_pubkey = get_signed_pubkey(request)
-    signed_operation = get_signed_operation(request)
-
+def verify_signed_operation(signed_operation: SignedOperation, signed_pubkey: SignedPubKeyHeader) -> str:
+    """Verify that the operation is signed by the ephemeral key authorized by the wallet."""
     if signed_pubkey.content.json_web_key.verify(
         data=signed_operation.payload,
         signature=signed_operation.signature,
@@ -176,7 +174,21 @@ async def authenticate_jwk(request: web.Request) -> str:
         logger.debug("Signature verified")
         return signed_pubkey.content.address
     else:
-        raise web.HTTPUnauthorized(reason="Signature could not verified") from None
+        raise web.HTTPUnauthorized(reason="Signature could not verified")
+
+
+async def authenticate_jwk(request: web.Request) -> str:
+    """Authenticate a request using the X-SignedPubKey and X-SignedOperation headers."""
+    signed_pubkey = get_signed_pubkey(request)
+    signed_operation = get_signed_operation(request)
+    return verify_signed_operation(signed_operation, signed_pubkey)
+
+
+async def authenicate_websocket_message(message) -> str:
+    """Authenticate a websocket message since JS cannot configure headers on WebSockets."""
+    signed_pubkey = SignedPubKeyHeader.parse_obj(message["X-SignedPubKey"])
+    signed_operation = SignedOperation.parse_obj(message["X-SignedOperation"])
+    return verify_signed_operation(signed_operation, signed_pubkey)
 
 
 def require_jwk_authentication(
@@ -219,24 +231,39 @@ def get_execution_or_404(ref: ItemHash, pool: VmPool) -> VmExecution:
         raise web.HTTPNotFound(body=f"No virtual machine with ref {ref}")
 
 
-@require_jwk_authentication
-async def stream_logs(request: web.Request, authenticated_sender: str) -> web.StreamResponse:
+async def stream_logs(request: web.Request) -> web.StreamResponse:
+    """Stream the logs of a VM.
+
+    The authentication method is slightly different because browsers do not
+    allow Javascript to set headers in WebSocket requests.
+    """
     vm_hash = get_itemhash_or_400(request.match_info)
     pool: VmPool = request.app["vm_pool"]
     execution = get_execution_or_404(vm_hash, pool=pool)
 
     if execution.vm is None:
         raise web.HTTPBadRequest(body=f"VM {vm_hash} is not running")
 
-    if execution.message.address != authenticated_sender:
-        logger.debug(f"Unauthorized sender {authenticated_sender} for {vm_hash}")
-        return web.Response(status=401, body="Unauthorized sender")
-
     queue: asyncio.Queue = asyncio.Queue(maxsize=1000)
     try:
         ws = web.WebSocketResponse()
         await ws.prepare(request)
+
         try:
+            # Authentication
+            first_message = await ws.receive_json()
+            credentials = first_message["auth"]
+            authenticated_sender = await authenicate_websocket_message(credentials)
+
+            if execution.message.address != authenticated_sender:
+                logger.debug(f"Denied request to access logs by {authenticated_sender} on {vm_hash}")
+                return web.Response(status=401, body="Unauthorized sender")
+            else:
+                logger.debug(f"Accepted request to access logs by {authenticated_sender} on {vm_hash}")
+                await ws.send_json({"status": "failed", "reason": "unauthorized sender"})
+
+            await ws.send_json({"status": "connected"})
+
             # Limit the number of queues per VM
             if len(execution.vm.fvm.log_queues) > 20:
                 logger.warning("Too many log queues, dropping the oldest one")