diff --git a/gns3server/api/routes/controller/nodes.py b/gns3server/api/routes/controller/nodes.py index ed17a703e..482c56146 100644 --- a/gns3server/api/routes/controller/nodes.py +++ b/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)) - async for msg in ws: - 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" - ) + tasks = [ + asyncio.ensure_future(ws_receive(ws)), + asyncio.ensure_future(ws_send(ws)), + ] + done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + for task in done: + if task.exception(): + log.warning(f"Exception while forwarding console WebSocket data: {task.exception()}") + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + # notify the client that the console 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 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)) - async for msg in ws: - 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" - ) + tasks = [ + asyncio.ensure_future(vnc_receive(ws)), + asyncio.ensure_future(vnc_send(ws)), + ] + done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + for task in done: + if task.exception(): + log.warning(f"Exception while forwarding VNC console WebSocket data: {task.exception()}") + for task in pending: + task.cancel() + 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}") diff --git a/tests/api/routes/controller/test_nodes.py b/tests/api/routes/controller/test_nodes.py index 02ca2c143..be115b1dd 100644 --- a/tests/api/routes/controller/test_nodes.py +++ b/tests/api/routes/controller/test_nodes.py @@ -16,9 +16,11 @@ # along with this program. If not, see . -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 - return {"type": "websocket.disconnect"} + # the client never sends anything else and never completes a close + # 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 + )