feat: add terminal support for guest execution agent and enhance lifecycle management

uvm/guest_terminal.py

@@ -0,0 +1,89 @@
+"""Interactive, process-pinned guest terminal transport over the existing vsock agent."""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import json
+import struct
+import time
+
+from .errors import GuestAgentUnavailableError, GuestExecutionError, ValidationError
+from .guest_agent import (
+    CONNECT_TIMEOUT_SECONDS, MAX_TERMINAL_FRAME, PROTOCOL_VERSION, ProtocolError,
+    receive_frame, send_frame, terminal_data, validate_terminal_size,
+)
+from .guest_exec import GuestSession
+
+
+def validate_size(rows: object, cols: object) -> None:
+    try:
+        validate_terminal_size(rows, cols)
+    except ProtocolError as error:
+        raise ValidationError(str(error)) from error
+
+
+class GuestTerminal(GuestSession):
+    def start(self, rows: int, cols: int) -> None:
+        validate_size(rows, cols)
+        try:
+            deadline = time.monotonic() + CONNECT_TIMEOUT_SECONDS
+            send_frame(self._connection, {"version": PROTOCOL_VERSION, "terminal": {"rows": rows, "cols": cols}},
+                       MAX_TERMINAL_FRAME, deadline)
+            response = receive_frame(self._connection, MAX_TERMINAL_FRAME, deadline)
+            if response != {"type": "ready"}:
+                raise ProtocolError("terminal support is unavailable")
+            self._connection.setblocking(False)
+            self._write_lock = asyncio.Lock()
+        except (OSError, EOFError, ProtocolError) as error:
+            self.close()
+            raise GuestAgentUnavailableError(
+                "Interactive terminal support is unavailable. Stop this VM, use Install / update terminal support "
+                "in its Sessions view, then start it and launch a new session."
+            ) from error
+
+    async def _receive_exact(self, size: int) -> bytes:
+        data = bytearray()
+        while len(data) < size:
+            chunk = await asyncio.get_running_loop().sock_recv(self._connection, size - len(data))
+            if not chunk:
+                raise EOFError("guest terminal connection closed")
+            data.extend(chunk)
+        return bytes(data)
+
+    async def read(self) -> dict:
+        try:
+            # An idle shell has no output deadline; partial messages are bounded.
+            first = await self._receive_exact(1)
+            async with asyncio.timeout(CONNECT_TIMEOUT_SECONDS):
+                size, = struct.unpack("!I", first + await self._receive_exact(3))
+                if not 0 < size <= MAX_TERMINAL_FRAME:
+                    raise ProtocolError("invalid terminal frame length")
+                message = json.loads(await self._receive_exact(size))
+            if isinstance(message, dict):
+                if set(message) == {"type", "data"} and message["type"] == "output":
+                    return {"type": "output", "data": terminal_data(message["data"])}
+                if (set(message) == {"type", "exit_code"} and message["type"] == "exit"
+                        and type(message["exit_code"]) is int and -64 <= message["exit_code"] <= 255):
+                    return message
+            raise ProtocolError("invalid guest terminal response")
+        except (OSError, EOFError, ValueError, UnicodeError, RecursionError) as error:
+            raise GuestExecutionError("The guest terminal connection ended unexpectedly.", outcome_unknown=True) from error
+
+    async def _send(self, message: dict) -> None:
+        payload = json.dumps(message, separators=(",", ":")).encode("ascii")
+        if len(payload) > MAX_TERMINAL_FRAME:
+            raise ValidationError("terminal input exceeds its size limit")
+        async with self._write_lock:
+            async with asyncio.timeout(CONNECT_TIMEOUT_SECONDS):
+                await asyncio.get_running_loop().sock_sendall(self._connection, struct.pack("!I", len(payload)) + payload)
+
+    async def write(self, data: bytes) -> None:
+        await self._send({"type": "input", "data": base64.b64encode(data).decode("ascii")})
+
+    async def resize(self, rows: int, cols: int) -> None:
+        validate_size(rows, cols)
+        await self._send({"type": "resize", "rows": rows, "cols": cols})
+
+    async def end(self) -> None:
+        await self._send({"type": "close"})