172 lines
7.0 KiB
Python
172 lines
7.0 KiB
Python
"""VM-scoped session history, creation, and authenticated browser terminal streams."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import hmac
|
|
import json
|
|
from contextlib import suppress
|
|
from typing import Literal
|
|
from urllib.parse import urlsplit
|
|
|
|
from fastapi import APIRouter, Depends, Request, WebSocket, WebSocketDisconnect
|
|
from pydantic import BaseModel, ConfigDict, Field, StrictInt
|
|
|
|
from uvm.app import Application
|
|
from uvm.errors import UvmError
|
|
from uvm.guest_agent import MAX_TERMINAL_FRAME, ProtocolError, terminal_data
|
|
from uvm.routers.dependencies import get_authorized_application
|
|
|
|
from .sessions import SessionManager
|
|
|
|
|
|
router = APIRouter(prefix="/vms/{vm_id}/sessions", tags=["sessions"])
|
|
|
|
|
|
class SessionCreate(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
rows: StrictInt = Field(default=24, ge=2, le=200)
|
|
cols: StrictInt = Field(default=80, ge=2, le=400)
|
|
|
|
|
|
class SessionResponse(BaseModel):
|
|
id: str
|
|
vm_id: str
|
|
name: str
|
|
username: str
|
|
status: Literal["starting", "active", "closed", "interrupted", "failed"]
|
|
created_at: float
|
|
updated_at: float
|
|
ended_at: float | None
|
|
rows: int
|
|
cols: int
|
|
exit_code: int | None
|
|
error: str | None
|
|
output_start: int
|
|
output_end: int
|
|
truncated: bool
|
|
|
|
|
|
def manager(request: Request, _application: Application = Depends(get_authorized_application)) -> SessionManager:
|
|
return request.app.state.nest_sessions
|
|
|
|
|
|
@router.get("", response_model=list[SessionResponse])
|
|
async def list_sessions(vm_id: str, sessions: SessionManager = Depends(manager)):
|
|
return await sessions.list(vm_id)
|
|
|
|
|
|
@router.post("", response_model=SessionResponse, status_code=201)
|
|
async def create_session(vm_id: str, body: SessionCreate, sessions: SessionManager = Depends(manager)):
|
|
return await sessions.create(vm_id, body.rows, body.cols)
|
|
|
|
|
|
@router.get("/{session_id}", response_model=SessionResponse)
|
|
async def get_session(vm_id: str, session_id: str, sessions: SessionManager = Depends(manager)):
|
|
return await sessions.get(vm_id, session_id)
|
|
|
|
|
|
@router.post("/{session_id}/close", response_model=SessionResponse)
|
|
async def close_session(vm_id: str, session_id: str, sessions: SessionManager = Depends(manager)):
|
|
return await sessions.end(vm_id, session_id)
|
|
|
|
|
|
def _same_origin(websocket: WebSocket) -> bool:
|
|
origin = websocket.headers.get("origin")
|
|
if origin is None:
|
|
return True # Non-browser API clients must still authenticate.
|
|
try:
|
|
parsed = urlsplit(origin)
|
|
except ValueError:
|
|
return False
|
|
scheme = "https" if websocket.url.scheme == "wss" else "http"
|
|
return parsed.scheme == scheme and parsed.netloc.lower() == websocket.headers.get("host", "").lower() and not parsed.path
|
|
|
|
|
|
async def _message(websocket: WebSocket, *, limit: int = MAX_TERMINAL_FRAME) -> dict:
|
|
text = await websocket.receive_text()
|
|
if len(text) > limit:
|
|
raise ProtocolError("terminal message exceeds its size limit")
|
|
message = json.loads(text)
|
|
if not isinstance(message, dict):
|
|
raise ProtocolError("terminal message must be an object")
|
|
return message
|
|
|
|
|
|
@router.websocket("/{session_id}/terminal")
|
|
async def terminal(websocket: WebSocket, vm_id: str, session_id: str):
|
|
if not _same_origin(websocket):
|
|
await websocket.close(code=1008)
|
|
return
|
|
await websocket.accept()
|
|
tasks = []
|
|
try:
|
|
# Browser WebSockets cannot set X-UVM-Token. Authenticate in the first
|
|
# frame so tokens never appear in URLs, access logs, or session history.
|
|
auth = await asyncio.wait_for(_message(websocket, limit=4096), timeout=10)
|
|
token = auth.get("token")
|
|
expected = websocket.app.state.uvm_application.settings.api_token
|
|
if (auth.get("type") != "auth" or not isinstance(token, str) or not token.isascii()
|
|
or not expected or not hmac.compare_digest(token, expected)):
|
|
await websocket.close(code=4401)
|
|
return
|
|
sessions = websocket.app.state.nest_sessions
|
|
await sessions.get(vm_id, session_id)
|
|
offset = auth.get("offset", 0)
|
|
acknowledgements = asyncio.Queue(maxsize=1)
|
|
|
|
async def output():
|
|
cursor = offset
|
|
previous = None
|
|
while True:
|
|
record, start, data = await sessions.read(vm_id, session_id, cursor)
|
|
if previous != record or start != cursor:
|
|
await websocket.send_json({"type": "session", "session": SessionResponse(**record).model_dump(), "offset": start})
|
|
previous = record
|
|
cursor = start
|
|
if data:
|
|
await websocket.send_json({"type": "output", "offset": cursor, "data": base64.b64encode(data).decode("ascii")})
|
|
cursor += len(data)
|
|
# Only one output chunk may be waiting for the terminal renderer.
|
|
acknowledged = await asyncio.wait_for(acknowledgements.get(), timeout=30)
|
|
if type(acknowledged) is not int or acknowledged != cursor:
|
|
raise ProtocolError("invalid terminal output acknowledgement")
|
|
continue
|
|
if record["status"] not in {"active", "starting"}:
|
|
return
|
|
await sessions.wait(session_id)
|
|
|
|
async def controls():
|
|
while True:
|
|
message = await _message(websocket)
|
|
if set(message) == {"type", "data"} and message["type"] == "input":
|
|
await sessions.write(vm_id, session_id, terminal_data(message["data"]))
|
|
elif set(message) == {"type", "rows", "cols"} and message["type"] == "resize":
|
|
await sessions.resize(vm_id, session_id, message["rows"], message["cols"])
|
|
elif set(message) == {"type", "offset"} and message["type"] == "ack":
|
|
if acknowledgements.full():
|
|
raise ProtocolError("unexpected terminal output acknowledgement")
|
|
acknowledgements.put_nowait(message["offset"])
|
|
else:
|
|
raise ProtocolError("invalid terminal control message")
|
|
|
|
tasks = [asyncio.create_task(output()), asyncio.create_task(controls())]
|
|
done, _pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
|
|
for task in done:
|
|
task.result()
|
|
except WebSocketDisconnect:
|
|
pass
|
|
except (ValueError, KeyError, TypeError, UnicodeError, RecursionError, TimeoutError):
|
|
with suppress(WebSocketDisconnect, RuntimeError):
|
|
await websocket.send_json({"type": "error", "message": "Invalid terminal message or authentication timeout."})
|
|
except UvmError as error:
|
|
with suppress(WebSocketDisconnect, RuntimeError):
|
|
await websocket.send_json({"type": "error", "message": str(error)})
|
|
finally:
|
|
for task in tasks:
|
|
task.cancel()
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
with suppress(WebSocketDisconnect, RuntimeError):
|
|
await websocket.close()
|