kstyagi
e1b934e35e
- Added `guest_tools.py` for standalone tool installer functionality. - Introduced `ToolJobStore` in `tool_store.py` to manage per-VM tool job records. - Created `ToolService` in `tools.py` for background installation workers using the guest execution channel. - Enhanced `LifecycleService` in `lifecycle.py` to support tool installation queuing and preparation. - Updated `guest_exec.py` to handle connection shutdown gracefully. - Added validation and normalization functions for tool specifications. - Implemented logging and error handling for tool installation processes. - Ensured compatibility with Ubuntu 22.04 and ARM64/x86_64 architectures.
121 lines
5.9 KiB
Python
121 lines
5.9 KiB
Python
"""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; wait for guest startup, or restart the VM "
|
|
"to apply automatic vsock guest-agent setup (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:
|
|
# Wake a thread waiting for a long-running tool command on server shutdown.
|
|
# The guest still owns its bounded command and package-manager cleanup.
|
|
try:
|
|
self._connection.shutdown(socket.SHUT_RDWR)
|
|
except OSError:
|
|
pass
|
|
self._connection.close()
|