Merge pull request #2867 from yueguobin/fix/controller-console-ws-disconnect-traceback

fix: tear down controller console WebSocket forwarding cleanly on client disconnect

gns3server/api/routes/controller/nodes.py

@@ -20,6 +20,7 @@ API routes for nodes.
 
 import aiohttp
 import asyncio
+import contextlib
 import ipaddress
 
 from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect, Request, Response, status, Query, HTTPException
@@ -707,14 +708,34 @@ async def ws_console(
                 elif "bytes" in msg and msg["bytes"]:
                     await ws_console_compute.send_bytes(msg["bytes"])
         except WebSocketDisconnect:
-            await ws_console_compute.close()
+            pass
+        log.info(
+            f"Client {websocket.client.host}:{websocket.client.port} has disconnected from controller"
+            f" console WebSocket"
+        )
+
+    async def ws_send(ws_console_compute):
+        """
+        Receive WebSocket data from compute console WebSocket and forward to client.
+        """
+
+        try:
+            async for msg in ws_console_compute:
+                if msg.type == aiohttp.WSMsgType.TEXT:
+                    await websocket.send_text(msg.data)
+                elif msg.type == aiohttp.WSMsgType.BINARY:
+                    await websocket.send_bytes(msg.data)
+                elif msg.type == aiohttp.WSMsgType.ERROR:
+                    break
+        except WebSocketDisconnect:
+            # the client disconnected while the compute was still streaming console output
             log.info(
                 f"Client {websocket.client.host}:{websocket.client.port} has disconnected from controller"
                 f" console WebSocket"
             )
 
     try:
-        # receive WebSocket data from compute console WebSocket and forward to client.
+        # forward WebSocket data in both directions between the client and the compute console WebSocket
         log.info(f"Forwarding console WebSocket to '{ws_console_compute_url}'")
         server_config = Config.instance().settings.Server
         user = server_config.compute_username
@@ -728,20 +749,20 @@ async def ws_console(
             auth = aiohttp.BasicAuth(user, "")
         ssl_context = Controller.instance().ssl_context()
         async with HTTPClient.get_client().ws_connect(ws_console_compute_url, auth=auth, ssl_context=ssl_context) as ws:
-            asyncio.ensure_future(ws_receive(ws))
+            tasks = [
-            async for msg in ws:
+                asyncio.ensure_future(ws_receive(ws)),
-                if msg.type == aiohttp.WSMsgType.TEXT:
+                asyncio.ensure_future(ws_send(ws)),
-                    await websocket.send_text(msg.data)
+            ]
-                elif msg.type == aiohttp.WSMsgType.BINARY:
+            done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
-                    await websocket.send_bytes(msg.data)
+            for task in done:
-                elif msg.type == aiohttp.WSMsgType.ERROR:
+                if task.exception():
-                    break
+                    log.warning(f"Exception while forwarding console WebSocket data: {task.exception()}")
-    except WebSocketDisconnect:
+            for task in pending:
-        # the client disconnected while the compute was still streaming console output
+                task.cancel()
-        log.info(
+            await asyncio.gather(*pending, return_exceptions=True)
-            f"Client {websocket.client.host}:{websocket.client.port} has disconnected from controller"
+            # notify the client that the console session has ended (ignore if the client is already gone)
-            f" console WebSocket"
+            with contextlib.suppress(WebSocketDisconnect):
-        )
+                await websocket.close()
     except aiohttp.ClientError as e:
         log.error(f"Client error received when forwarding to compute console WebSocket: {e}")
 
@@ -786,18 +807,39 @@ async def vnc_console(
 
         try:
             while True:
-                data = await websocket.receive_bytes()
+                msg = await websocket.receive()
+                if msg["type"] == "websocket.disconnect":
+                    break
+                data = msg.get("bytes")
                 if data:
                     await vnc_console_compute.send_bytes(data)
         except WebSocketDisconnect:
-            await vnc_console_compute.close()
+            pass
+        log.info(
+            f"Client {websocket.client.host}:{websocket.client.port} has disconnected from controller"
+            f" VNC console WebSocket"
+        )
+
+    async def vnc_send(vnc_console_compute):
+        """
+        Receive binary WebSocket data from compute VNC console WebSocket and forward to client.
+        """
+
+        try:
+            async for msg in vnc_console_compute:
+                if msg.type == aiohttp.WSMsgType.BINARY:
+                    await websocket.send_bytes(msg.data)
+                elif msg.type == aiohttp.WSMsgType.ERROR:
+                    break
+        except WebSocketDisconnect:
+            # the client disconnected while the compute was still streaming VNC console output
             log.info(
                 f"Client {websocket.client.host}:{websocket.client.port} has disconnected from controller"
                 f" VNC console WebSocket"
             )
 
     try:
-        # receive binary data from compute VNC console WebSocket and forward to client
+        # forward WebSocket data in both directions between the client and the compute VNC console WebSocket
         log.info(f"Forwarding VNC console WebSocket to '{vnc_console_compute_url}'")
         server_config = Config.instance().settings.Server
         user = server_config.compute_username
@@ -811,18 +853,20 @@ async def vnc_console(
             auth = aiohttp.BasicAuth(user, "")
         ssl_context = Controller.instance().ssl_context()
         async with HTTPClient.get_client().ws_connect(vnc_console_compute_url, auth=auth, ssl_context=ssl_context) as ws:
-            asyncio.ensure_future(vnc_receive(ws))
+            tasks = [
-            async for msg in ws:
+                asyncio.ensure_future(vnc_receive(ws)),
-                if msg.type == aiohttp.WSMsgType.BINARY:
+                asyncio.ensure_future(vnc_send(ws)),
-                    await websocket.send_bytes(msg.data)
+            ]
-                elif msg.type == aiohttp.WSMsgType.ERROR:
+            done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
-                    break
+            for task in done:
-    except WebSocketDisconnect:
+                if task.exception():
-        # the client disconnected while the compute was still streaming VNC console output
+                    log.warning(f"Exception while forwarding VNC console WebSocket data: {task.exception()}")
-        log.info(
+            for task in pending:
-            f"Client {websocket.client.host}:{websocket.client.port} has disconnected from controller"
+                task.cancel()
-            f" VNC console WebSocket"
+            await asyncio.gather(*pending, return_exceptions=True)
-        )
+            # notify the client that the VNC session has ended (ignore if the client is already gone)
+            with contextlib.suppress(WebSocketDisconnect):
+                await websocket.close()
     except aiohttp.ClientError as e:
         log.error(f"Client error received when forwarding to compute VNC console WebSocket: {e}")
 

tests/api/routes/controller/test_nodes.py

@@ -16,9 +16,11 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 
-import aiohttp
+import asyncio
+import contextlib
 import logging
 
+import aiohttp
 import pytest
 
 from types import SimpleNamespace
@@ -27,8 +29,11 @@ from typing import List, Optional
 from fastapi import FastAPI, HTTPException, WebSocketDisconnect, status
 from httpx import AsyncClient
 from pydantic import SecretStr
+from httpx_ws import aconnect_ws
+from httpx_ws import WebSocketDisconnect as HttpxWebSocketDisconnect
+from httpx_ws.transport import ASGIWebSocketTransport, ASGIWebSocketAsyncNetworkStream
 
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 from tests.utils import AsyncioMagicMock
 
 from gns3server.config import Config
@@ -37,6 +42,8 @@ from gns3server.controller.project import Project
 from gns3server.controller.compute import Compute
 from gns3server.utils.http_client import HTTPClient
 from gns3server.api.routes.controller.nodes import ws_console, vnc_console
+from gns3server.services import auth_service
+from gns3server.services.authentication import DEFAULT_JWT_SECRET_KEY
 
 pytestmark = pytest.mark.asyncio
 
@@ -681,6 +688,8 @@ class TestNodeRoutes:
     #     compute.http_query.assert_called_with("POST", "/projects/{project_id}/files/project-files/vpcs/{node_id}/hello/nested".format(project_id=project.id, node_id=node.id), data=b'hello', timeout=None, raw=True)
 
 
+
+
 class FakeComputeConsoleWebSocket:
     """
     Stand-in for the aiohttp ClientWebSocketResponse returned when connecting
@@ -738,13 +747,21 @@ class FakeClientWebSocket:
 
     async def receive(self) -> dict:
 
-        # the client is already gone from the receive side
+        # the client never sends anything else and never completes a close
-        return {"type": "websocket.disconnect"}
+        # handshake: block until the forwarding task is cancelled, like a real
+        # receive side waiting for a disconnect message that only arrives on
+        # the send path
+        await asyncio.Event().wait()
 
     async def receive_bytes(self) -> bytes:
 
         raise WebSocketDisconnect(code=1006)
 
+    async def close(self, code: int = 1000, reason: Optional[str] = None) -> None:
+
+        # notification that the console session ended; nothing to record
+        pass
+
     async def send_text(self, data: str) -> None:
 
         self._check_client_alive()
@@ -855,3 +872,218 @@ class TestNodeConsoleWebSocketRoutes:
 
         assert websocket.sent == []
         assert compute_ws.closed
+
+
+class _FakeComputeWebSocket:
+    """
+    Mimics the aiohttp client WebSocket the controller forwards console traffic to.
+
+    Pass a list of WSMessage to emulate a compute that sends data then closes the
+    session, or None to emulate a busy console streaming binary frames forever
+    (until the controller cancels the forwarding task).
+    """
+
+    def __init__(self, messages=None) -> None:
+        self._messages = messages
+        self.sent = []  # frames received from the client
+        self.closed = False
+        self.stream_cancelled = False
+
+    async def send_str(self, data: str) -> None:
+        self.sent.append(("text", data))
+
+    async def send_bytes(self, data: bytes) -> None:
+        self.sent.append(("bytes", data))
+
+    async def close(self) -> None:
+        self.closed = True
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        if self._messages is not None:
+            if not self._messages:
+                raise StopAsyncIteration
+            return self._messages.pop(0)
+        try:
+            await asyncio.sleep(0.05)
+        except asyncio.CancelledError:
+            self.stream_cancelled = True
+            raise
+        return aiohttp.WSMessage(aiohttp.WSMsgType.BINARY, b"console output", "")
+
+
+class _FakeComputeWebSocketContext:
+    """Async context manager mimicking the object returned by aiohttp ws_connect()."""
+
+    def __init__(self, websocket: _FakeComputeWebSocket) -> None:
+        self._websocket = websocket
+
+    async def __aenter__(self) -> _FakeComputeWebSocket:
+        return self._websocket
+
+    async def __aexit__(self, *exc_info) -> bool:
+        await self._websocket.close()
+        return False
+
+
+# The in-memory ASGI WebSocket transport notifies the app of a client close with a
+# non-conformant "websocket.close" message, which starlette's receive() rejects.
+# Patch it to deliver "websocket.disconnect" like a real ASGI server (uvicorn) does.
+_original_stream_send = ASGIWebSocketAsyncNetworkStream.send
+
+
+async def _conforming_stream_send(self, message):
+    if message.get("type") == "websocket.close":
+        message = {"type": "websocket.disconnect", "code": message.get("code") or 1000}
+    await _original_stream_send(self, message)
+
+
+class TestConsoleWebSocketRoutes:
+    """
+    Walk the console/VNC WebSocket endpoints through the ASGI stack: data is
+    forwarded in both directions and both sides tear down cleanly when the
+    other goes away, instead of letting a WebSocketDisconnect escape as an
+    ASGI error.
+    """
+
+    @pytest.fixture
+    def node(self, project: Project, compute: Compute) -> Node:
+
+        node = Node(project, compute, "test", node_type="qemu")
+        project._nodes[node.id] = node
+        return node
+
+    @staticmethod
+    def _patches(fake_ws: _FakeComputeWebSocket):
+        """
+        Make HTTPClient.get_client().ws_connect() yield the fake compute WebSocket,
+        and make the in-memory transport deliver a conformant disconnect message.
+        """
+
+        stack = contextlib.ExitStack()
+        http_client = MagicMock()
+        http_client.ws_connect = MagicMock(return_value=_FakeComputeWebSocketContext(fake_ws))
+        stack.enter_context(patch(
+            "gns3server.api.routes.controller.nodes.HTTPClient.get_client",
+            return_value=http_client
+        ))
+        stack.enter_context(patch.object(ASGIWebSocketAsyncNetworkStream, "send", _conforming_stream_send))
+        return stack
+
+    @staticmethod
+    def _admin_token() -> str:
+
+        return auth_service.create_access_token("admin", secret_key=DEFAULT_JWT_SECRET_KEY)
+
+    @staticmethod
+    async def _wait_for_teardown(fake_ws: _FakeComputeWebSocket) -> None:
+        # the handler keeps running on the event loop after the client is gone
+        for _ in range(100):
+            if fake_ws.closed:
+                await asyncio.sleep(0.1)  # give cancelled tasks a chance to finish
+                return
+            await asyncio.sleep(0.05)
+
+    async def test_console_ws_client_disconnect(
+            self,
+            app: FastAPI,
+            client: AsyncClient,
+            project: Project,
+            compute: Compute,
+            node: Node,
+            caplog
+    ) -> None:
+        """
+        A client disconnecting while the compute keeps streaming console output
+        must not raise: the forwarding tasks are cancelled and the compute
+        console WebSocket is closed.
+        """
+
+        fake_ws = _FakeComputeWebSocket(messages=None)
+        with self._patches(fake_ws), caplog.at_level(logging.INFO):
+            async with AsyncClient(base_url="http://test-api", transport=ASGIWebSocketTransport(app=app)) as ws_client:
+                async with aconnect_ws(
+                    app.url_path_for("ws_console", project_id=project.id, node_id=node.id),
+                    ws_client,
+                    params={"token": self._admin_token()},
+                ) as ws:
+                    # compute -> client
+                    assert await ws.receive_bytes() == b"console output"
+                    # client -> compute
+                    await ws.send_text("dir")
+                    await ws.send_bytes(b"\x01\x02")
+            await self._wait_for_teardown(fake_ws)
+
+        assert fake_ws.sent == [("text", "dir"), ("bytes", b"\x01\x02")]
+        assert fake_ws.stream_cancelled, "the compute -> client forwarding task should have been cancelled"
+        assert fake_ws.closed, "the compute console WebSocket should have been closed"
+        assert any(
+            "has disconnected from controller console WebSocket" in record.getMessage()
+            for record in caplog.records
+        )
+
+    async def test_console_ws_compute_closes_session(
+            self,
+            app: FastAPI,
+            client: AsyncClient,
+            project: Project,
+            compute: Compute,
+            node: Node
+    ) -> None:
+        """
+        When the compute closes the console WebSocket the client should receive
+        the frames sent before the close, then the close itself.
+        """
+
+        fake_ws = _FakeComputeWebSocket(messages=[
+            aiohttp.WSMessage(aiohttp.WSMsgType.BINARY, b"output", ""),
+            aiohttp.WSMessage(aiohttp.WSMsgType.TEXT, "done", ""),
+        ])
+        with self._patches(fake_ws):
+            async with AsyncClient(base_url="http://test-api", transport=ASGIWebSocketTransport(app=app)) as ws_client:
+                async with aconnect_ws(
+                    app.url_path_for("ws_console", project_id=project.id, node_id=node.id),
+                    ws_client,
+                    params={"token": self._admin_token()},
+                ) as ws:
+                    assert await ws.receive_bytes() == b"output"
+                    assert await ws.receive_text() == "done"
+                    with pytest.raises(HttpxWebSocketDisconnect):
+                        await ws.receive_bytes()
+
+        assert fake_ws.closed
+
+    async def test_vnc_console_ws_client_disconnect(
+            self,
+            app: FastAPI,
+            client: AsyncClient,
+            project: Project,
+            compute: Compute,
+            node: Node,
+            caplog
+    ) -> None:
+        """
+        Same as the console test, for the VNC endpoint (binary frames only).
+        """
+
+        fake_ws = _FakeComputeWebSocket(messages=None)
+        with self._patches(fake_ws), caplog.at_level(logging.INFO):
+            async with AsyncClient(base_url="http://test-api", transport=ASGIWebSocketTransport(app=app)) as ws_client:
+                async with aconnect_ws(
+                    app.url_path_for("vnc_console", project_id=project.id, node_id=node.id),
+                    ws_client,
+                    params={"token": self._admin_token()},
+                ) as ws:
+                    assert await ws.receive_bytes() == b"console output"
+                    await ws.send_bytes(b"\x01\x02")
+            await self._wait_for_teardown(fake_ws)
+
+        assert fake_ws.sent == [("bytes", b"\x01\x02")]
+        assert fake_ws.stream_cancelled, "the compute -> client forwarding task should have been cancelled"
+        assert fake_ws.closed, "the compute VNC console WebSocket should have been closed"
+        assert any(
+            "has disconnected from controller VNC console WebSocket" in record.getMessage()
+            for record in caplog.records
+        )