474 lines
21 KiB
Python
474 lines
21 KiB
Python
"""Standalone, stdlib-only guest execution agent (Python 3.10+).
|
||
|
||
This file is copied into guest images and must not import the host application.
|
||
The host also imports its wire-format helpers; importing it never starts a server.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import errno
|
||
import fcntl
|
||
import json
|
||
import math
|
||
import os
|
||
import pwd
|
||
import re
|
||
import selectors
|
||
import signal
|
||
import socket
|
||
import stat
|
||
import struct
|
||
import subprocess
|
||
import sys
|
||
import termios
|
||
import threading
|
||
import time
|
||
import uuid
|
||
from pathlib import Path
|
||
|
||
|
||
PROTOCOL_VERSION = 1
|
||
AGENT_PORT = 1070
|
||
DEFAULT_TIMEOUT_SECONDS = 60.0
|
||
MAX_TIMEOUT_SECONDS = 3600.0
|
||
MAX_COMMAND_BYTES = 65536
|
||
MAX_OUTPUT_BYTES = 1024 * 1024 # Combined stdout and stderr, before UTF-8 decoding.
|
||
MAX_REQUEST_BYTES = 6 * MAX_COMMAND_BYTES + 4096
|
||
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")
|
||
|
||
|
||
class ProtocolError(ValueError):
|
||
"""An invalid or oversized wire message; never include command text."""
|
||
|
||
|
||
class ExecutionFailure(Exception):
|
||
def __init__(self, message: str, *, outcome_unknown: bool = False) -> None:
|
||
super().__init__(message)
|
||
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")
|
||
try:
|
||
size = len(command.encode("utf-8"))
|
||
except UnicodeError as error:
|
||
raise ProtocolError("command must be valid UTF-8 text") from error
|
||
if size > MAX_COMMAND_BYTES:
|
||
raise ProtocolError(f"command exceeds {MAX_COMMAND_BYTES} UTF-8 bytes")
|
||
if type(sudo) is not bool:
|
||
raise ProtocolError("sudo must be a boolean")
|
||
try:
|
||
valid_timeout = (
|
||
type(timeout_seconds) in (int, float)
|
||
and math.isfinite(timeout_seconds)
|
||
and 0 < timeout_seconds <= MAX_TIMEOUT_SECONDS
|
||
)
|
||
except OverflowError:
|
||
valid_timeout = False
|
||
if not valid_timeout:
|
||
raise ProtocolError(f"timeout_seconds must be finite and greater than 0, up to {MAX_TIMEOUT_SECONDS:g}")
|
||
|
||
|
||
def _remaining(deadline: float) -> float:
|
||
remaining = deadline - time.monotonic()
|
||
if remaining <= 0:
|
||
raise TimeoutError("guest protocol deadline exceeded")
|
||
return remaining
|
||
|
||
|
||
def receive_exact(connection: socket.socket, size: int, deadline: float) -> bytes:
|
||
parts = bytearray()
|
||
while len(parts) < size:
|
||
connection.settimeout(_remaining(deadline))
|
||
chunk = connection.recv(size - len(parts))
|
||
if not chunk:
|
||
raise EOFError("guest connection closed")
|
||
parts.extend(chunk)
|
||
return bytes(parts)
|
||
|
||
|
||
def receive_frame(connection: socket.socket, limit: int, deadline: float) -> dict:
|
||
size, = struct.unpack("!I", receive_exact(connection, 4, deadline))
|
||
if not 0 < size <= limit:
|
||
raise ProtocolError("invalid guest message length")
|
||
try:
|
||
message = json.loads(receive_exact(connection, size, deadline))
|
||
except (ValueError, UnicodeError, RecursionError) as error:
|
||
raise ProtocolError("invalid guest JSON message") from error
|
||
if not isinstance(message, dict):
|
||
raise ProtocolError("guest message must be an object")
|
||
return message
|
||
|
||
|
||
def send_frame(connection: socket.socket, message: dict, limit: int, deadline: float) -> None:
|
||
payload = json.dumps(message, ensure_ascii=True, allow_nan=False, separators=(",", ":")).encode("ascii")
|
||
if len(payload) > limit:
|
||
raise ProtocolError("guest message exceeds its size limit")
|
||
connection.settimeout(_remaining(deadline))
|
||
connection.sendall(struct.pack("!I", len(payload)) + payload)
|
||
|
||
|
||
class JobScope:
|
||
"""A delegated guest cgroup, keeping detached descendants in the same job."""
|
||
|
||
def __init__(self) -> None:
|
||
entries = Path("/proc/self/cgroup").read_text(encoding="ascii").splitlines()
|
||
unified = next((entry[3:] for entry in entries if entry.startswith("0::")), None)
|
||
if unified is None or ".." in Path(unified).parts:
|
||
raise ExecutionFailure("guest execution requires a delegated cgroup v2 hierarchy")
|
||
parent = Path("/sys/fs/cgroup") / unified.lstrip("/")
|
||
self.path = parent / f"uvm-exec-{uuid.uuid4().hex}"
|
||
self.path.mkdir(mode=0o700)
|
||
if not (self.path / "cgroup.kill").is_file():
|
||
self.path.rmdir()
|
||
raise ExecutionFailure("guest execution requires cgroup.kill (Linux 5.14 or newer)")
|
||
|
||
def kill(self) -> None:
|
||
(self.path / "cgroup.kill").write_text("1", encoding="ascii")
|
||
|
||
def close(self) -> None:
|
||
# cgroup.events reports descendants as well as processes in this cgroup.
|
||
deadline = time.monotonic() + 1.0
|
||
while "populated 1" in (self.path / "cgroup.events").read_text(encoding="ascii"):
|
||
if time.monotonic() >= deadline:
|
||
raise ExecutionFailure("guest job cleanup could not be confirmed", outcome_unknown=True)
|
||
time.sleep(0.01)
|
||
for child in sorted(self.path.rglob("*"), key=lambda path: len(path.parts), reverse=True):
|
||
if child.is_dir():
|
||
child.rmdir()
|
||
self.path.rmdir()
|
||
|
||
|
||
def _signal_group(process: subprocess.Popen, number: int) -> None:
|
||
# Do not reap the group leader until after signaling: its PID pins the PGID.
|
||
try:
|
||
os.killpg(process.pid, number)
|
||
except ProcessLookupError:
|
||
pass
|
||
|
||
|
||
def _exited(process: subprocess.Popen) -> bool:
|
||
return os.waitid(os.P_PID, process.pid, os.WEXITED | os.WNOHANG | os.WNOWAIT) is not None
|
||
|
||
|
||
def _cleanup_job(process: subprocess.Popen | None, scope: JobScope) -> None:
|
||
failed = False
|
||
try:
|
||
scope.kill()
|
||
except (OSError, ExecutionFailure):
|
||
failed = True
|
||
if process is not None:
|
||
try:
|
||
_signal_group(process, signal.SIGKILL)
|
||
process.wait(timeout=1.0)
|
||
except (OSError, subprocess.TimeoutExpired):
|
||
failed = True
|
||
finally:
|
||
for name in ("stdout", "stderr"):
|
||
stream = getattr(process, name)
|
||
if stream is not None:
|
||
stream.close()
|
||
try:
|
||
scope.close()
|
||
except (OSError, ExecutionFailure):
|
||
failed = True
|
||
if failed:
|
||
raise ExecutionFailure("guest job cleanup could not be confirmed", outcome_unknown=process is not None)
|
||
|
||
|
||
def run_command(command: str, username: str, timeout_seconds: float, *, scope_factory=JobScope) -> dict:
|
||
"""Run one non-interactive guest job, drain both pipes, and bound its lifetime."""
|
||
scope = scope_factory()
|
||
process = None
|
||
buffers = {"stdout": bytearray(), "stderr": bytearray()}
|
||
captured = 0
|
||
truncated = False
|
||
timed_out = False
|
||
try:
|
||
# The fresh child joins its scope BEFORE dropping credentials or executing
|
||
# user text. Avoid preexec_fn in the threaded agent.
|
||
process = subprocess.Popen(
|
||
[sys.executable, "-I", str(Path(__file__).resolve()), "--child", str(scope.path), username, command],
|
||
stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
||
start_new_session=True, env={"PATH": "/usr/sbin:/usr/bin:/sbin:/bin", "LANG": "C.UTF-8"},
|
||
)
|
||
deadline = time.monotonic() + timeout_seconds
|
||
terminating = None
|
||
killed = False
|
||
with selectors.DefaultSelector() as selector:
|
||
for name in buffers:
|
||
stream = getattr(process, name)
|
||
os.set_blocking(stream.fileno(), False)
|
||
selector.register(stream, selectors.EVENT_READ, name)
|
||
while True:
|
||
now = time.monotonic()
|
||
exited = _exited(process)
|
||
if terminating is None and (exited or now >= deadline):
|
||
timed_out = not exited
|
||
terminating = now
|
||
_signal_group(process, signal.SIGTERM if timed_out else signal.SIGKILL)
|
||
if terminating is not None:
|
||
if not killed and (exited or now >= terminating + 0.5):
|
||
scope.kill()
|
||
_signal_group(process, signal.SIGKILL)
|
||
killed = True
|
||
if killed and exited and not selector.get_map():
|
||
break
|
||
if now >= terminating + 3.0:
|
||
raise ExecutionFailure("guest job cleanup could not be confirmed", outcome_unknown=True)
|
||
for key, _events in selector.select(0.05):
|
||
chunk = os.read(key.fileobj.fileno(), 65536)
|
||
if not chunk:
|
||
selector.unregister(key.fileobj)
|
||
continue
|
||
keep = min(len(chunk), MAX_OUTPUT_BYTES - captured)
|
||
buffers[key.data].extend(chunk[:keep])
|
||
captured += keep
|
||
truncated |= keep < len(chunk)
|
||
# Finally below kills the scope and reaps the leader before returning.
|
||
except OSError as error:
|
||
raise ExecutionFailure("guest command execution failed", outcome_unknown=process is not None) from error
|
||
finally:
|
||
_cleanup_job(process, scope)
|
||
return {
|
||
"exit_code": None if timed_out else process.returncode,
|
||
"stdout": buffers["stdout"].decode("utf-8", errors="replace"),
|
||
"stderr": buffers["stderr"].decode("utf-8", errors="replace"),
|
||
"timed_out": timed_out,
|
||
"truncated": truncated,
|
||
}
|
||
|
||
|
||
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)
|
||
os.setgid(account.pw_gid)
|
||
os.setuid(account.pw_uid)
|
||
elif (account.pw_uid, account.pw_gid) != (os.geteuid(), os.getegid()):
|
||
raise PermissionError("guest agent cannot assume the requested account")
|
||
os.chdir(account.pw_dir)
|
||
os.umask(0o022)
|
||
environment = {
|
||
"HOME": account.pw_dir, "USER": account.pw_name, "LOGNAME": account.pw_name,
|
||
"SHELL": "/bin/sh", "PATH": "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
|
||
"LANG": "C.UTF-8",
|
||
}
|
||
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 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.
|
||
for _ in range(16):
|
||
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:
|
||
raise ProtocolError("unsupported guest protocol version")
|
||
validate_command(request["command"], request["sudo"], request["timeout_seconds"])
|
||
username = "root" if request["sudo"] else config["username"]
|
||
result = runner(request["command"], username, request["timeout_seconds"])
|
||
send_frame(connection, result, MAX_RESPONSE_BYTES, time.monotonic() + CONNECT_TIMEOUT_SECONDS)
|
||
except (ProtocolError, ExecutionFailure, OSError, KeyError) as error:
|
||
# Never log request text, guest output, or credentials to the console.
|
||
try:
|
||
send_frame(connection, {
|
||
"error": "guest execution failed: " + str(error),
|
||
"outcome_unknown": isinstance(error, ExecutionFailure) and error.outcome_unknown,
|
||
}, 4096, time.monotonic() + CONNECT_TIMEOUT_SECONDS)
|
||
except (OSError, ProtocolError):
|
||
pass
|
||
except EOFError:
|
||
pass
|
||
|
||
|
||
def load_config(path: Path = CONFIG_PATH) -> dict:
|
||
descriptor = os.open(path, os.O_RDONLY | os.O_NOFOLLOW)
|
||
with os.fdopen(descriptor, encoding="utf-8") as source:
|
||
metadata = os.fstat(source.fileno())
|
||
if (not stat.S_ISREG(metadata.st_mode) or metadata.st_uid != 0
|
||
or metadata.st_mode & 0o077 or metadata.st_size > 4096):
|
||
raise ValueError("guest agent configuration must be a root-owned private regular file")
|
||
config = json.load(source)
|
||
if (not isinstance(config, dict) or set(config) != {"vm_id", "username"}
|
||
or not isinstance(config["vm_id"], str)
|
||
or re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,80}", config["vm_id"]) is None
|
||
or not isinstance(config["username"], str)
|
||
or re.fullmatch(r"[a-z_][a-z0-9_-]{0,31}\$?", config["username"]) is None):
|
||
raise ValueError("invalid guest agent configuration")
|
||
pwd.getpwnam(config["username"])
|
||
return config
|
||
|
||
|
||
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])
|
||
except (OSError, KeyError) as error:
|
||
print(f"uvm-agent: could not launch guest command: {error}", file=sys.stderr)
|
||
sys.exit(126)
|
||
return
|
||
if len(sys.argv) != 1 or os.geteuid() != 0:
|
||
raise SystemExit("uvm-agent must run as a root systemd service")
|
||
config = load_config()
|
||
slots = threading.BoundedSemaphore(MAX_CONCURRENT_COMMANDS)
|
||
|
||
def handle(connection: socket.socket) -> None:
|
||
try:
|
||
serve_connection(connection, config)
|
||
finally:
|
||
slots.release()
|
||
|
||
with socket.socket(socket.AF_VSOCK, socket.SOCK_STREAM) as listener:
|
||
listener.bind((socket.VMADDR_CID_ANY, AGENT_PORT))
|
||
listener.listen(MAX_CONCURRENT_COMMANDS)
|
||
while True:
|
||
connection, peer = listener.accept()
|
||
# Guest-local callers must not use the root agent to elevate themselves.
|
||
if peer[0] != socket.VMADDR_CID_HOST:
|
||
connection.close()
|
||
continue
|
||
if not slots.acquire(blocking=False):
|
||
connection.close()
|
||
continue
|
||
threading.Thread(target=handle, args=(connection,), daemon=True).start()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|