Add HTTPS console streaming and tokens

Author: msinhore

firecracker.py

@@ -367,6 +367,13 @@ def op_console(ctx):
     """POST /v1/vms/{name}/console — obtain VNC bridge connection info."""
     url = f"{ctx.agent.base_url}/vms/{ctx.vm_name}/console"
     data = _json_or_fail(_req("POST", url, ctx.agent))
+    console_obj = data.get("console")
+    if isinstance(console_obj, dict):
+        path = console_obj.get("path")
+        if path and not console_obj.get("url"):
+            base = ctx.agent.base_url.rsplit("/v1", 1)[0]
+            console_obj["url"] = f"{base}{path}"
+        _ok(data)
     console_host = getattr(ctx.agent, "console_host", "") or ""
     resp_host = data.get("host")
     if console_host and (not resp_host or resp_host in {"0.0.0.0", "127.0.0.1", "::"}):

host-agent/api/handlers.py

@@ -25,19 +25,30 @@
 from utils.validation import validate_name
 from utils.vnc_console import VNCConsoleManager
 
+try:
+    from utils.console_tokens import ConsoleTokenManager
+except ImportError:  # pragma: no cover
+    ConsoleTokenManager = None  # type: ignore
+
 logger = logging.getLogger("fc-agent")
 
 
 class APIHandlers:
 
-    def __init__(self, agent_defaults: Dict[str, Any], ui_config: Optional[Dict[str, Any]] = None):
+    def __init__(
+        self,
+        agent_defaults: Dict[str, Any],
+        ui_config: Optional[Dict[str, Any]] = None,
+        console_token_manager: Optional["ConsoleTokenManager"] = None,
+    ):
         self.agent_defaults = agent_defaults
         self.ui_config = ui_config or {"enabled": True, "session_timeout_seconds": 1800}
         self.vm_manager = VMManager()
         self.vm_lifecycle = VMLifecycle(agent_defaults)
         self.config_manager = ConfigManager(agent_defaults)
         self.state_manager = StateManager(agent_defaults)
         self.vnc_console = VNCConsoleManager(agent_defaults)
+        self.console_tokens = console_token_manager
 
     def v1_ui_config(self) -> Dict[str, Any]:
         cfg = self.ui_config or {}
@@ -477,14 +488,33 @@ def v1_vm_console_start(self, vm_name: str) -> Dict[str, Any]:
         """Start (or reuse) the VNC bridge for a VM console."""
         try:
             info = self.vnc_console.ensure_console(vm_name)
+            console_obj: Dict[str, Any]
+            if info.get("direct_proxy") and self.console_tokens:
+                token = self.console_tokens.issue(
+                    vm_name,
+                    info.get("port"),
+                    info.get("password"),
+                    info.get("session_name", ""),
+                )
+                console_obj = {
+                    "path": f"/console/ws/{token}",
+                    "protocol": "https",
+                    "passwordonetimeuseonly": True,
+                }
+            else:
+                console_obj = {
+                    "host": info.get("bind_host") or self.vnc_console.bind_host,
+                    "port": info.get("port"),
+                    "password": info.get("password"),
+                    "protocol": "vnc",
+                    "passwordonetimeuseonly": False,
+                }
             return {
                 "status": "success",
                 "message": f"VNC console ready for VM {vm_name}",
                 "vm_name": vm_name,
-                "host": info.get("host"),
-                "port": info.get("port"),
-                "password": info.get("password"),
                 "created_at": info.get("created_at"),
+                "console": console_obj,
             }
         except RuntimeError as exc:
             raise HTTPException(status_code=400, detail=str(exc)) from exc

host-agent/api/routes.py

@@ -14,9 +14,10 @@ def register_routes(
     agent_defaults: Dict[str, Any],
     auth_dependency: Optional[Any] = None,
     ui_config: Optional[Dict[str, Any]] = None,
+    console_token_manager: Optional[Any] = None,
 ) -> None:
     """Register all API routes with the FastAPI application."""
-    handlers = APIHandlers(agent_defaults, ui_config=ui_config)
+    handlers = APIHandlers(agent_defaults, ui_config=ui_config, console_token_manager=console_token_manager)
     deps = [Depends(auth_dependency)] if auth_dependency else []
     protected = {"dependencies": deps} if deps else {}
 

host-agent/debian/firecracker-agent.json

@@ -26,7 +26,9 @@
       "geometry": "1024x768x24",
       "xterm_geometry": "127x45",
       "font_family": "Monospace",
-      "font_size": 10
+      "font_size": 10,
+      "direct_proxy_enabled": false,
+      "direct_token_ttl": 300
     }
   },
   "security": {

host-agent/debian/install

@@ -25,6 +25,7 @@ host-agent/utils/validation.py                      usr/lib/firecracker-cloudsta
 host-agent/utils/filesystem.py                      usr/lib/firecracker-cloudstack-agent/utils/
 host-agent/utils/tmux.py                            usr/lib/firecracker-cloudstack-agent/utils/
 host-agent/utils/vnc_console.py                     usr/lib/firecracker-cloudstack-agent/utils/
+host-agent/utils/console_tokens.py                  usr/lib/firecracker-cloudstack-agent/utils/
 host-agent/tools/fc_shutdown_service.py             usr/lib/firecracker-cloudstack-agent/tools/
 host-agent/api/__init__.py                          usr/lib/firecracker-cloudstack-agent/api/
 host-agent/api/handlers.py                          usr/lib/firecracker-cloudstack-agent/api/

host-agent/firecracker-agent.json-file-example

@@ -18,6 +18,18 @@
     "net": {
       "driver": "linux-bridge-vlan",
       "host_bridge": "cloudbr1"
+    },
+    "console": {
+      "bind_host": "0.0.0.0",
+      "port_min": 5900,
+      "port_max": 5999,
+      "geometry": "1024x768x24",
+      "xterm_geometry": "127x45",
+      "font_family": "Monospace",
+      "font_size": 10,
+      "read_only": false,
+      "direct_proxy_enabled": true,
+      "direct_token_ttl": 300
     }
   },
   "security": {
@@ -34,4 +46,3 @@
     "service": "firecracker-agent"
   }
 }
-

host-agent/firecracker-agent.py

@@ -18,6 +18,8 @@
 # under the License.
 from __future__ import annotations
 
+import asyncio
+import contextlib
 import logging
 import os
 import ssl
@@ -26,7 +28,7 @@
 
 import typer
 import uvicorn
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect
 from fastapi.staticfiles import StaticFiles
 from fastapi.exceptions import RequestValidationError
 from starlette.responses import FileResponse, JSONResponse, RedirectResponse
@@ -39,6 +41,7 @@
 from orchestration import VMLifecycle
 from utils.auth import PamError, build_auth_dependency
 from utils.filesystem import set_agent_defaults
+from utils.console_tokens import ConsoleTokenManager
 
 # Global variables
 logger = logging.getLogger("fc-agent")
@@ -58,6 +61,7 @@
 IS_API_MODE = True
 UI_STATIC_MOUNTED = False
 UI_CONFIG: Dict[str, Any] = {"enabled": True, "session_timeout_seconds": 1800}
+CONSOLE_TOKEN_MANAGER: Optional[ConsoleTokenManager] = None
 # Initialize FastAPI app
 app = FastAPI(title="Firecracker Agent", version="1.0.0")
 
@@ -198,6 +202,67 @@ def v1_config_effective() -> Dict[str, Any]:
     return {"status": "success", "config": AGENT_DEFAULTS}
 
 
+@app.websocket("/console/ws/{token}")
+async def console_websocket(token: str, websocket: WebSocket):
+    if not CONSOLE_TOKEN_MANAGER:
+        await websocket.close(code=1008)
+        return
+    info = CONSOLE_TOKEN_MANAGER.consume(token)
+    if not info:
+        await websocket.close(code=1008)
+        return
+    port = int(info.get("port", 0) or 0)
+    if port <= 0:
+        await websocket.close(code=1008)
+        return
+    try:
+        reader, writer = await asyncio.open_connection("127.0.0.1", port)
+    except Exception:
+        await websocket.close(code=1011)
+        return
+
+    await websocket.accept()
+    close_event = asyncio.Event()
+
+    async def ws_to_tcp():
+        try:
+            while True:
+                msg = await websocket.receive()
+                if msg["type"] == "websocket.disconnect":
+                    break
+                data = msg.get("bytes")
+                if data is None:
+                    text = msg.get("text")
+                    if text is None:
+                        continue
+                    data = text.encode("utf-8")
+                writer.write(data)
+                await writer.drain()
+        except WebSocketDisconnect:
+            pass
+        finally:
+            close_event.set()
+            writer.close()
+            with contextlib.suppress(Exception):
+                await writer.wait_closed()
+
+    async def tcp_to_ws():
+        try:
+            while not close_event.is_set():
+                data = await reader.read(4096)
+                if not data:
+                    break
+                await websocket.send_bytes(data)
+        except Exception:
+            pass
+        finally:
+            close_event.set()
+            with contextlib.suppress(Exception):
+                await websocket.close()
+
+    await asyncio.gather(ws_to_tcp(), tcp_to_ws())
+
+
 # FastAPI event handlers
 @app.on_event("startup")
 async def startup_event():
@@ -210,6 +275,11 @@ async def startup_event():
     AGENT_DEFAULTS = AGENT_CFG.get("defaults", {})
     config_manager.agent_defaults = AGENT_DEFAULTS
     set_agent_defaults(AGENT_DEFAULTS)
+    console_cfg = AGENT_DEFAULTS.get("console", {}) if isinstance(AGENT_DEFAULTS, dict) else {}
+    token_ttl = int(console_cfg.get("direct_token_ttl", 300))
+    direct_enabled = bool(console_cfg.get("direct_proxy_enabled", False))
+    global CONSOLE_TOKEN_MANAGER
+    CONSOLE_TOKEN_MANAGER = ConsoleTokenManager(ttl_seconds=token_ttl) if direct_enabled else None
 
     logger.info("Configuration loaded successfully")
     logger.info("AGENT_CFG keys: %s", list(AGENT_CFG.keys()))
@@ -224,7 +294,7 @@ async def startup_event():
     # Register API routes with loaded configuration
     if AUTH_DEPENDENCY is None:
         AUTH_DEPENDENCY = _configure_auth_dependency(AGENT_CFG.get("auth", {}))
-    register_routes(app, AGENT_DEFAULTS, AUTH_DEPENDENCY, UI_CONFIG)
+    register_routes(app, AGENT_DEFAULTS, AUTH_DEPENDENCY, UI_CONFIG, console_token_manager=CONSOLE_TOKEN_MANAGER)
     if UI_CONFIG.get("enabled"):
         _mount_ui_static_if_available()
     else:

host-agent/utils/console_tokens.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""Simple in-memory console token manager."""
+import secrets
+import threading
+import time
+from typing import Dict, Optional
+
+
+class ConsoleTokenManager:
+    def __init__(self, ttl_seconds: int = 300):
+        self.ttl = int(ttl_seconds or 300)
+        self._tokens: Dict[str, Dict[str, object]] = {}
+        self._lock = threading.Lock()
+
+    def issue(self, vm_name: str, port: int, password: Optional[str], session_name: str) -> str:
+        token = secrets.token_urlsafe(16)
+        entry = {
+            "vm_name": vm_name,
+            "port": int(port),
+            "password": password,
+            "session": session_name,
+            "expires_at": time.time() + self.ttl,
+        }
+        with self._lock:
+            self._tokens[token] = entry
+        return token
+
+    def consume(self, token: str) -> Optional[Dict[str, object]]:
+        if not token:
+            return None
+        with self._lock:
+            entry = self._tokens.pop(token, None)
+        if not entry:
+            return None
+        if entry["expires_at"] < time.time():
+            return None
+        return entry
+
+    def purge_expired(self) -> None:
+        cutoff = time.time()
+        with self._lock:
+            expired = [tok for tok, info in self._tokens.items() if info["expires_at"] < cutoff]
+            for tok in expired:
+                self._tokens.pop(tok, None)

host-agent/utils/vnc_console.py

@@ -48,6 +48,10 @@ def __init__(self, agent_defaults: Dict[str, Any]):
         self.font_family = console_defaults.get("font_family") or "Monospace"
         self.font_size = int(console_defaults.get("font_size", 10))
         self.read_only = bool(console_defaults.get("read_only", False))
+        self.direct_proxy = bool(console_defaults.get("direct_proxy_enabled", False))
+        if self.direct_proxy:
+            self.bind_host = "127.0.0.1"
+        self.require_password = not self.direct_proxy
 
         self.tmux = TmuxManager()
 
@@ -58,7 +62,7 @@ def ensure_console(self, vm_name: str) -> Dict[str, Any]:
         current_state = self._load_state(vm_state_path)
         if current_state and self._state_active(current_state):
             logger.debug("Reusing existing VNC console for %s (state=%s)", vm_name, current_state)
-            return self._response_payload(current_state)
+            return current_state
 
         if current_state:
             logger.debug("Cleaning up stale VNC console for %s", vm_name)
@@ -70,8 +74,11 @@ def ensure_console(self, vm_name: str) -> Dict[str, Any]:
             raise RuntimeError(f"tmux session {session_name} not found; VM console is not available")
 
         port = self._allocate_port()
-        password = self._generate_password()
-        password_file = self._write_password_file(vm_name, password)
+        password = None
+        password_file = None
+        if self.require_password:
+            password = self._generate_password()
+            password_file = self._write_password_file(vm_name, password)
         display, xvfb_proc = self._start_xvfb(vm_name)
         xterm_proc = self._start_xterm(display, vm_name, session_name)
         x11vnc_proc = self._start_x11vnc(vm_name, display, port, password_file)
@@ -85,14 +92,15 @@ def ensure_console(self, vm_name: str) -> Dict[str, Any]:
             "x11vnc_pid": x11vnc_proc.pid,
             "port": port,
             "password": password,
-            "password_file": str(password_file),
+            "password_file": str(password_file) if password_file else None,
             "bind_host": self.bind_host,
             "session_name": session_name,
+            "direct_proxy": self.direct_proxy,
         }
         self._write_state(vm_state_path, state)
         logger.debug("VNC console state stored for %s: %s", vm_name, state)
         self._monitor_console(vm_name, state)
-        return self._response_payload(state)
+        return state
 
     def stop_console(self, vm_name: str) -> Dict[str, Any]:
         """Terminate VNC bridge for the VM."""
@@ -150,14 +158,7 @@ def _state_active(self, state: Dict[str, Any]) -> bool:
         return True
 
     def _response_payload(self, state: Dict[str, Any]) -> Dict[str, Any]:
-        return {
-            "status": "success",
-            "vm_name": state.get("vm_name"),
-            "host": state.get("bind_host", self.bind_host),
-            "port": int(state.get("port")),
-            "password": state.get("password"),
-            "created_at": state.get("created_at"),
-        }
+        return state
 
     def _cleanup_state(self, state: Dict[str, Any]) -> None:
         for key in ("x11vnc_pid", "xterm_pid", "xvfb_pid"):
@@ -339,15 +340,18 @@ def _start_x11vnc(self, vm_name: str, display: str, port: int, password_file: Pa
             display,
             "-rfbport",
             str(port),
-            "-rfbauth",
-            str(password_file),
             "-once",
             "-shared",
             "-noxdamage",
             "-nolookup",
             "-o",
             str(log_file),
         ]
+        if self.direct_proxy:
+            cmd.append("-localhost")
+            cmd.append("-nopw")
+        elif password_file:
+            cmd.extend(["-rfbauth", str(password_file)])
         if self.bind_host not in ("127.0.0.1", "::1"):
             cmd.extend(["-listen", self.bind_host])
         else: