from __future__ import annotations import asyncio import base64 import os import tempfile import unittest from pathlib import Path from types import SimpleNamespace from nest.server import create_app from nest.session_store import SessionStore from nest.sessions import SessionManager from tests.test_server import FASTAPI_AVAILABLE, fake_application from uvm.errors import ConflictError, GuestAgentUnavailableError, NotFoundError if FASTAPI_AVAILABLE: from fastapi.testclient import TestClient from starlette.websockets import WebSocketDisconnect class FakeTerminal: def __init__(self) -> None: self.events = asyncio.Queue() self.events.put_nowait({"type": "output", "data": b"guest$ "}) self.input = [] self.sizes = [] self.closed = False async def read(self): return await self.events.get() async def write(self, data): self.input.append(data) await self.events.put({"type": "output", "data": data + b"\r\nguest$ "}) async def resize(self, rows, cols): self.sizes.append((rows, cols)) async def end(self): await self.events.put({"type": "exit", "exit_code": 0}) def close(self): self.closed = True def session_application(base): application = fake_application(api_token="test-token") application.settings.base = Path(base) application.terminals = [] application.vm_ids = {"vm-test", "vm-other"} def get(vm_id): if vm_id not in application.vm_ids: raise NotFoundError("VM not found") return SimpleNamespace(id=vm_id) def open_terminal(vm_id, **_size): get(vm_id) terminal = FakeTerminal() application.terminals.append(terminal) return terminal, "worker" application.state_store = SimpleNamespace(get=get) application.lifecycle = SimpleNamespace(open_terminal=open_terminal) return application class SessionStoreTests(unittest.TestCase): def setUp(self): temporary = tempfile.TemporaryDirectory() self.addCleanup(temporary.cleanup) self.path = Path(temporary.name) / "nest" / "sessions.sqlite3" self.store = SessionStore(self.path, max_output=8) self.addCleanup(self.store.close) def test_history_is_scoped_by_exact_vm_id_and_survives_reopening(self): record = self.store.create("vm-a", 24, 80) self.store.activate(record["id"], "worker") self.store.append(record["id"], b"hello") self.store.finish(record["id"], "closed", 7) self.assertEqual(self.store.list("vm-b"), []) with self.assertRaises(NotFoundError): self.store.get("vm-b", record["id"]) self.store.close() reopened = SessionStore(self.path) self.addCleanup(reopened.close) saved, start, output = reopened.read("vm-a", record["id"], 0) self.assertEqual((saved["status"], saved["exit_code"], saved["username"], start, output), ("closed", 7, "worker", 0, b"hello")) self.assertEqual(self.path.stat().st_mode & 0o777, 0o600) self.assertEqual(self.path.parent.stat().st_mode & 0o777, 0o700) def test_restarted_server_marks_only_live_sessions_interrupted(self): live = self.store.create("vm-a", 24, 80) ended = self.store.create("vm-a", 24, 80) self.store.activate(live["id"], "worker") self.store.finish(ended["id"], "closed", 0) self.store.close() reopened = SessionStore(self.path) self.addCleanup(reopened.close) self.assertEqual(reopened.get("vm-a", live["id"])["status"], "interrupted") self.assertIsNotNone(reopened.get("vm-a", live["id"])["ended_at"]) self.assertEqual(reopened.get("vm-a", ended["id"])["status"], "closed") def test_transcript_is_bounded_and_offsets_support_exact_replay(self): record = self.store.create("vm-a", 24, 80) self.store.append(record["id"], b"abcdefghij") self.store.append(record["id"], b"kl") saved, start, output = self.store.read("vm-a", record["id"], 0) self.assertEqual((start, output, saved["output_end"], saved["truncated"]), (4, b"efghijkl", 12, True)) self.assertEqual(self.store.read("vm-a", record["id"], 9)[1:], (9, b"jkl")) self.assertEqual(self.store.read("vm-a", record["id"], 12)[2], b"") def test_second_server_cannot_reconcile_another_servers_live_sessions(self): record = self.store.create("vm-a", 24, 80) self.store.activate(record["id"], "worker") with self.assertRaises(ConflictError): SessionStore(self.path) self.assertEqual(self.store.get("vm-a", record["id"])["status"], "active") def test_symlinked_history_file_is_not_opened_or_rewritten(self): self.store.close() self.path.unlink() target = self.path.parent / "unrelated" target.write_bytes(b"keep this") self.path.symlink_to(target) with self.assertRaises(OSError): SessionStore(self.path) self.assertEqual(target.read_bytes(), b"keep this") class SessionManagerTests(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.temporary = tempfile.TemporaryDirectory() self.addCleanup(self.temporary.cleanup) self.application = session_application(self.temporary.name) self.manager = SessionManager(self.application) self.addAsyncCleanup(self.manager.shutdown) async def test_failed_launch_is_saved_in_history(self): def unavailable(*_args, **_kwargs): raise GuestAgentUnavailableError("update the guest agent") self.application.lifecycle.open_terminal = unavailable with self.assertRaises(GuestAgentUnavailableError): await self.manager.create("vm-test", 24, 80) history = await self.manager.list("vm-test") self.assertEqual(history[0]["status"], "failed") self.assertEqual(history[0]["error"], "update the guest agent") async def test_shutdown_closes_live_channels_and_saves_interrupted_history(self): record = await self.manager.create("vm-test", 24, 80) await self.manager.shutdown() self.assertTrue(self.application.terminals[0].closed) store = SessionStore(Path(self.temporary.name) / "nest" / "sessions.sqlite3") self.addCleanup(store.close) self.assertEqual(store.get("vm-test", record["id"])["status"], "interrupted") async def test_cancelling_the_create_request_does_not_orphan_the_shell(self): import threading entered, release = threading.Event(), threading.Event() connect = self.application.lifecycle.open_terminal def slow_connect(*args, **kwargs): entered.set() release.wait(2) return connect(*args, **kwargs) self.application.lifecycle.open_terminal = slow_connect request = asyncio.create_task(self.manager.create("vm-test", 24, 80)) await asyncio.to_thread(entered.wait, 2) request.cancel() release.set() await asyncio.gather(request, return_exceptions=True) await self.manager.shutdown() self.assertTrue(self.application.terminals[0].closed) @unittest.skipUnless(FASTAPI_AVAILABLE, "FastAPI and httpx are required") class SessionApiTests(unittest.TestCase): def setUp(self): self.temporary = self.enterContext(tempfile.TemporaryDirectory()) self.application = session_application(self.temporary) self.api = create_app(self.application, host="127.0.0.1") self.client = self.enterContext(TestClient(self.api)) self.headers = {"X-UVM-Token": "test-token"} self.base = "/vms/vm-test/sessions" def create(self): response = self.client.post(self.base, headers=self.headers, json={}) self.assertEqual(response.status_code, 201, response.text) return response.json() def connect(self, record, *, vm_id="vm-test", offset=0): ws = self.enterContext(self.client.websocket_connect(f"/vms/{vm_id}/sessions/{record['id']}/terminal")) ws.send_json({"type": "auth", "token": "test-token", "offset": offset}) return ws def output(self, ws): while True: message = ws.receive_json() if message["type"] == "output": data = base64.b64decode(message["data"]) end = message["offset"] + len(data) ws.send_json({"type": "ack", "offset": end}) return data, end self.assertEqual(message["type"], "session", message) def test_all_http_operations_and_websocket_output_require_authentication(self): for method, path in (("GET", self.base), ("POST", self.base), ("GET", self.base + "/session-unknown"), ("POST", self.base + "/session-unknown/close")): self.assertEqual(self.client.request(method, path, json={}).status_code, 401) for token in ("wrong", "", "\u00e9"): with self.client.websocket_connect(self.base + "/session-unknown/terminal") as ws: ws.send_json({"type": "auth", "token": token}) with self.assertRaises(WebSocketDisconnect) as raised: ws.receive_json() self.assertEqual(raised.exception.code, 4401) self.assertEqual(self.application.terminals, []) self.assertFalse((Path(self.temporary) / "nest").exists()) def test_foreign_browser_origin_is_rejected(self): with self.assertRaises(WebSocketDisconnect): with self.client.websocket_connect(self.base + "/unknown/terminal", headers={"Origin": "https://foreign.example"}): self.fail("foreign origin accepted") def test_create_rejects_unknown_vm_aliases_and_invalid_sizes(self): for body in ({"rows": 0}, {"cols": True}, {"cols": 401}, {"command": "id"}): self.assertEqual(self.client.post(self.base, headers=self.headers, json=body).status_code, 422) self.assertEqual(self.client.post("/vms/name-alias/sessions", headers=self.headers, json={}).status_code, 404) self.assertEqual(self.application.terminals, []) def test_sessions_are_nested_under_the_selected_vm(self): record = self.create() self.assertEqual(record["username"], "worker") other = "/vms/vm-other/sessions" self.assertEqual(self.client.get(other, headers=self.headers).json(), []) self.assertEqual(self.client.get(other + "/" + record["id"], headers=self.headers).status_code, 404) self.assertEqual(self.client.post(other + "/" + record["id"] + "/close", headers=self.headers).status_code, 404) ws = self.connect(record, vm_id="vm-other") self.assertEqual(ws.receive_json()["type"], "error") self.assertFalse(self.application.terminals[0].closed) def test_disconnect_reconnect_input_resize_and_closed_history(self): record = self.create() ws = self.connect(record) output, offset = self.output(ws) self.assertEqual(output, b"guest$ ") ws.send_json({"type": "resize", "rows": 35, "cols": 120}) ws.send_json({"type": "input", "data": base64.b64encode(b"pwd\r").decode()}) output, offset = self.output(ws) self.assertEqual(output, b"pwd\r\r\nguest$ ") self.assertEqual(self.application.terminals[0].sizes, [(35, 120)]) ws.close() self.assertFalse(self.application.terminals[0].closed) reopened = self.connect(record, offset=offset) self.assertEqual(reopened.receive_json()["session"]["status"], "active") self.assertEqual(len(self.application.terminals), 1) closed = self.client.post(self.base + "/" + record["id"] + "/close", headers=self.headers) self.assertEqual(closed.json()["status"], "closed") self.assertEqual(closed.json()["exit_code"], 0) replay = self.connect(record) transcript, _end = self.output(replay) self.assertEqual(transcript, b"guest$ pwd\r\r\nguest$ ") self.assertEqual(self.client.get(self.base, headers=self.headers).json()[0]["status"], "closed") self.assertTrue(self.application.terminals[0].closed)