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 committed 2026-08-31 22:43:41 +08:00
1 parent 861b7e86ff
commit eed981fee3
2 files changed
+311 -35

No files matched your search

+75 -31
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}")