feat: add terminal support for guest execution agent and enhance lifecycle management

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])