From 9541e29b763a2efac27fea8100892adb127393f8 Mon Sep 17 00:00:00 2001 From: kstyagi Date: Fri, 25 Sep 2026 19:52:26 +0530 Subject: [PATCH] feat: implement session management and terminal handling for VM sessions --- nest/server.py | 18 +++++ nest/session_routes.py | 159 +++++++++++++++++++++++++++++++++++++ nest/session_store.py | 169 +++++++++++++++++++++++++++++++++++++++ nest/sessions.py | 174 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 520 insertions(+) create mode 100644 nest/session_routes.py create mode 100644 nest/session_store.py create mode 100644 nest/sessions.py diff --git a/nest/server.py b/nest/server.py index cb036b0..bf9f4cc 100644 --- a/nest/server.py +++ b/nest/server.py @@ -2,6 +2,7 @@ from __future__ import annotations +from contextlib import asynccontextmanager from pathlib import Path from typing import Any @@ -21,12 +22,29 @@ def create_app(application: Application | None = None, *, host: str | None = Non from fastapi.staticfiles import StaticFiles from .dashboard.router import router as dashboard_router + from .session_routes import router as session_router + from .sessions import SessionManager + + sessions = SessionManager(api.state.uvm_application) + api.state.nest_sessions = sessions + previous_lifespan = api.router.lifespan_context + + @asynccontextmanager + async def lifespan(app): + async with previous_lifespan(app): + try: + yield + finally: + await sessions.shutdown() + + api.router.lifespan_context = lifespan api.title = "Nest" api.description = "Nest workspace API, with local Firecracker machine management powered by UVM." static_dir = Path(__file__).parent / "dashboard" / "static" api.mount("/dashboard/static", StaticFiles(directory=static_dir), name="dashboard_static") api.include_router(dashboard_router, prefix="/dashboard") + api.include_router(session_router) @api.get("/") async def root() -> Any: diff --git a/nest/session_routes.py b/nest/session_routes.py new file mode 100644 index 0000000..5edb255 --- /dev/null +++ b/nest/session_routes.py @@ -0,0 +1,159 @@ +"""VM-scoped session history, creation, and authenticated browser terminal streams.""" + +from __future__ import annotations + +import asyncio +import base64 +import hmac +import json +from contextlib import suppress +from typing import Literal +from urllib.parse import urlsplit + +from fastapi import APIRouter, Depends, Request, WebSocket, WebSocketDisconnect +from pydantic import BaseModel, ConfigDict, Field, StrictInt + +from uvm.app import Application +from uvm.errors import UvmError +from uvm.guest_agent import MAX_TERMINAL_FRAME, ProtocolError, terminal_data +from uvm.routers.dependencies import get_authorized_application + +from .sessions import SessionManager + + +router = APIRouter(prefix="/vms/{vm_id}/sessions", tags=["sessions"]) + + +class SessionCreate(BaseModel): + model_config = ConfigDict(extra="forbid") + rows: StrictInt = Field(default=24, ge=2, le=200) + cols: StrictInt = Field(default=80, ge=2, le=400) + + +class SessionResponse(BaseModel): + id: str + vm_id: str + name: str + username: str + status: Literal["starting", "active", "closed", "interrupted", "failed"] + created_at: float + updated_at: float + ended_at: float | None + rows: int + cols: int + exit_code: int | None + error: str | None + output_start: int + output_end: int + truncated: bool + + +def manager(request: Request, _application: Application = Depends(get_authorized_application)) -> SessionManager: + return request.app.state.nest_sessions + + +@router.get("", response_model=list[SessionResponse]) +async def list_sessions(vm_id: str, sessions: SessionManager = Depends(manager)): + return await sessions.list(vm_id) + + +@router.post("", response_model=SessionResponse, status_code=201) +async def create_session(vm_id: str, body: SessionCreate, sessions: SessionManager = Depends(manager)): + return await sessions.create(vm_id, body.rows, body.cols) + + +@router.get("/{session_id}", response_model=SessionResponse) +async def get_session(vm_id: str, session_id: str, sessions: SessionManager = Depends(manager)): + return await sessions.get(vm_id, session_id) + + +@router.post("/{session_id}/close", response_model=SessionResponse) +async def close_session(vm_id: str, session_id: str, sessions: SessionManager = Depends(manager)): + return await sessions.end(vm_id, session_id) + + +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) + 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 + + +async def _message(websocket: WebSocket, *, limit: int = MAX_TERMINAL_FRAME) -> dict: + text = await websocket.receive_text() + if len(text) > limit: + raise ProtocolError("terminal message exceeds its size limit") + message = json.loads(text) + if not isinstance(message, dict): + raise ProtocolError("terminal message must be an object") + return message + + +@router.websocket("/{session_id}/terminal") +async def terminal(websocket: WebSocket, vm_id: str, session_id: str): + if not _same_origin(websocket): + await websocket.close(code=1008) + return + await websocket.accept() + tasks = [] + try: + # Browser WebSockets cannot set X-UVM-Token. Authenticate in the first + # frame so tokens never appear in URLs, access logs, or session history. + auth = await asyncio.wait_for(_message(websocket, limit=4096), timeout=10) + token = auth.get("token") + expected = websocket.app.state.uvm_application.settings.api_token + if (auth.get("type") != "auth" or not isinstance(token, str) or not token.isascii() + or not expected or not hmac.compare_digest(token, expected)): + await websocket.close(code=4401) + return + sessions = websocket.app.state.nest_sessions + await sessions.get(vm_id, session_id) + offset = auth.get("offset", 0) + + async def output(): + cursor = offset + previous = None + while True: + record, start, data = await sessions.read(vm_id, session_id, cursor) + if previous != record or start != cursor: + await websocket.send_json({"type": "session", "session": record, "offset": start}) + previous = record + cursor = start + if data: + await websocket.send_json({"type": "output", "offset": cursor, "data": base64.b64encode(data).decode("ascii")}) + cursor += len(data) + continue + if record["status"] not in {"active", "starting"}: + return + await sessions.wait(session_id) + + async def controls(): + while True: + message = await _message(websocket) + if set(message) == {"type", "data"} and message["type"] == "input": + 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"]) + else: + raise ProtocolError("invalid terminal control message") + + tasks = [asyncio.create_task(output()), asyncio.create_task(controls())] + done, _pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + for task in done: + task.result() + except WebSocketDisconnect: + pass + except (ValueError, KeyError, TypeError, UnicodeError, RecursionError, TimeoutError): + with suppress(WebSocketDisconnect, RuntimeError): + await websocket.send_json({"type": "error", "message": "Invalid terminal message or authentication timeout."}) + except UvmError as error: + with suppress(WebSocketDisconnect, RuntimeError): + await websocket.send_json({"type": "error", "message": str(error)}) + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + with suppress(WebSocketDisconnect, RuntimeError): + await websocket.close() diff --git a/nest/session_store.py b/nest/session_store.py new file mode 100644 index 0000000..bd7a1bc --- /dev/null +++ b/nest/session_store.py @@ -0,0 +1,169 @@ +"""Private, durable session metadata and bounded PTY transcripts, separate from VM state.""" + +from __future__ import annotations + +import fcntl +import os +import sqlite3 +import stat +import threading +import time +import uuid +from contextlib import contextmanager +from pathlib import Path + +from uvm.errors import ConflictError, NotFoundError, StateError, ValidationError + + +MAX_TRANSCRIPT_BYTES = 8 * 1024 * 1024 + + +class SessionStore: + def __init__(self, path: Path, *, max_output: int = MAX_TRANSCRIPT_BYTES) -> None: + self.path = path + self.max_output = max_output + self._mutex = threading.RLock() + self._owner = None + try: + path.parent.mkdir(parents=True, mode=0o700, exist_ok=True) + if path.parent.is_symlink(): + raise StateError("session history directory must not be a symlink") + path.parent.chmod(0o700) + self._owner = self._private_file(path.with_suffix(".lock")) + try: + fcntl.flock(self._owner, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError as error: + raise ConflictError("Another Nest server is already managing sessions for this VM registry.") from error + descriptor = self._private_file(path) + os.close(descriptor) + with self._db() as db: + db.executescript(""" + CREATE TABLE IF NOT EXISTS sessions ( + id TEXT PRIMARY KEY, vm_id TEXT NOT NULL, name TEXT NOT NULL, + username TEXT NOT NULL, status TEXT NOT NULL, + created_at REAL NOT NULL, updated_at REAL NOT NULL, ended_at REAL, + rows INTEGER NOT NULL, cols INTEGER NOT NULL, exit_code INTEGER, error TEXT, + output_start INTEGER NOT NULL DEFAULT 0, output_end INTEGER NOT NULL DEFAULT 0, + truncated INTEGER NOT NULL DEFAULT 0 + ); + CREATE INDEX IF NOT EXISTS sessions_vm ON sessions(vm_id, created_at); + CREATE TABLE IF NOT EXISTS output ( + session_id TEXT NOT NULL REFERENCES sessions(id), offset INTEGER NOT NULL, + data BLOB NOT NULL, PRIMARY KEY (session_id, offset) + ); + """) + now = time.time() + db.execute("""UPDATE sessions SET status='interrupted', ended_at=?, updated_at=?, + error='Nest restarted; this shell is no longer connected.' + WHERE status IN ('starting', 'active')""", (now, now)) + except BaseException: + self.close() + raise + + @staticmethod + def _private_file(path: Path) -> int: + descriptor = os.open(path, os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW, 0o600) + metadata = os.fstat(descriptor) + if not stat.S_ISREG(metadata.st_mode) or metadata.st_nlink != 1: + os.close(descriptor) + raise StateError("session history must use private regular files") + os.fchmod(descriptor, 0o600) + return descriptor + + @contextmanager + def _db(self): + with self._mutex: + db = None + try: + db = sqlite3.connect(self.path, timeout=5) + db.row_factory = sqlite3.Row + with db: + yield db + except sqlite3.Error as error: + raise StateError("Could not read or save terminal session history.") from error + finally: + if db is not None: + db.close() + + @staticmethod + def _record(row) -> dict: + if row is None: + raise NotFoundError("Session not found for this VM.") + result = dict(row) + result["truncated"] = bool(result["truncated"]) + return result + + def create(self, vm_id: str, rows: int, cols: int) -> dict: + session_id = "session-" + uuid.uuid4().hex + now = time.time() + with self._db() as db: + number = db.execute("SELECT COUNT(*) FROM sessions WHERE vm_id=?", (vm_id,)).fetchone()[0] + 1 + db.execute("""INSERT INTO sessions + (id, vm_id, name, username, status, created_at, updated_at, rows, cols) + VALUES (?, ?, ?, '', 'starting', ?, ?, ?, ?)""", + (session_id, vm_id, f"Session {number}", now, now, rows, cols)) + return self.get(vm_id, session_id) + + def list(self, vm_id: str) -> list[dict]: + with self._db() as db: + return [self._record(row) for row in db.execute( + "SELECT * FROM sessions WHERE vm_id=? ORDER BY created_at DESC, id DESC", (vm_id,), + )] + + def get(self, vm_id: str, session_id: str) -> dict: + with self._db() as db: + return self._record(db.execute("SELECT * FROM sessions WHERE vm_id=? AND id=?", (vm_id, session_id)).fetchone()) + + def activate(self, session_id: str, username: str) -> None: + with self._db() as db: + db.execute("UPDATE sessions SET status='active', username=?, updated_at=? WHERE id=?", + (username, time.time(), session_id)) + + def resize(self, session_id: str, rows: int, cols: int) -> None: + with self._db() as db: + db.execute("UPDATE sessions SET rows=?, cols=? WHERE id=?", (rows, cols, session_id)) + + def finish(self, session_id: str, status: str, exit_code: int | None = None, error: str | None = None) -> None: + now = time.time() + with self._db() as db: + db.execute("""UPDATE sessions SET status=?, ended_at=?, updated_at=?, exit_code=?, error=? + WHERE id=? AND status IN ('starting', 'active')""", (status, now, now, exit_code, error, session_id)) + + def append(self, session_id: str, data: bytes) -> None: + if not data: + return + with self._db() as db: + row = db.execute("SELECT output_end FROM sessions WHERE id=?", (session_id,)).fetchone() + if row is None: + raise NotFoundError("Session not found.") + offset = row[0] + end = offset + len(data) + start = max(0, end - self.max_output) + db.execute("INSERT INTO output VALUES (?, ?, ?)", (session_id, offset, data)) + db.execute("DELETE FROM output WHERE session_id=? AND offset + length(data)<=?", (session_id, start)) + first = db.execute("SELECT offset, data FROM output WHERE session_id=? AND offset 0), time.time(), session_id)) + + def read(self, vm_id: str, session_id: str, offset: int, limit: int = 65536) -> tuple[dict, int, bytes]: + with self._db() as db: + session = self._record(db.execute("SELECT * FROM sessions WHERE vm_id=? AND id=?", (vm_id, session_id)).fetchone()) + if type(offset) is not int or not 0 <= offset <= session["output_end"]: + raise ValidationError("invalid terminal output offset") + start = max(offset, session["output_start"]) + data = bytearray() + for row in db.execute("""SELECT offset, data FROM output + WHERE session_id=? AND offset + length(data)>? ORDER BY offset""", (session_id, start)): + part = row["data"][max(0, start - row["offset"]):] + data.extend(part[:limit - len(data)]) + if len(data) >= limit: + break + return session, start, bytes(data) + + def close(self) -> None: + if self._owner is not None: + os.close(self._owner) + self._owner = None diff --git a/nest/sessions.py b/nest/sessions.py new file mode 100644 index 0000000..884457a --- /dev/null +++ b/nest/sessions.py @@ -0,0 +1,174 @@ +"""Nest-owned live terminal sessions; browser attachments do not own guest shells.""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from pathlib import Path + +from uvm.app import Application +from uvm.errors import ConflictError, GuestExecutionError, StateError, UvmError +from uvm.guest_terminal import GuestTerminal, validate_size + +from .session_store import SessionStore + + +@dataclass +class _LiveSession: + terminal: GuestTerminal + changed: asyncio.Event = field(default_factory=asyncio.Event) + task: asyncio.Task | None = None + + +class SessionManager: + def __init__(self, application: Application) -> None: + self.application = application + self._store: SessionStore | None = None + self._live: dict[str, _LiveSession] = {} + self._pending: set[asyncio.Task] = set() + self._create_lock = asyncio.Lock() + self._closing = False + + def _history(self) -> SessionStore: + if self._closing: + raise ConflictError("Nest is shutting down.") + if self._store is None: + try: + self._store = SessionStore(Path(self.application.settings.base) / "nest" / "sessions.sqlite3") + except OSError as error: + raise StateError("Could not open private terminal history; run Nest with access to the VM registry.") from error + return self._store + + async def _vm(self, vm_id: str) -> None: + # Exact IDs prevent a session from following a reused VM name or IP. + await asyncio.to_thread(self.application.state_store.get, vm_id) + + async def list(self, vm_id: str) -> list[dict]: + await self._vm(vm_id) + return await asyncio.to_thread(self._history().list, vm_id) + + async def get(self, vm_id: str, session_id: str) -> dict: + await self._vm(vm_id) + return await asyncio.to_thread(self._history().get, vm_id, session_id) + + async def create(self, vm_id: str, rows: int, cols: int) -> dict: + validate_size(rows, cols) + await self._vm(vm_id) + history = self._history() + # A dropped HTTP request must not abandon a just-created guest connection. + task = asyncio.create_task(self._create(history, vm_id, rows, cols)) + self._pending.add(task) + + def completed(future): + self._pending.discard(future) + if not future.cancelled(): + future.exception() + + task.add_done_callback(completed) + return await asyncio.shield(task) + + async def _create(self, history: SessionStore, vm_id: str, rows: int, cols: int) -> dict: + async with self._create_lock: + if self._closing: + raise ConflictError("Nest is shutting down.") + if len(self._live) >= 32: + raise ConflictError("There are already 32 active terminal sessions. End a session before launching another.") + record = await asyncio.to_thread(history.create, vm_id, rows, cols) + terminal = None + 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) + self._live[record["id"]] = live + live.task = asyncio.create_task(self._read(history, record["id"], live)) + except BaseException as error: + if terminal is not None: + terminal.close() + message = str(error) if isinstance(error, UvmError) else "Could not launch the guest terminal." + await asyncio.to_thread(history.finish, record["id"], "failed", None, message) + raise + return await asyncio.to_thread(history.get, vm_id, record["id"]) + + async def _read(self, history: SessionStore, session_id: str, live: _LiveSession) -> None: + status, exit_code, error = "interrupted", None, None + try: + while True: + message = await live.terminal.read() + if message["type"] == "exit": + status, exit_code = "closed", message["exit_code"] + break + await asyncio.to_thread(history.append, session_id, message["data"]) + live.changed.set() + except asyncio.CancelledError: + error = "Nest stopped; this shell is no longer connected." if self._closing else "The guest terminal connection was interrupted." + except (UvmError, OSError) as exception: + error = str(exception) if isinstance(exception, UvmError) else "The guest terminal connection was interrupted." + finally: + live.terminal.close() + try: + await asyncio.to_thread(history.finish, session_id, status, exit_code, error) + finally: + live.changed.set() + self._live.pop(session_id, None) + + async def read(self, vm_id: str, session_id: str, offset: int) -> tuple[dict, int, bytes]: + return await asyncio.to_thread(self._history().read, vm_id, session_id, offset) + + async def wait(self, session_id: str) -> None: + live = self._live.get(session_id) + if live is None: + await asyncio.sleep(0.1) + return + try: + await asyncio.wait_for(live.changed.wait(), timeout=1) + except TimeoutError: + pass + 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: + raise ConflictError("This session has ended. Launch a new session to run commands.") + return live + + async def write(self, vm_id: str, session_id: str, data: bytes) -> None: + live = await self._active(vm_id, session_id) + try: + await live.terminal.write(data) + except OSError as error: + live.task.cancel() + raise GuestExecutionError("The terminal connection was interrupted.", outcome_unknown=True) from error + + async def resize(self, vm_id: str, session_id: str, rows: int, cols: int) -> None: + validate_size(rows, cols) + live = await self._active(vm_id, session_id) + try: + await live.terminal.resize(rows, cols) + await asyncio.to_thread(self._history().resize, session_id, rows, cols) + except OSError as error: + live.task.cancel() + raise GuestExecutionError("The terminal connection was interrupted.", outcome_unknown=True) from error + + async def end(self, vm_id: str, session_id: str) -> dict: + record = await self.get(vm_id, session_id) + live = self._live.get(session_id) + if live is not None: + try: + await live.terminal.end() + await asyncio.wait_for(asyncio.shield(live.task), timeout=8) + except (OSError, UvmError): + live.task.cancel() + await asyncio.gather(live.task, return_exceptions=True) + record = await self.get(vm_id, session_id) + return record + + async def shutdown(self) -> None: + self._closing = True + await asyncio.gather(*self._pending, return_exceptions=True) + tasks = [live.task for live in self._live.values()] + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + if self._store is not None: + self._store.close()