feat: Add guest execution agent and vsock support
- Introduced a standalone guest execution agent to facilitate command execution within the guest VM over Firecracker's vsock. - Enhanced FirecrackerClient to support vsock configuration, validating guest CID and UDS path. - Updated configuration to include vsock settings by default. - Refactored FirecrackerProcessManager to manage runtime sockets more effectively. - Implemented guest agent installation and provisioning in the image store. - Added API endpoints for executing commands in the VM and installing the guest agent. - Improved error handling for guest execution failures and agent availability.
This commit is contained in:
@@ -0,0 +1,342 @@
|
||||
"""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 json
|
||||
import math
|
||||
import os
|
||||
import pwd
|
||||
import re
|
||||
import selectors
|
||||
import signal
|
||||
import socket
|
||||
import stat
|
||||
import struct
|
||||
import subprocess
|
||||
import sys
|
||||
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
|
||||
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_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"):
|
||||
getattr(process, name).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:
|
||||
"""Enter the job cgroup and account in a fresh interpreter, then replace it."""
|
||||
(scope / "cgroup.procs").write_text(str(os.getpid()), encoding="ascii")
|
||||
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",
|
||||
}
|
||||
os.execve("/bin/sh", ["sh", "-c", command], environment)
|
||||
|
||||
|
||||
def serve_connection(connection: socket.socket, config: dict, *, runner=run_command) -> 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", "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) == 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()
|
||||
Reference in New Issue
Block a user