diff --git a/nest/session_routes.py b/nest/session_routes.py index 5edb255..35f681d 100644 --- a/nest/session_routes.py +++ b/nest/session_routes.py @@ -76,7 +76,10 @@ def _same_origin(websocket: WebSocket) -> bool: origin = websocket.headers.get("origin") if origin is None: return True # Non-browser API clients must still authenticate. - parsed = urlsplit(origin) + try: + parsed = urlsplit(origin) + except ValueError: + return False scheme = "https" if websocket.url.scheme == "wss" else "http" return parsed.scheme == scheme and parsed.netloc.lower() == websocket.headers.get("host", "").lower() and not parsed.path @@ -111,6 +114,7 @@ async def terminal(websocket: WebSocket, vm_id: str, session_id: str): sessions = websocket.app.state.nest_sessions await sessions.get(vm_id, session_id) offset = auth.get("offset", 0) + acknowledgements = asyncio.Queue(maxsize=1) async def output(): cursor = offset @@ -124,6 +128,10 @@ async def terminal(websocket: WebSocket, vm_id: str, session_id: str): if data: await websocket.send_json({"type": "output", "offset": cursor, "data": base64.b64encode(data).decode("ascii")}) cursor += len(data) + # Only one output chunk may be waiting for the terminal renderer. + acknowledged = await asyncio.wait_for(acknowledgements.get(), timeout=30) + if type(acknowledged) is not int or acknowledged != cursor: + raise ProtocolError("invalid terminal output acknowledgement") continue if record["status"] not in {"active", "starting"}: return @@ -136,6 +144,10 @@ async def terminal(websocket: WebSocket, vm_id: str, session_id: str): await sessions.write(vm_id, session_id, terminal_data(message["data"])) elif set(message) == {"type", "rows", "cols"} and message["type"] == "resize": await sessions.resize(vm_id, session_id, message["rows"], message["cols"]) + elif set(message) == {"type", "offset"} and message["type"] == "ack": + if acknowledgements.full(): + raise ProtocolError("unexpected terminal output acknowledgement") + acknowledgements.put_nowait(message["offset"]) else: raise ProtocolError("invalid terminal control message") diff --git a/nest/sessions.py b/nest/sessions.py index 884457a..d9f84c5 100644 --- a/nest/sessions.py +++ b/nest/sessions.py @@ -7,7 +7,7 @@ from dataclasses import dataclass, field from pathlib import Path from uvm.app import Application -from uvm.errors import ConflictError, GuestExecutionError, StateError, UvmError +from uvm.errors import ConflictError, GuestExecutionError, NotFoundError, StateError, UvmError from uvm.guest_terminal import GuestTerminal, validate_size from .session_store import SessionStore @@ -15,6 +15,7 @@ from .session_store import SessionStore @dataclass class _LiveSession: + vm_id: str terminal: GuestTerminal changed: asyncio.Event = field(default_factory=asyncio.Event) task: asyncio.Task | None = None @@ -78,7 +79,7 @@ class SessionManager: try: terminal, username = await asyncio.to_thread(self.application.lifecycle.open_terminal, vm_id, rows=rows, cols=cols) await asyncio.to_thread(history.activate, record["id"], username) - live = _LiveSession(terminal) + live = _LiveSession(vm_id, terminal) self._live[record["id"]] = live live.task = asyncio.create_task(self._read(history, record["id"], live)) except BaseException as error: @@ -126,10 +127,12 @@ class SessionManager: live.changed.clear() async def _active(self, vm_id: str, session_id: str) -> _LiveSession: - record = await self.get(vm_id, session_id) live = self._live.get(session_id) - if record["status"] != "active" or live is None: + if live is None: + await self.get(vm_id, session_id) raise ConflictError("This session has ended. Launch a new session to run commands.") + if live.vm_id != vm_id: + raise NotFoundError("Session not found for this VM.") return live async def write(self, vm_id: str, session_id: str, data: bytes) -> None: diff --git a/uvm/guest_agent.py b/uvm/guest_agent.py index 19b2911..e25c405 100644 --- a/uvm/guest_agent.py +++ b/uvm/guest_agent.py @@ -353,7 +353,7 @@ def run_terminal(connection: socket.socket, username: str, rows: int, cols: int, selector.modify(master, selectors.EVENT_READ | (selectors.EVENT_WRITE if pending else 0), "pty") if _exited(process): # Drain the leader's final output before terminating any background jobs. - while True: + for _ in range(16): try: data = os.read(master, MAX_TERMINAL_BYTES) except (BlockingIOError, OSError):