mirror of
https://github.com/GNS3/gns3-server.git
synced 2026-09-03 00:25:17 +03:00
fix: tear down controller console WebSocket forwarding cleanly on client disconnect
Closing a web console while the node kept streaming output crashed the ws_console/vnc_console handlers with an uncaught WebSocketDisconnect from the compute-to-client send path (only the opposite direction was guarded), producing a full ASGI traceback on every console close. Restructure both endpoints as symmetric forwarding tasks managed with asyncio.wait(FIRST_COMPLETED): exceptions from either direction are collected as task exceptions, the peer task is cancelled, the compute WebSocket is closed, and the client is notified. The receive loops now close and log on every exit path instead of only in the exception branch. Tests patch the in-memory ASGI transport to deliver a conformant websocket.disconnect on client close (it sends a non-conformant websocket.close that starlette receive() rejects).
This commit is contained in:
YueGuobin
parent
861b7e86ff
commit
eed981fee3
@ -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}")
|
||||
|
||||
|
||||
@ -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
|
||||
)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user