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,289 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user