Nexus Agent
341b130d9b
独立 guacd reader 避免 WS 竞争丢数据；ping 回显与 sync 已生产验证；前端 90s 超时与 OPEN 后发尺寸。 Co-authored-by: Cursor <cursoragent@cursor.com>
128 lines
4.0 KiB
Python
128 lines
4.0 KiB
Python
"""Guacamole WebSocket tunnel: handshake with guacd then relay (not raw TCP bridge)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import logging
|
|
|
|
from fastapi import WebSocket, WebSocketDisconnect
|
|
|
|
from server.infrastructure.guacamole.protocol import (
|
|
GuacdStream,
|
|
guacd_rdp_handshake,
|
|
is_internal_instruction,
|
|
iter_raw_instructions,
|
|
tunnel_open_instruction,
|
|
)
|
|
|
|
logger = logging.getLogger("nexus.guacamole.ws_tunnel")
|
|
|
|
|
|
async def _guacd_reader_loop(stream: GuacdStream, out: asyncio.Queue[bytes | None]) -> None:
|
|
"""Continuously read complete guacd instructions into a queue (never cancel mid-read)."""
|
|
try:
|
|
while True:
|
|
batch = await stream.read_raw_batch()
|
|
if not batch:
|
|
break
|
|
await out.put(batch)
|
|
except ConnectionError:
|
|
pass
|
|
except Exception as exc:
|
|
logger.warning("guacd reader ended: %s", exc)
|
|
finally:
|
|
await out.put(None)
|
|
|
|
|
|
async def _relay_ws_guacd(
|
|
websocket: WebSocket,
|
|
stream: GuacdStream,
|
|
writer: asyncio.StreamWriter,
|
|
) -> None:
|
|
"""Single WS consumer; guacd reads run in a dedicated task (Starlette-safe)."""
|
|
guacd_queue: asyncio.Queue[bytes | None] = asyncio.Queue(maxsize=512)
|
|
reader_task = asyncio.create_task(_guacd_reader_loop(stream, guacd_queue))
|
|
ws_lock = asyncio.Lock()
|
|
|
|
async def ws_send(text: str) -> None:
|
|
async with ws_lock:
|
|
await websocket.send_text(text)
|
|
|
|
async def forward_client_message(data: str) -> None:
|
|
for inst in iter_raw_instructions(data):
|
|
if is_internal_instruction(inst):
|
|
await ws_send(inst)
|
|
else:
|
|
writer.write(inst.encode("utf-8"))
|
|
await writer.drain()
|
|
|
|
try:
|
|
while True:
|
|
ws_recv_task = asyncio.create_task(websocket.receive_text())
|
|
guacd_recv_task = asyncio.create_task(guacd_queue.get())
|
|
done, pending = await asyncio.wait(
|
|
{ws_recv_task, guacd_recv_task},
|
|
return_when=asyncio.FIRST_COMPLETED,
|
|
)
|
|
for task in pending:
|
|
task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await task
|
|
|
|
if ws_recv_task in done:
|
|
exc = ws_recv_task.exception()
|
|
if exc is not None:
|
|
if isinstance(exc, WebSocketDisconnect):
|
|
return
|
|
raise exc
|
|
await forward_client_message(ws_recv_task.result())
|
|
|
|
if guacd_recv_task in done:
|
|
batch = guacd_recv_task.result()
|
|
if batch is None:
|
|
return
|
|
await ws_send(batch.decode("utf-8", errors="replace"))
|
|
finally:
|
|
reader_task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await reader_task
|
|
|
|
|
|
async def run_guacamole_rdp_tunnel(
|
|
websocket: WebSocket,
|
|
guacd_host: str,
|
|
guacd_port: int,
|
|
connect_params: dict[str, str],
|
|
*,
|
|
width: int = 1024,
|
|
height: int = 768,
|
|
) -> None:
|
|
"""Handshake with guacd using server-side credentials, then bridge WS ↔ guacd."""
|
|
reader, writer = await asyncio.open_connection(guacd_host, guacd_port)
|
|
stream = GuacdStream(reader)
|
|
|
|
try:
|
|
conn_id = await guacd_rdp_handshake(
|
|
stream, writer, connect_params, width=width, height=height,
|
|
)
|
|
await websocket.send_text(tunnel_open_instruction(conn_id))
|
|
except Exception as exc:
|
|
logger.warning("guacd handshake failed: %s", exc)
|
|
writer.close()
|
|
with contextlib.suppress(Exception):
|
|
await writer.wait_closed()
|
|
await websocket.close(code=1011, reason=str(exc)[:120])
|
|
return
|
|
|
|
try:
|
|
await _relay_ws_guacd(websocket, stream, writer)
|
|
except WebSocketDisconnect:
|
|
pass
|
|
except Exception as exc:
|
|
logger.warning("RDP relay ended: %s", exc)
|
|
finally:
|
|
writer.close()
|
|
with contextlib.suppress(Exception):
|
|
await writer.wait_closed()
|