from __future__ import annotations import os import pwd import shlex import socket import struct import sys import tempfile import threading import time import unittest from contextlib import contextmanager from pathlib import Path from unittest.mock import Mock, patch from uvm.errors import GuestAgentUnavailableError, GuestExecutionError, ValidationError from uvm import guest_agent as agent from uvm.guest_exec import GuestSession, validate_exec class FakeScope: """Use real subprocesses/process groups without touching the host's cgroups.""" def __init__(self) -> None: self.temporary = tempfile.TemporaryDirectory() self.path = Path(self.temporary.name) (self.path / "cgroup.procs").touch() self.kills = 0 self.joined_pid = None def kill(self) -> None: self.kills += 1 self.joined_pid = (self.path / "cgroup.procs").read_text() def close(self) -> None: self.temporary.cleanup() class GuestCommandTests(unittest.TestCase): def setUp(self) -> None: self.account = pwd.getpwuid(os.geteuid()) self.scope = FakeScope() self.addCleanup(self.scope.temporary.cleanup) def run_command(self, command, timeout=5): return agent.run_command(command, self.account.pw_name, timeout, scope_factory=lambda: self.scope) def test_multiline_script_has_separate_output_exit_status_and_no_stdin(self) -> None: result = self.run_command("printf 'hello\\n'\nprintf 'problem\\n' >&2\nread input || exit 7") self.assertEqual(result, { "exit_code": 7, "stdout": "hello\n", "stderr": "problem\n", "timed_out": False, "truncated": False, }) self.assertTrue(self.scope.joined_pid.isdigit()) self.assertGreater(self.scope.kills, 0) def test_command_enters_scope_before_running_and_uses_account_home_and_environment(self) -> None: with patch.dict(os.environ, {"UVM_API_TOKEN": "must-not-reach-guest-command"}): result = self.run_command('printf "%s\\n" "$$" "$HOME" "$USER" "$UVM_API_TOKEN"; pwd; id -u') self.assertEqual(result["exit_code"], 0) self.assertEqual(result["stdout"].splitlines(), [ self.scope.joined_pid, self.account.pw_dir, self.account.pw_name, "", self.account.pw_dir, str(os.geteuid()), ]) def test_timeout_kills_the_shell_and_children_that_ignore_sigterm(self) -> None: start = time.monotonic() result = self.run_command("trap '' TERM; sleep 30 & printf '%s\\n' \"$!\"; wait", timeout=0.25) self.assertTrue(result["timed_out"]) self.assertIsNone(result["exit_code"]) self.assertLess(time.monotonic() - start, 4) child = int(result["stdout"]) try: state = Path(f"/proc/{child}/stat").read_text().split(") ", 1)[1].split()[0] except FileNotFoundError: state = "gone" self.assertIn(state, {"gone", "Z"}) def test_large_stdout_and_stderr_are_drained_without_unbounded_capture(self) -> None: script = "import os; [(os.write(1, b'o'*65536), os.write(2, b'e'*65536)) for _ in range(32)]" result = self.run_command(f"{shlex.quote(sys.executable)} -c {shlex.quote(script)}") self.assertEqual(result["exit_code"], 0) self.assertEqual(len(result["stdout"]) + len(result["stderr"]), agent.MAX_OUTPUT_BYTES) self.assertTrue(result["truncated"]) self.assertFalse(result["timed_out"]) def test_non_utf8_output_is_replaced_without_losing_the_exit_code(self) -> None: result = self.run_command("printf '\\377\\000'; exit 255") self.assertEqual(result["stdout"], "\ufffd\0") self.assertEqual(result["exit_code"], 255) def test_cgroup_cleanup_failure_is_an_unknown_outcome_and_still_reaps_the_process(self) -> None: with patch.object(self.scope, "kill", side_effect=OSError("cannot kill scope")): with self.assertRaises(agent.ExecutionFailure) as raised: self.run_command("exit 0") self.assertTrue(raised.exception.outcome_unknown) def test_privilege_drop_initializes_groups_before_uid_and_exec(self) -> None: account = Mock(pw_name="worker", pw_uid=1001, pw_gid=1002, pw_dir="/home/worker") events = [] with ( patch.object(agent.pwd, "getpwnam", return_value=account), patch.object(agent.os, "geteuid", return_value=0), patch.object(agent.os, "initgroups", side_effect=lambda *args: events.append(("groups", args))), patch.object(agent.os, "setgid", side_effect=lambda value: events.append(("gid", value))), patch.object(agent.os, "setuid", side_effect=lambda value: events.append(("uid", value))), patch.object(agent.os, "chdir") as chdir, patch.object(agent.os, "umask"), patch.object(agent.os, "execve") as execute, ): agent.child_main(self.scope.path, "worker", "printf x > file") self.assertEqual(events, [("groups", ("worker", 1002)), ("gid", 1002), ("uid", 1001)]) chdir.assert_called_once_with("/home/worker") self.assertEqual(execute.call_args.args[:2], ("/bin/sh", ["sh", "-c", "printf x > file"])) self.assertEqual(execute.call_args.args[2]["HOME"], "/home/worker") class GuestProtocolTests(unittest.TestCase): def test_validates_byte_limits_types_finite_deadlines_and_text_before_transport(self) -> None: for command, sudo, timeout in ( ("", False, 1), (" \n", False, 1), ("a\0b", False, 1), ("\ud800", False, 1), ("x" * (agent.MAX_COMMAND_BYTES + 1), False, 1), ("\u00e9" * agent.MAX_COMMAND_BYTES, False, 1), ("id", "false", 1), ("id", 0, 1), ("id", False, True), ("id", False, 0), ("id", False, -1), ("id", False, float("nan")), ("id", False, float("inf")), ("id", False, 3601), ("id", False, "60"), ): with self.subTest(command=command[:20], sudo=sudo, timeout=timeout), self.assertRaises(ValidationError): validate_exec(command, sudo, timeout) validate_exec("id", False, 0.1) def test_oversized_frame_is_rejected_before_reading_a_body(self) -> None: left, right = socket.socketpair() with left, right: left.sendall(struct.pack("!I", agent.MAX_REQUEST_BYTES + 1)) with self.assertRaises(agent.ProtocolError): agent.receive_frame(right, agent.MAX_REQUEST_BYTES, time.monotonic() + 1) def test_read_deadline_applies_to_the_whole_frame(self) -> None: left, right = socket.socketpair() with left, right: left.sendall(struct.pack("!I", 5) + b"{") start = time.monotonic() with self.assertRaises(TimeoutError): agent.receive_frame(right, 1024, time.monotonic() + 0.05) self.assertLess(time.monotonic() - start, 0.5) def test_agent_maps_sudo_to_root_and_rejects_bad_requests_without_launching(self) -> None: config = {"vm_id": "vm-test", "username": "worker"} for sudo in (False, True, "true"): with self.subTest(sudo=sudo): left, right = socket.socketpair() result = {"exit_code": 0, "stdout": "", "stderr": "", "timed_out": False, "truncated": False} runner = Mock(return_value=result) thread = threading.Thread(target=agent.serve_connection, args=(right, config), kwargs={"runner": runner}) thread.start() with left: deadline = time.monotonic() + 2 self.assertEqual(agent.receive_frame(left, 4096, deadline), {"version": 1, **config}) agent.send_frame(left, {"version": 1, "command": "id", "sudo": sudo, "timeout_seconds": 2}, 4096, deadline) reply = agent.receive_frame(left, 4096, deadline) thread.join(2) self.assertFalse(thread.is_alive()) if type(sudo) is bool: self.assertEqual(reply, result) runner.assert_called_once_with("id", "root" if sudo else "worker", 2) else: self.assertFalse(reply["outcome_unknown"]) runner.assert_not_called() def test_guest_local_connections_are_rejected_before_a_worker_is_started(self) -> None: connection = Mock() listener = Mock() listener.__enter__ = Mock(return_value=listener) listener.__exit__ = Mock(return_value=False) listener.accept.side_effect = [(connection, (3, 9999)), KeyboardInterrupt()] with ( patch.object(agent.sys, "argv", ["agent.py"]), patch.object(agent.os, "geteuid", return_value=0), patch.object(agent, "load_config", return_value={"vm_id": "vm-test", "username": "root"}), patch.object(agent.socket, "socket", return_value=listener), patch.object(agent.threading, "Thread") as worker, self.assertRaises(KeyboardInterrupt), ): agent.main() connection.close.assert_called_once() worker.assert_not_called() class GuestTransportTests(unittest.TestCase): @contextmanager def firecracker(self, handler): with tempfile.TemporaryDirectory() as temporary: path = Path(temporary) / "vsock.sock" errors = [] with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as listener: listener.bind(str(path)) listener.listen(1) listener.settimeout(3) def serve(): try: connection, _peer = listener.accept() with connection: handler(connection) except (BrokenPipeError, EOFError, ConnectionResetError): pass except BaseException as error: errors.append(error) thread = threading.Thread(target=serve, daemon=True) thread.start() try: yield path finally: thread.join(4) self.assertFalse(thread.is_alive()) if errors: raise errors[0] def handshake(self, connection, **changes): deadline = time.monotonic() + 2 self.assertEqual(agent.receive_exact(connection, len(b"CONNECT 1070\n"), deadline), b"CONNECT 1070\n") # The acknowledgement and JSON frame can be coalesced on the wire. connection.sendall(b"OK 1234\n") agent.send_frame(connection, {"version": 1, "vm_id": "vm-test", "username": "worker", **changes}, 4096, deadline) def test_round_trip_preserves_command_text_and_a_nonzero_exit(self) -> None: command = "printf '%s\\n' 'quotes; $(literal)'\nexit 7" expected = {"exit_code": 7, "stdout": "quotes; $(literal)\n", "stderr": "", "timed_out": False, "truncated": False} def handle(connection): self.handshake(connection) request = agent.receive_frame(connection, agent.MAX_REQUEST_BYTES, time.monotonic() + 2) self.assertEqual(request, {"version": 1, "command": command, "sudo": True, "timeout_seconds": 3}) agent.send_frame(connection, expected, agent.MAX_RESPONSE_BYTES, time.monotonic() + 2) with self.firecracker(handle) as path: session = GuestSession.connect(path, "vm-test", "worker", os.getpid()) try: result = session.execute(command, sudo=True, timeout_seconds=3) finally: session.close() self.assertEqual(result.exit_code, 7) self.assertEqual(result.stdout, expected["stdout"]) def test_wrong_vm_or_process_identity_prevents_command_submission(self) -> None: for changes, pid in (({"vm_id": "vm-other"}, os.getpid()), ({"version": 2}, os.getpid()), ({"username": "root"}, os.getpid()), ({}, os.getpid() + 100000)): with self.subTest(changes=changes, pid=pid): def handle(connection): self.handshake(connection, **changes) self.assertEqual(connection.recv(1), b"") with self.firecracker(handle) as path, self.assertRaises(GuestAgentUnavailableError): GuestSession.connect(path, "vm-test", "worker", pid) def test_connection_loss_after_submission_has_unknown_outcome_without_retry(self) -> None: submissions = [] def handle(connection): self.handshake(connection) submissions.append(agent.receive_frame(connection, agent.MAX_REQUEST_BYTES, time.monotonic() + 2)) with self.firecracker(handle) as path: session = GuestSession.connect(path, "vm-test", "worker", os.getpid()) with self.assertRaises(GuestExecutionError) as raised: try: session.execute("touch example", sudo=False, timeout_seconds=1) finally: session.close() self.assertTrue(raised.exception.outcome_unknown) self.assertEqual(len(submissions), 1) def test_malformed_result_is_not_reported_as_success(self) -> None: for result in ({"exit_code": True, "stdout": "", "stderr": "", "timed_out": False, "truncated": False}, {"error": "bad", "outcome_unknown": "false"}, {"error": "\ud800", "outcome_unknown": False}, {"exit_code": 0}): with self.subTest(result=result): def handle(connection): self.handshake(connection) agent.receive_frame(connection, agent.MAX_REQUEST_BYTES, time.monotonic() + 2) agent.send_frame(connection, result, agent.MAX_RESPONSE_BYTES, time.monotonic() + 2) with self.firecracker(handle) as path: session = GuestSession.connect(path, "vm-test", "worker", os.getpid()) with self.assertRaises(GuestExecutionError) as raised: try: session.execute("id", sudo=False, timeout_seconds=1) finally: session.close() self.assertTrue(raised.exception.outcome_unknown)