NgTech-LTD/gns3-server
2
0
Fork 0
mirror of https://github.com/GNS3/gns3-server.git synced 2026-09-01 15:54:03 +03:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
fix: tear down controller console WebSocket forwarding cleanly on client disconnect
This commit is contained in:
Jeremy Grossmann 2026-08-31 19:02:02 +02:00 committed by GitHub
commit 20819f35ae
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 311 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

106
gns3server/api/routes/controller/nodes.py
View File

@ -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}")

240
tests/api/routes/controller/test_nodes.py
View File

@ -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
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
)