173 lines
7.9 KiB
Python
173 lines
7.9 KiB
Python
"""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)
|
|
try:
|
|
metadata = os.fstat(descriptor)
|
|
if not stat.S_ISREG(metadata.st_mode) or metadata.st_nlink != 1:
|
|
raise StateError("session history must use private regular files")
|
|
os.fchmod(descriptor, 0o600)
|
|
return descriptor
|
|
except BaseException:
|
|
os.close(descriptor)
|
|
raise
|
|
|
|
@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<?", (session_id, start)).fetchone()
|
|
if first is not None:
|
|
db.execute("UPDATE output SET offset=?, data=? WHERE session_id=? AND offset=?",
|
|
(start, first["data"][start - first["offset"]:], session_id, first["offset"]))
|
|
db.execute("UPDATE sessions SET output_start=?, output_end=?, truncated=?, updated_at=? WHERE id=?",
|
|
(start, end, int(start > 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
|