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.
2026-09-25 17:39:06 +05:30
parent dca7f78dae
commit b0bf99a6ee
25 changed files with 1893 additions and 30 deletions
@@ -0,0 +1,114 @@
+"""Host-side, one-shot guest execution over Firecracker's vsock Unix socket."""
+
+from __future__ import annotations
+
+import socket
+import stat
+import struct
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+from .errors import GuestAgentUnavailableError, GuestExecutionError, ValidationError
+from .guest_agent import (
+    AGENT_PORT, CLEANUP_SECONDS, CONNECT_TIMEOUT_SECONDS, MAX_OUTPUT_BYTES,
+    MAX_REQUEST_BYTES, MAX_RESPONSE_BYTES, PROTOCOL_VERSION, ProtocolError,
+    receive_exact, receive_frame, send_frame, validate_command,
+)
+
+
+@dataclass(frozen=True, slots=True)
+class ExecResult:
+    exit_code: int | None
+    stdout: str
+    stderr: str
+    timed_out: bool
+    truncated: bool
+
+
+def validate_exec(command: object, sudo: object, timeout_seconds: object) -> None:
+    try:
+        validate_command(command, sudo, timeout_seconds)
+    except ProtocolError as error:
+        raise ValidationError(str(error)) from error
+
+
+class GuestSession:
+    """A connected session pins the selected VMM; never reconnect after submission."""
+
+    def __init__(self, connection: socket.socket) -> None:
+        self._connection = connection
+
+    @classmethod
+    def connect(cls, path: Path, vm_id: str, username: str, pid: int) -> GuestSession:
+        connection = None
+        try:
+            connection = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+            metadata = path.lstat()
+            if not stat.S_ISSOCK(metadata.st_mode):
+                raise ProtocolError("guest socket is not a Unix socket")
+            deadline = time.monotonic() + CONNECT_TIMEOUT_SECONDS
+            connection.settimeout(CONNECT_TIMEOUT_SECONDS)
+            connection.connect(str(path))
+            peer_pid, _uid, _gid = struct.unpack("3i", connection.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12))
+            if peer_pid != pid:
+                raise ProtocolError("guest socket does not belong to the selected VM process")
+            connection.sendall(f"CONNECT {AGENT_PORT}\n".encode("ascii"))
+            acknowledgement = bytearray()
+            while not acknowledgement.endswith(b"\n") and len(acknowledgement) < 32:
+                acknowledgement.extend(receive_exact(connection, 1, deadline))
+            parts = bytes(acknowledgement).split()
+            if (not acknowledgement.endswith(b"\n") or len(parts) != 2
+                    or parts[0] != b"OK" or not parts[1].isdigit()):
+                raise ProtocolError("invalid Firecracker vsock acknowledgement")
+            hello = receive_frame(connection, 4096, deadline)
+            if (set(hello) != {"version", "vm_id", "username"}
+                    or type(hello["version"]) is not int or hello["version"] != PROTOCOL_VERSION
+                    or hello["vm_id"] != vm_id or hello["username"] != username):
+                raise ProtocolError("guest agent identity or protocol version does not match this VM")
+            return cls(connection)
+        except (OSError, EOFError, ProtocolError) as error:
+            if connection is not None:
+                connection.close()
+            raise GuestAgentUnavailableError(
+                "guest agent is unavailable or not ready; install it with POST /vms/{vm_id}/agent/install "
+                "while stopped, then start the VM (vsock-capable kernel required)"
+            ) from error
+
+    def execute(self, command: str, *, sudo: bool, timeout_seconds: float) -> ExecResult:
+        validate_exec(command, sudo, timeout_seconds)
+        try:
+            send_frame(self._connection, {
+                "version": PROTOCOL_VERSION, "command": command,
+                "sudo": sudo, "timeout_seconds": timeout_seconds,
+            }, MAX_REQUEST_BYTES, time.monotonic() + CONNECT_TIMEOUT_SECONDS)
+            result = receive_frame(
+                self._connection, MAX_RESPONSE_BYTES,
+                time.monotonic() + timeout_seconds + CLEANUP_SECONDS + CONNECT_TIMEOUT_SECONDS,
+            )
+            if "error" in result:
+                if (set(result) != {"error", "outcome_unknown"} or not isinstance(result["error"], str)
+                        or len(result["error"]) > 4096 or type(result["outcome_unknown"]) is not bool):
+                    raise ProtocolError("invalid guest failure response")
+                result["error"].encode("utf-8")  # Reject invalid Unicode before rendering an HTTP error.
+                raise GuestExecutionError(result["error"], outcome_unknown=result["outcome_unknown"])
+            if set(result) != {"exit_code", "stdout", "stderr", "timed_out", "truncated"}:
+                raise ProtocolError("invalid guest result fields")
+            if (type(result["timed_out"]) is not bool or type(result["truncated"]) is not bool
+                    or not isinstance(result["stdout"], str) or not isinstance(result["stderr"], str)
+                    or (result["timed_out"] and result["exit_code"] is not None)
+                    or (not result["timed_out"] and (
+                        type(result["exit_code"]) is not int or not -64 <= result["exit_code"] <= 255))):
+                raise ProtocolError("invalid guest result types")
+            if sum(len(result[name].encode("utf-8")) for name in ("stdout", "stderr")) > 3 * MAX_OUTPUT_BYTES:
+                raise ProtocolError("guest output exceeds its limit")
+            return ExecResult(**result)
+        except (OSError, EOFError, ProtocolError, UnicodeError) as error:
+            raise GuestExecutionError(
+                "guest execution connection was lost or returned an invalid result; "
+                "the command may have run and will not be retried",
+                outcome_unknown=True,
+            ) from error
+
+    def close(self) -> None:
+        self._connection.close()