From 3b94ac1b48470da76e895bc6dbf356eb6190c892 Mon Sep 17 00:00:00 2001 From: kstyagi Date: Fri, 25 Sep 2026 19:51:35 +0530 Subject: [PATCH] feat: add terminal support for guest execution agent and enhance lifecycle management --- pyproject.toml | 1 + uvm/guest_agent.py | 139 ++++++++++++++++++++++++++++++++++++++++-- uvm/guest_terminal.py | 89 +++++++++++++++++++++++++++ uvm/images.py | 4 +- uvm/install.py | 2 +- uvm/lifecycle.py | 21 +++++++ 6 files changed, 250 insertions(+), 6 deletions(-) create mode 100644 uvm/guest_terminal.py diff --git a/pyproject.toml b/pyproject.toml index 2611ebd..aa7b709 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,7 @@ dependencies = [ "fastapi>=0.115,<1", "pydantic>=2,<3", "uvicorn>=0.30,<1", + "websockets>=13,<17", ] [project.optional-dependencies] diff --git a/uvm/guest_agent.py b/uvm/guest_agent.py index 0611e4f..19b2911 100644 --- a/uvm/guest_agent.py +++ b/uvm/guest_agent.py @@ -6,6 +6,9 @@ The host also imports its wire-format helpers; importing it never starts a serve from __future__ import annotations +import base64 +import errno +import fcntl import json import math import os @@ -18,6 +21,7 @@ import stat import struct import subprocess import sys +import termios import threading import time import uuid @@ -35,6 +39,8 @@ MAX_RESPONSE_BYTES = 6 * MAX_OUTPUT_BYTES + 4096 CONNECT_TIMEOUT_SECONDS = 5.0 CLEANUP_SECONDS = 5.0 MAX_CONCURRENT_COMMANDS = 8 +MAX_TERMINAL_BYTES = 8192 +MAX_TERMINAL_FRAME = 16384 CONFIG_PATH = Path("/etc/uvm-agent.json") @@ -48,6 +54,23 @@ class ExecutionFailure(Exception): self.outcome_unknown = outcome_unknown +def validate_terminal_size(rows: object, cols: object) -> None: + if type(rows) is not int or type(cols) is not int or not (2 <= rows <= 200 and 2 <= cols <= 400): + raise ProtocolError("terminal dimensions must be integers: 2–200 rows and 2–400 columns") + + +def terminal_data(value: object) -> bytes: + if not isinstance(value, str) or len(value) > 4 * ((MAX_TERMINAL_BYTES + 2) // 3): + raise ProtocolError("terminal data exceeds its size limit") + try: + data = base64.b64decode(value, validate=True) + except (ValueError, UnicodeError) as error: + raise ProtocolError("invalid terminal data") from error + if len(data) > MAX_TERMINAL_BYTES: + raise ProtocolError("terminal data exceeds its size limit") + return data + + def validate_command(command: object, sudo: object, timeout_seconds: object) -> None: if not isinstance(command, str) or not command.strip() or "\0" in command: raise ProtocolError("command must be non-empty text without NUL characters") @@ -167,7 +190,9 @@ def _cleanup_job(process: subprocess.Popen | None, scope: JobScope) -> None: failed = True finally: for name in ("stdout", "stderr"): - getattr(process, name).close() + stream = getattr(process, name) + if stream is not None: + stream.close() try: scope.close() except (OSError, ExecutionFailure): @@ -239,9 +264,12 @@ def run_command(command: str, username: str, timeout_seconds: float, *, scope_fa } -def child_main(scope: Path, username: str, command: str) -> None: +def child_main(scope: Path, username: str, command: str | None) -> None: """Enter the job cgroup and account in a fresh interpreter, then replace it.""" (scope / "cgroup.procs").write_text(str(os.getpid()), encoding="ascii") + if command is None: + # Popen's start_new_session already performed setsid in this fresh child. + fcntl.ioctl(0, termios.TIOCSCTTY, 0) account = pwd.getpwnam(username) if os.geteuid() == 0: os.initgroups(account.pw_name, account.pw_gid) @@ -256,15 +284,111 @@ def child_main(scope: Path, username: str, command: str) -> None: "SHELL": "/bin/sh", "PATH": "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin", "LANG": "C.UTF-8", } - os.execve("/bin/sh", ["sh", "-c", command], environment) + if command is None: + shell = account.pw_shell or "/bin/sh" + environment.update(SHELL=shell, TERM="xterm-256color", COLORTERM="truecolor") + os.execve(shell, [shell, "-i"], environment) + else: + os.execve("/bin/sh", ["sh", "-c", command], environment) -def serve_connection(connection: socket.socket, config: dict, *, runner=run_command) -> None: +def run_terminal(connection: socket.socket, username: str, rows: int, cols: int, *, scope_factory=JobScope) -> None: + """Bridge one real guest PTY until shell exit, explicit close, or host loss.""" + validate_terminal_size(rows, cols) + scope = scope_factory() + process = None + master = slave = None + try: + master, slave = os.openpty() + fcntl.ioctl(slave, termios.TIOCSWINSZ, struct.pack("HHHH", rows, cols, 0, 0)) + process = subprocess.Popen( + [sys.executable, "-I", str(Path(__file__).resolve()), "--terminal-child", str(scope.path), username], + stdin=slave, stdout=slave, stderr=slave, start_new_session=True, + env={"PATH": "/usr/sbin:/usr/bin:/sbin:/bin", "LANG": "C.UTF-8"}, + ) + os.close(slave) + slave = None + os.set_blocking(master, False) + send_frame(connection, {"type": "ready"}, MAX_TERMINAL_FRAME, time.monotonic() + CONNECT_TIMEOUT_SECONDS) + pending = bytearray() + with selectors.DefaultSelector() as selector: + selector.register(connection, selectors.EVENT_READ, "host") + selector.register(master, selectors.EVENT_READ, "pty") + finished = False + while not finished: + for key, events in selector.select(0.1): + if key.data == "host": + message = receive_frame(connection, MAX_TERMINAL_FRAME, time.monotonic() + CONNECT_TIMEOUT_SECONDS) + if message == {"type": "close"}: + finished = True + break + if set(message) == {"type", "data"} and message["type"] == "input": + pending.extend(terminal_data(message["data"])) + if len(pending) > 65536: + raise ProtocolError("terminal input backlog exceeded") + elif set(message) == {"type", "rows", "cols"} and message["type"] == "resize": + validate_terminal_size(message["rows"], message["cols"]) + fcntl.ioctl(master, termios.TIOCSWINSZ, struct.pack("HHHH", message["rows"], message["cols"], 0, 0)) + else: + raise ProtocolError("invalid terminal control message") + selector.modify(master, selectors.EVENT_READ | (selectors.EVENT_WRITE if pending else 0), "pty") + else: + if events & selectors.EVENT_READ: + try: + data = os.read(master, MAX_TERMINAL_BYTES) + except OSError as error: + if error.errno != errno.EIO: + raise + data = b"" + if not data: + finished = True + break + send_frame(connection, {"type": "output", "data": base64.b64encode(data).decode("ascii")}, + MAX_TERMINAL_FRAME, time.monotonic() + CONNECT_TIMEOUT_SECONDS) + if events & selectors.EVENT_WRITE and pending: + try: + del pending[:os.write(master, pending)] + except BlockingIOError: + pass + 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: + try: + data = os.read(master, MAX_TERMINAL_BYTES) + except (BlockingIOError, OSError): + break + if not data: + break + send_frame(connection, {"type": "output", "data": base64.b64encode(data).decode("ascii")}, + MAX_TERMINAL_FRAME, time.monotonic() + CONNECT_TIMEOUT_SECONDS) + break + finally: + try: + _cleanup_job(process, scope) + finally: + for descriptor in (master, slave): + if descriptor is not None: + os.close(descriptor) + send_frame(connection, {"type": "exit", "exit_code": process.returncode}, MAX_TERMINAL_FRAME, + time.monotonic() + CONNECT_TIMEOUT_SECONDS) + + +def serve_connection(connection: socket.socket, config: dict, *, runner=run_command, terminal_runner=run_terminal) -> None: with connection: try: deadline = time.monotonic() + CONNECT_TIMEOUT_SECONDS send_frame(connection, {"version": PROTOCOL_VERSION, **config}, 4096, deadline) request = receive_frame(connection, MAX_REQUEST_BYTES, deadline) + if set(request) == {"version", "terminal"}: + if type(request["version"]) is not int or request["version"] != PROTOCOL_VERSION: + raise ProtocolError("unsupported guest protocol version") + size = request["terminal"] + if not isinstance(size, dict) or set(size) != {"rows", "cols"}: + raise ProtocolError("invalid terminal request") + validate_terminal_size(size["rows"], size["cols"]) + terminal_runner(connection, config["username"], size["rows"], size["cols"]) + return if set(request) != {"version", "command", "sudo", "timeout_seconds"}: raise ProtocolError("invalid execution request fields") if type(request["version"]) is not int or request["version"] != PROTOCOL_VERSION: @@ -305,6 +429,13 @@ def load_config(path: Path = CONFIG_PATH) -> dict: def main() -> None: + if len(sys.argv) == 4 and sys.argv[1] == "--terminal-child": + try: + child_main(Path(sys.argv[2]), sys.argv[3], None) + except (OSError, KeyError) as error: + print(f"uvm-agent: could not launch guest terminal: {error}", file=sys.stderr) + sys.exit(126) + return if len(sys.argv) == 5 and sys.argv[1] == "--child": try: child_main(Path(sys.argv[2]), sys.argv[3], sys.argv[4]) diff --git a/uvm/guest_terminal.py b/uvm/guest_terminal.py new file mode 100644 index 0000000..bfb9822 --- /dev/null +++ b/uvm/guest_terminal.py @@ -0,0 +1,89 @@ +"""Interactive, process-pinned guest terminal transport over the existing vsock agent.""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import struct +import time + +from .errors import GuestAgentUnavailableError, GuestExecutionError, ValidationError +from .guest_agent import ( + CONNECT_TIMEOUT_SECONDS, MAX_TERMINAL_FRAME, PROTOCOL_VERSION, ProtocolError, + receive_frame, send_frame, terminal_data, validate_terminal_size, +) +from .guest_exec import GuestSession + + +def validate_size(rows: object, cols: object) -> None: + try: + validate_terminal_size(rows, cols) + except ProtocolError as error: + raise ValidationError(str(error)) from error + + +class GuestTerminal(GuestSession): + def start(self, rows: int, cols: int) -> None: + validate_size(rows, cols) + try: + deadline = time.monotonic() + CONNECT_TIMEOUT_SECONDS + send_frame(self._connection, {"version": PROTOCOL_VERSION, "terminal": {"rows": rows, "cols": cols}}, + MAX_TERMINAL_FRAME, deadline) + response = receive_frame(self._connection, MAX_TERMINAL_FRAME, deadline) + if response != {"type": "ready"}: + raise ProtocolError("terminal support is unavailable") + self._connection.setblocking(False) + self._write_lock = asyncio.Lock() + except (OSError, EOFError, ProtocolError) as error: + self.close() + raise GuestAgentUnavailableError( + "Interactive terminal support is unavailable. Stop this VM, use Install / update terminal support " + "in its Sessions view, then start it and launch a new session." + ) from error + + async def _receive_exact(self, size: int) -> bytes: + data = bytearray() + while len(data) < size: + chunk = await asyncio.get_running_loop().sock_recv(self._connection, size - len(data)) + if not chunk: + raise EOFError("guest terminal connection closed") + data.extend(chunk) + return bytes(data) + + async def read(self) -> dict: + try: + # An idle shell has no output deadline; partial messages are bounded. + first = await self._receive_exact(1) + async with asyncio.timeout(CONNECT_TIMEOUT_SECONDS): + size, = struct.unpack("!I", first + await self._receive_exact(3)) + if not 0 < size <= MAX_TERMINAL_FRAME: + raise ProtocolError("invalid terminal frame length") + message = json.loads(await self._receive_exact(size)) + if isinstance(message, dict): + if set(message) == {"type", "data"} and message["type"] == "output": + return {"type": "output", "data": terminal_data(message["data"])} + if (set(message) == {"type", "exit_code"} and message["type"] == "exit" + and type(message["exit_code"]) is int and -64 <= message["exit_code"] <= 255): + return message + raise ProtocolError("invalid guest terminal response") + except (OSError, EOFError, ValueError, UnicodeError, RecursionError) as error: + raise GuestExecutionError("The guest terminal connection ended unexpectedly.", outcome_unknown=True) from error + + async def _send(self, message: dict) -> None: + payload = json.dumps(message, separators=(",", ":")).encode("ascii") + if len(payload) > MAX_TERMINAL_FRAME: + raise ValidationError("terminal input exceeds its size limit") + async with self._write_lock: + async with asyncio.timeout(CONNECT_TIMEOUT_SECONDS): + await asyncio.get_running_loop().sock_sendall(self._connection, struct.pack("!I", len(payload)) + payload) + + async def write(self, data: bytes) -> None: + await self._send({"type": "input", "data": base64.b64encode(data).decode("ascii")}) + + async def resize(self, rows: int, cols: int) -> None: + validate_size(rows, cols) + await self._send({"type": "resize", "rows": rows, "cols": cols}) + + async def end(self) -> None: + await self._send({"type": "close"}) diff --git a/uvm/images.py b/uvm/images.py index 6abe778..6eee8bf 100644 --- a/uvm/images.py +++ b/uvm/images.py @@ -317,7 +317,9 @@ class ImageStore: if not install: marker = self._read_guest_file(disk, "/etc/uvm-image", files / "marker", required=False, regular_only=True) - if marker is None or marker.path.read_text(encoding="utf-8") != "uvm-ubuntu-jammy-v4\n": + if marker is None or marker.path.read_text(encoding="utf-8") not in { + "uvm-ubuntu-jammy-v4\n", "uvm-ubuntu-jammy-v5\n", + }: return False descriptor, candidate = tempfile.mkstemp(prefix=".rootfs-agent-", suffix=".ext4", dir=disk.parent) os.close(descriptor) diff --git a/uvm/install.py b/uvm/install.py index b662893..d2a28d0 100644 --- a/uvm/install.py +++ b/uvm/install.py @@ -21,7 +21,7 @@ from .system import CommandRunner, check_kvm, ensure_data_directories, require_r _UBUNTU_RELEASE = "jammy" -_ROOTFS_RECIPE = "uvm-ubuntu-jammy-v4\n" +_ROOTFS_RECIPE = "uvm-ubuntu-jammy-v5\n" _GUEST_PACKAGES = ( "ca-certificates", "openssh-server", "sudo", "systemd-sysv", "udev", "dbus", "libpam-systemd", "iproute2", "iputils-ping", "curl", "python3", diff --git a/uvm/lifecycle.py b/uvm/lifecycle.py index 236521b..f507ba6 100644 --- a/uvm/lifecycle.py +++ b/uvm/lifecycle.py @@ -19,6 +19,7 @@ from .firecracker.config import build_config, write_config from .firecracker.process import FirecrackerProcessManager, ProcessInfo from .guest_agent import DEFAULT_TIMEOUT_SECONDS from .guest_exec import ExecResult, GuestSession, validate_exec +from .guest_terminal import GuestTerminal, validate_size from .images import ImageStore from .integrity import load_manifest, verify_file from .known_hosts import KnownHostsManager @@ -61,6 +62,7 @@ class LifecycleService: known_hosts: KnownHostsManager, client_factory: Callable[[Path, float], FirecrackerClient] = FirecrackerClient, guest_connector: Callable[[Path, str, str, int], GuestSession] = GuestSession.connect, + terminal_connector: Callable[[Path, str, str, int], GuestTerminal] = GuestTerminal.connect, clock: Callable[[], float] = time.time, ) -> None: self._settings = settings @@ -71,6 +73,7 @@ class LifecycleService: self._known_hosts = known_hosts self._client_factory = client_factory self._guest_connector = guest_connector + self._terminal_connector = terminal_connector self._clock = clock def create(self, spec: VmSpec) -> VmRecord: @@ -520,6 +523,24 @@ class LifecycleService: finally: session.close() + def open_terminal(self, vm_id: str, *, rows: int = 24, cols: int = 80) -> tuple[GuestTerminal, str]: + """Connect to the exact running VMM under the lock, then open its guest PTY.""" + require_root() + validate_size(rows, cols) + with self._state_store.operation_lock(): + vm = self._load_for_operation(vm_id, exact_id=True) + self._not_deleting(vm) + if (vm.status != "running" or vm.operation is not None or vm.pending_resize is not None + or self._process.observe(vm) != "alive"): + raise ConflictError("terminal sessions require a running VM with no pending lifecycle operation") + terminal = self._terminal_connector(self._settings.guest_socket_path(vm.id), vm.id, vm.username, vm.pid) + try: + terminal.start(rows, cols) + return terminal, vm.username + except BaseException: + terminal.close() + raise + def install_guest_agent(self, vm_id: str) -> VmRecord: """Explicitly install/upgrade the agent on a cleanly stopped private disk.""" require_root()