add remote fleet terminal sessions
This commit is contained in:
@@ -1,7 +1,7 @@
|
||||
# Caddy template for wakey-control-plane with Cloudflare Access.
|
||||
#
|
||||
# Security model:
|
||||
# - Public endpoints for agents: /healthz, /api/v1/agents/enroll, /api/v1/agent/ws
|
||||
# - Public endpoints for agents: health, enrollment, control WS, and scoped terminal relay WS
|
||||
# - Private admin surface: /ui/* and /api/v1/control/* (requires CF Access headers)
|
||||
#
|
||||
# Replace cp.example.com with your public control-plane domain.
|
||||
@@ -10,7 +10,7 @@ wakey.ldlda.com {
|
||||
encode zstd gzip
|
||||
|
||||
# Public agent-facing endpoints.
|
||||
@public path /healthz /api/v1/agents/enroll /api/v1/agent/ws
|
||||
@public path /healthz /api/v1/agents/enroll /api/v1/agent/ws /api/v1/agent/terminals/*
|
||||
handle @public {
|
||||
reverse_proxy 127.0.0.1:6767
|
||||
}
|
||||
|
||||
@@ -39,7 +39,7 @@ Use `pty-process` with its async feature as the initial PTY abstraction because
|
||||
- child session-leader and controlling-terminal setup; and
|
||||
- a small implementation built on `rustix`.
|
||||
|
||||
Adoption remains gated on successfully cross-compiling for the router target and exercising `/bin/ash` or the configured shell on a real device. Wakey does not call `forkpty` or add its own unsafe PTY implementation.
|
||||
The PTY wrapper is cross-compiled for `armv7-unknown-linux-musleabihf` and its input, output, resize, and exit test has passed on the target router. That focused test remains part of the remote compatibility workflow. Wakey does not call `forkpty` or add its own unsafe PTY implementation.
|
||||
|
||||
The executable, working directory, environment, UID, and GID are agent-controlled configuration. The browser cannot supply an arbitrary executable or process environment.
|
||||
|
||||
|
||||
@@ -15,6 +15,8 @@
|
||||
"@base-ui/react": "^1.4.0",
|
||||
"@fontsource-variable/inter": "^5.2.8",
|
||||
"@tailwindcss/vite": "^4.2.2",
|
||||
"@xterm/addon-fit": "^0.11.0",
|
||||
"@xterm/xterm": "^6.0.0",
|
||||
"class-variance-authority": "^0.7.1",
|
||||
"clsx": "^2.1.1",
|
||||
"lucide-react": "^1.8.0",
|
||||
|
||||
Generated
+22
@@ -16,6 +16,12 @@ importers:
|
||||
"@tailwindcss/vite":
|
||||
specifier: ^4.2.2
|
||||
version: 4.2.2([email protected](@types/[email protected])([email protected]))
|
||||
"@xterm/addon-fit":
|
||||
specifier: ^0.11.0
|
||||
version: 0.11.0
|
||||
"@xterm/xterm":
|
||||
specifier: ^6.0.0
|
||||
version: 6.0.0
|
||||
class-variance-authority:
|
||||
specifier: ^0.7.1
|
||||
version: 0.7.1
|
||||
@@ -1242,6 +1248,18 @@ packages:
|
||||
peerDependencies:
|
||||
vite: ^4.2.0 || ^5.0.0 || ^6.0.0 || ^7.0.0
|
||||
|
||||
"@xterm/[email protected]":
|
||||
resolution:
|
||||
{
|
||||
integrity: sha512-jYcgT6xtVYhnhgxh3QgYDnnNMYTcf8ElbxxFzX0IZo+vabQqSPAjC3c1wJrKB5E19VwQei89QCiZZP86DCPF7g==,
|
||||
}
|
||||
|
||||
"@xterm/[email protected]":
|
||||
resolution:
|
||||
{
|
||||
integrity: sha512-TQwDdQGtwwDt+2cgKDLn0IRaSxYu1tSUjgKarSDkUM0ZNiSRXFpjxEsvc/Zgc5kq5omJ+V0a8/kIM2WD3sMOYg==,
|
||||
}
|
||||
|
||||
[email protected]:
|
||||
resolution:
|
||||
{
|
||||
@@ -4203,6 +4221,10 @@ snapshots:
|
||||
transitivePeerDependencies:
|
||||
- supports-color
|
||||
|
||||
"@xterm/[email protected]": {}
|
||||
|
||||
"@xterm/[email protected]": {}
|
||||
|
||||
[email protected]:
|
||||
dependencies:
|
||||
mime-types: 3.0.2
|
||||
|
||||
+24
-1
@@ -1,4 +1,4 @@
|
||||
import { useEffect, useState } from "react";
|
||||
import { lazy, Suspense, useEffect, useState } from "react";
|
||||
import { Navigate, Route, Routes } from "react-router-dom";
|
||||
import { toast } from "sonner";
|
||||
|
||||
@@ -15,6 +15,7 @@ import {
|
||||
setAgentNickname,
|
||||
} from "@/api";
|
||||
import { Toaster } from "@/components/ui/sonner";
|
||||
import { Skeleton } from "@/components/ui/skeleton";
|
||||
import { AppLayout } from "@/layout/AppLayout";
|
||||
import { AgentsPage } from "@/pages/AgentsPage";
|
||||
import { AlertsPage } from "@/pages/AlertsPage";
|
||||
@@ -23,6 +24,12 @@ import { CommandsPage } from "@/pages/CommandsPage";
|
||||
import { DashboardPage } from "@/pages/DashboardPage";
|
||||
import { DevicesPage } from "@/pages/DevicesPage";
|
||||
import { TokensPage } from "@/pages/TokensPage";
|
||||
|
||||
const TerminalPage = lazy(() =>
|
||||
import("@/pages/TerminalPage").then((module) => ({
|
||||
default: module.TerminalPage,
|
||||
})),
|
||||
);
|
||||
import { WakeToolsPage } from "@/pages/WakeToolsPage";
|
||||
|
||||
type LoadState = "idle" | "loading" | "ready" | "error";
|
||||
@@ -215,6 +222,22 @@ export function App() {
|
||||
}
|
||||
/>
|
||||
<Route path="tokens" element={<TokensPage />} />
|
||||
<Route
|
||||
path="terminal"
|
||||
element={
|
||||
<Suspense
|
||||
fallback={
|
||||
<Skeleton className="h-[calc(100dvh-2.5rem)] w-full" />
|
||||
}
|
||||
>
|
||||
<TerminalPage
|
||||
agents={agents}
|
||||
selectedAgentId={selectedAgentId}
|
||||
onSelectAgent={setSelectedAgentId}
|
||||
/>
|
||||
</Suspense>
|
||||
}
|
||||
/>
|
||||
|
||||
<Route path="*" element={<Navigate to="/" replace />} />
|
||||
</Route>
|
||||
|
||||
@@ -2,6 +2,17 @@ export type Agent = {
|
||||
agent_id: string;
|
||||
connected: boolean;
|
||||
nickname?: string | null;
|
||||
capabilities: "terminal"[];
|
||||
};
|
||||
|
||||
export type TerminalSession = {
|
||||
terminal_id: string;
|
||||
agent_id: string;
|
||||
created_at_unix: number;
|
||||
agent_attached: boolean;
|
||||
operator_attached: boolean;
|
||||
websocket_url: string;
|
||||
attachment_token?: string;
|
||||
};
|
||||
|
||||
export type Alert = {
|
||||
@@ -225,6 +236,34 @@ export function fetchAgents(): Promise<Agent[]> {
|
||||
return request<Agent[]>("/api/v1/control/agents");
|
||||
}
|
||||
|
||||
export function createTerminal(
|
||||
agentId: string,
|
||||
rows: number,
|
||||
cols: number,
|
||||
): Promise<TerminalSession> {
|
||||
return request<TerminalSession>("/api/v1/control/terminals", {
|
||||
method: "POST",
|
||||
body: JSON.stringify({ agent_id: agentId, rows, cols }),
|
||||
});
|
||||
}
|
||||
|
||||
export function attachTerminal(terminalId: string): Promise<TerminalSession> {
|
||||
return request<TerminalSession>(
|
||||
`/api/v1/control/terminals/${encodeURIComponent(terminalId)}/attach`,
|
||||
{ method: "POST" },
|
||||
);
|
||||
}
|
||||
|
||||
export async function closeTerminal(terminalId: string): Promise<void> {
|
||||
const response = await fetch(
|
||||
`/api/v1/control/terminals/${encodeURIComponent(terminalId)}`,
|
||||
{ method: "DELETE" },
|
||||
);
|
||||
if (!response.ok) {
|
||||
throw new Error(`Failed to close terminal: ${response.status}`);
|
||||
}
|
||||
}
|
||||
|
||||
export function fetchAlerts(): Promise<Alert[]> {
|
||||
return request<Alert[]>("/api/v1/control/alerts");
|
||||
}
|
||||
|
||||
@@ -12,6 +12,7 @@ import {
|
||||
FileText,
|
||||
Key,
|
||||
Terminal,
|
||||
SquareTerminal,
|
||||
PanelLeftClose,
|
||||
PanelLeft,
|
||||
} from "lucide-react";
|
||||
@@ -57,6 +58,7 @@ const navSections: NavSection[] = [
|
||||
{ to: "/audit", label: "Audit", icon: FileText },
|
||||
{ to: "/tokens", label: "Tokens", icon: Key },
|
||||
{ to: "/commands", label: "Commands", icon: Terminal },
|
||||
{ to: "/terminal", label: "Terminal", icon: SquareTerminal },
|
||||
],
|
||||
},
|
||||
];
|
||||
|
||||
@@ -4,6 +4,7 @@ import { BrowserRouter } from "react-router-dom";
|
||||
import { ThemeProvider } from "next-themes";
|
||||
import { TooltipProvider } from "@/components/ui/tooltip";
|
||||
import { App } from "@/App";
|
||||
import "@xterm/xterm/css/xterm.css";
|
||||
import "./styles.css";
|
||||
|
||||
createRoot(document.getElementById("root")!).render(
|
||||
|
||||
@@ -0,0 +1,380 @@
|
||||
import { useCallback, useEffect, useRef, useState } from "react";
|
||||
import { FitAddon } from "@xterm/addon-fit";
|
||||
import { Terminal as XTerm } from "@xterm/xterm";
|
||||
import { Eraser, PlugZap, RotateCcw, Square, Terminal } from "lucide-react";
|
||||
import { toast } from "sonner";
|
||||
|
||||
import {
|
||||
type Agent,
|
||||
type TerminalSession,
|
||||
attachTerminal,
|
||||
closeTerminal,
|
||||
createTerminal,
|
||||
} from "@/api";
|
||||
import { AgentSelector, displayAgentLabel } from "@/components/AgentSelector";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import {
|
||||
Tooltip,
|
||||
TooltipContent,
|
||||
TooltipTrigger,
|
||||
} from "@/components/ui/tooltip";
|
||||
|
||||
type Props = {
|
||||
agents: Agent[];
|
||||
selectedAgentId: string;
|
||||
onSelectAgent: (agentId: string) => void;
|
||||
};
|
||||
|
||||
type ConnectionState =
|
||||
| "idle"
|
||||
| "connecting"
|
||||
| "ready"
|
||||
| "disconnected"
|
||||
| "exited";
|
||||
|
||||
function websocketUrl(path: string): string {
|
||||
const scheme = window.location.protocol === "https:" ? "wss:" : "ws:";
|
||||
return `${scheme}//${window.location.host}${path}`;
|
||||
}
|
||||
|
||||
export function TerminalPage({
|
||||
agents,
|
||||
selectedAgentId,
|
||||
onSelectAgent,
|
||||
}: Props) {
|
||||
const hostRef = useRef<HTMLDivElement>(null);
|
||||
const terminalRef = useRef<XTerm | null>(null);
|
||||
const fitRef = useRef<FitAddon | null>(null);
|
||||
const socketRef = useRef<WebSocket | null>(null);
|
||||
const intentionalCloseRef = useRef(false);
|
||||
const lastTerminalSizeRef = useRef({ rows: 0, cols: 0 });
|
||||
const [session, setSession] = useState<TerminalSession | null>(null);
|
||||
const [connection, setConnection] = useState<ConnectionState>("idle");
|
||||
|
||||
const selectedAgent = agents.find(
|
||||
(agent) => agent.agent_id === selectedAgentId,
|
||||
);
|
||||
const canStart =
|
||||
selectedAgent?.connected &&
|
||||
selectedAgent.capabilities.includes("terminal") &&
|
||||
connection === "idle";
|
||||
|
||||
const sendResize = useCallback(() => {
|
||||
const terminal = terminalRef.current;
|
||||
const socket = socketRef.current;
|
||||
if (!terminal || socket?.readyState !== WebSocket.OPEN) return;
|
||||
if (
|
||||
lastTerminalSizeRef.current.rows === terminal.rows &&
|
||||
lastTerminalSizeRef.current.cols === terminal.cols
|
||||
) {
|
||||
return;
|
||||
}
|
||||
lastTerminalSizeRef.current = {
|
||||
rows: terminal.rows,
|
||||
cols: terminal.cols,
|
||||
};
|
||||
socket.send(
|
||||
JSON.stringify({
|
||||
type: "resize",
|
||||
rows: terminal.rows,
|
||||
cols: terminal.cols,
|
||||
}),
|
||||
);
|
||||
}, []);
|
||||
|
||||
const connect = useCallback(
|
||||
(nextSession: TerminalSession) => {
|
||||
if (!nextSession.attachment_token) {
|
||||
throw new Error(
|
||||
"Control plane did not issue a terminal attachment token",
|
||||
);
|
||||
}
|
||||
socketRef.current?.close();
|
||||
intentionalCloseRef.current = false;
|
||||
setConnection("connecting");
|
||||
const socket = new WebSocket(websocketUrl(nextSession.websocket_url));
|
||||
socket.binaryType = "arraybuffer";
|
||||
socketRef.current = socket;
|
||||
|
||||
socket.onopen = () => {
|
||||
socket.send(
|
||||
JSON.stringify({
|
||||
type: "attach",
|
||||
attachment_token: nextSession.attachment_token,
|
||||
}),
|
||||
);
|
||||
fitRef.current?.fit();
|
||||
sendResize();
|
||||
};
|
||||
socket.onmessage = (event) => {
|
||||
if (typeof event.data !== "string") {
|
||||
terminalRef.current?.write(new Uint8Array(event.data as ArrayBuffer));
|
||||
return;
|
||||
}
|
||||
try {
|
||||
const control = JSON.parse(event.data) as {
|
||||
type: string;
|
||||
exit_code?: number | null;
|
||||
message?: string;
|
||||
};
|
||||
if (control.type === "ready") {
|
||||
setConnection("ready");
|
||||
window.requestAnimationFrame(() => {
|
||||
fitRef.current?.fit();
|
||||
lastTerminalSizeRef.current = { rows: 0, cols: 0 };
|
||||
sendResize();
|
||||
terminalRef.current?.focus();
|
||||
});
|
||||
} else if (control.type === "exited") {
|
||||
setConnection("exited");
|
||||
terminalRef.current?.writeln(
|
||||
`\r\n[process exited${control.exit_code == null ? "" : ` ${control.exit_code}`}]`,
|
||||
);
|
||||
} else if (control.type === "error") {
|
||||
setConnection("exited");
|
||||
terminalRef.current?.writeln(
|
||||
`\r\n[terminal error: ${control.message ?? "unknown error"}]`,
|
||||
);
|
||||
}
|
||||
} catch {
|
||||
terminalRef.current?.writeln("\r\n[invalid terminal control frame]");
|
||||
}
|
||||
};
|
||||
socket.onerror = () => {
|
||||
terminalRef.current?.writeln("\r\n[terminal transport error]");
|
||||
};
|
||||
socket.onclose = () => {
|
||||
socketRef.current = null;
|
||||
setConnection((current) => {
|
||||
if (intentionalCloseRef.current || current === "exited")
|
||||
return current;
|
||||
return "disconnected";
|
||||
});
|
||||
};
|
||||
},
|
||||
[sendResize],
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
if (!hostRef.current) return;
|
||||
const terminal = new XTerm({
|
||||
cursorBlink: true,
|
||||
convertEol: false,
|
||||
fontFamily: '"SFMono-Regular", Consolas, "Liberation Mono", monospace',
|
||||
fontSize: 14,
|
||||
lineHeight: 1.1,
|
||||
scrollback: 5000,
|
||||
theme: {
|
||||
background: "#0b1117",
|
||||
foreground: "#e5e7eb",
|
||||
cursor: "#6ee7a8",
|
||||
},
|
||||
});
|
||||
const fit = new FitAddon();
|
||||
terminal.loadAddon(fit);
|
||||
terminal.open(hostRef.current);
|
||||
fit.fit();
|
||||
terminalRef.current = terminal;
|
||||
fitRef.current = fit;
|
||||
|
||||
let cancelled = false;
|
||||
void document.fonts.ready.then(() => {
|
||||
if (cancelled) return;
|
||||
fit.fit();
|
||||
lastTerminalSizeRef.current = { rows: 0, cols: 0 };
|
||||
sendResize();
|
||||
});
|
||||
|
||||
const input = terminal.onData((data) => {
|
||||
const socket = socketRef.current;
|
||||
if (socket?.readyState === WebSocket.OPEN) {
|
||||
socket.send(new TextEncoder().encode(data));
|
||||
}
|
||||
});
|
||||
let resizeFrame = 0;
|
||||
let lastHostWidth = 0;
|
||||
let lastHostHeight = 0;
|
||||
const resizeObserver = new ResizeObserver(([entry]) => {
|
||||
const width = Math.round(entry.contentRect.width);
|
||||
const height = Math.round(entry.contentRect.height);
|
||||
if (width === lastHostWidth && height === lastHostHeight) return;
|
||||
lastHostWidth = width;
|
||||
lastHostHeight = height;
|
||||
window.cancelAnimationFrame(resizeFrame);
|
||||
resizeFrame = window.requestAnimationFrame(() => {
|
||||
fit.fit();
|
||||
sendResize();
|
||||
});
|
||||
});
|
||||
resizeObserver.observe(hostRef.current);
|
||||
|
||||
return () => {
|
||||
cancelled = true;
|
||||
input.dispose();
|
||||
resizeObserver.disconnect();
|
||||
window.cancelAnimationFrame(resizeFrame);
|
||||
socketRef.current?.close();
|
||||
terminal.dispose();
|
||||
terminalRef.current = null;
|
||||
fitRef.current = null;
|
||||
};
|
||||
}, [sendResize]);
|
||||
|
||||
async function start() {
|
||||
if (!selectedAgentId || !terminalRef.current) return;
|
||||
setConnection("connecting");
|
||||
terminalRef.current.clear();
|
||||
try {
|
||||
fitRef.current?.fit();
|
||||
const created = await createTerminal(
|
||||
selectedAgentId,
|
||||
terminalRef.current.rows,
|
||||
terminalRef.current.cols,
|
||||
);
|
||||
setSession(created);
|
||||
connect(created);
|
||||
} catch (error) {
|
||||
setConnection("idle");
|
||||
toast.error("Could not start terminal", { description: String(error) });
|
||||
}
|
||||
}
|
||||
|
||||
async function reconnect() {
|
||||
if (!session) return;
|
||||
try {
|
||||
const attached = await attachTerminal(session.terminal_id);
|
||||
setSession(attached);
|
||||
connect(attached);
|
||||
} catch (error) {
|
||||
setConnection("exited");
|
||||
toast.error("Terminal session is no longer available", {
|
||||
description: String(error),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
async function close() {
|
||||
if (!session) return;
|
||||
intentionalCloseRef.current = true;
|
||||
socketRef.current?.send(JSON.stringify({ type: "close" }));
|
||||
socketRef.current?.close();
|
||||
try {
|
||||
await closeTerminal(session.terminal_id);
|
||||
} catch (error) {
|
||||
toast.error("Terminal cleanup failed", { description: String(error) });
|
||||
} finally {
|
||||
setSession(null);
|
||||
setConnection("idle");
|
||||
terminalRef.current?.clear();
|
||||
}
|
||||
}
|
||||
|
||||
return (
|
||||
<section className="terminal-page" aria-label="Remote terminal">
|
||||
<header className="terminal-toolbar">
|
||||
<div className="terminal-heading">
|
||||
<Terminal className="size-5 text-primary" aria-hidden />
|
||||
<h1>Remote Terminal</h1>
|
||||
</div>
|
||||
|
||||
<div className="terminal-controls">
|
||||
<AgentSelector
|
||||
agents={agents.filter(
|
||||
(agent) =>
|
||||
agent.connected && agent.capabilities.includes("terminal"),
|
||||
)}
|
||||
value={selectedAgentId}
|
||||
onChange={onSelectAgent}
|
||||
disabled={connection !== "idle"}
|
||||
className="w-full min-w-0 sm:w-64"
|
||||
/>
|
||||
{connection === "idle" ? (
|
||||
<Button type="button" onClick={start} disabled={!canStart}>
|
||||
<PlugZap className="size-4" aria-hidden />
|
||||
Connect
|
||||
</Button>
|
||||
) : null}
|
||||
</div>
|
||||
</header>
|
||||
|
||||
{!selectedAgent && connection === "idle" ? (
|
||||
<p className="terminal-notice">
|
||||
No connected terminal-capable agent is selected.
|
||||
</p>
|
||||
) : selectedAgent && !selectedAgent.capabilities.includes("terminal") ? (
|
||||
<p className="terminal-notice">
|
||||
This agent has not enabled remote terminal capability.
|
||||
</p>
|
||||
) : null}
|
||||
|
||||
<div className="terminal-frame">
|
||||
<div className="terminal-framebar">
|
||||
<div className="terminal-session-label">
|
||||
<Terminal className="size-4" aria-hidden />
|
||||
<span>
|
||||
{session && selectedAgent
|
||||
? displayAgentLabel(selectedAgent)
|
||||
: "No active session"}
|
||||
</span>
|
||||
<Badge
|
||||
variant="outline"
|
||||
className="terminal-status"
|
||||
data-state={connection}
|
||||
>
|
||||
<span aria-hidden />
|
||||
{connection}
|
||||
</Badge>
|
||||
</div>
|
||||
|
||||
<div className="terminal-frame-actions">
|
||||
{connection === "disconnected" ? (
|
||||
<Button
|
||||
type="button"
|
||||
size="sm"
|
||||
variant="ghost"
|
||||
onClick={reconnect}
|
||||
>
|
||||
<RotateCcw className="size-4" aria-hidden />
|
||||
Reconnect
|
||||
</Button>
|
||||
) : null}
|
||||
<Tooltip>
|
||||
<TooltipTrigger>
|
||||
<Button
|
||||
type="button"
|
||||
size="icon"
|
||||
variant="ghost"
|
||||
disabled={connection === "idle"}
|
||||
aria-label="Clear terminal"
|
||||
onClick={() => terminalRef.current?.clear()}
|
||||
>
|
||||
<Eraser className="size-4" aria-hidden />
|
||||
</Button>
|
||||
</TooltipTrigger>
|
||||
<TooltipContent>Clear terminal</TooltipContent>
|
||||
</Tooltip>
|
||||
{session ? (
|
||||
<Tooltip>
|
||||
<TooltipTrigger>
|
||||
<Button
|
||||
type="button"
|
||||
size="icon"
|
||||
variant="ghost"
|
||||
className="text-destructive hover:text-destructive"
|
||||
aria-label="Close terminal session"
|
||||
onClick={close}
|
||||
>
|
||||
<Square className="size-3.5 fill-current" aria-hidden />
|
||||
</Button>
|
||||
</TooltipTrigger>
|
||||
<TooltipContent>Close session</TooltipContent>
|
||||
</Tooltip>
|
||||
) : null}
|
||||
</div>
|
||||
</div>
|
||||
<div className="terminal-surface" ref={hostRef} />
|
||||
</div>
|
||||
</section>
|
||||
);
|
||||
}
|
||||
@@ -192,6 +192,12 @@
|
||||
transition: background 100ms ease-in-out;
|
||||
}
|
||||
|
||||
/* xterm focuses an input near the active row. Keep that focus operation from
|
||||
scrolling the page shell and moving the terminal toolbar off-screen. */
|
||||
.app-main:has(.terminal-page) {
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
/* ── Sidebar internals ────────────────────────────────────────────────── */
|
||||
|
||||
.sidebar-brand {
|
||||
@@ -316,6 +322,172 @@
|
||||
padding: 0.5rem;
|
||||
}
|
||||
|
||||
/* ── Remote terminal ─────────────────────────────────────────────────── */
|
||||
|
||||
.terminal-page {
|
||||
display: flex;
|
||||
height: calc(100dvh - 2.5rem);
|
||||
min-height: 0;
|
||||
flex-direction: column;
|
||||
gap: 1rem;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.terminal-toolbar {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: space-between;
|
||||
gap: 1rem;
|
||||
flex-wrap: wrap;
|
||||
}
|
||||
|
||||
.terminal-heading {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 0.65rem;
|
||||
}
|
||||
|
||||
.terminal-heading h1 {
|
||||
margin: 0;
|
||||
font-size: 1.0625rem;
|
||||
font-weight: 650;
|
||||
}
|
||||
|
||||
.terminal-controls {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: flex-end;
|
||||
gap: 0.5rem;
|
||||
flex-wrap: wrap;
|
||||
}
|
||||
|
||||
.terminal-status {
|
||||
gap: 0.4rem;
|
||||
height: 1.4rem;
|
||||
padding-inline: 0.45rem;
|
||||
font-size: 0.6875rem;
|
||||
text-transform: capitalize;
|
||||
}
|
||||
|
||||
.terminal-status > span {
|
||||
width: 0.45rem;
|
||||
height: 0.45rem;
|
||||
border-radius: 999px;
|
||||
background: var(--presence-unknown);
|
||||
}
|
||||
|
||||
.terminal-status[data-state="ready"] > span {
|
||||
background: var(--presence-online);
|
||||
}
|
||||
|
||||
.terminal-status[data-state="connecting"] > span {
|
||||
background: var(--presence-likely);
|
||||
}
|
||||
|
||||
.terminal-status[data-state="disconnected"] > span,
|
||||
.terminal-status[data-state="exited"] > span {
|
||||
background: var(--presence-offline);
|
||||
}
|
||||
|
||||
.terminal-notice {
|
||||
margin: 0;
|
||||
color: var(--muted-foreground);
|
||||
font-size: 0.8125rem;
|
||||
}
|
||||
|
||||
.terminal-frame {
|
||||
display: flex;
|
||||
flex: 1;
|
||||
min-height: 0;
|
||||
flex-direction: column;
|
||||
overflow: hidden;
|
||||
border: 1px solid var(--border);
|
||||
border-radius: var(--radius);
|
||||
background: #0b1117;
|
||||
}
|
||||
|
||||
.terminal-framebar {
|
||||
display: flex;
|
||||
min-height: 2.5rem;
|
||||
flex-shrink: 0;
|
||||
align-items: center;
|
||||
justify-content: space-between;
|
||||
gap: 0.75rem;
|
||||
padding: 0.35rem 0.45rem 0.35rem 0.75rem;
|
||||
border-bottom: 1px solid color-mix(in oklab, white 9%, transparent);
|
||||
background: #101820;
|
||||
color: #cbd5e1;
|
||||
}
|
||||
|
||||
.terminal-session-label,
|
||||
.terminal-frame-actions {
|
||||
display: flex;
|
||||
min-width: 0;
|
||||
align-items: center;
|
||||
gap: 0.5rem;
|
||||
}
|
||||
|
||||
.terminal-session-label > span:not(.terminal-status) {
|
||||
overflow: hidden;
|
||||
color: #dbe4ee;
|
||||
font-size: 0.75rem;
|
||||
font-weight: 550;
|
||||
text-overflow: ellipsis;
|
||||
white-space: nowrap;
|
||||
}
|
||||
|
||||
.terminal-framebar button {
|
||||
color: #aab7c4;
|
||||
}
|
||||
|
||||
.terminal-framebar button:hover {
|
||||
color: #f1f5f9;
|
||||
background: rgb(255 255 255 / 7%);
|
||||
}
|
||||
|
||||
.terminal-surface {
|
||||
flex: 1;
|
||||
min-height: 0;
|
||||
overflow: hidden;
|
||||
background: #0b1117;
|
||||
padding: 0;
|
||||
}
|
||||
|
||||
.terminal-surface .xterm {
|
||||
height: 100%;
|
||||
box-sizing: border-box;
|
||||
background: #0b1117;
|
||||
padding: 0.75rem;
|
||||
}
|
||||
|
||||
.terminal-surface .xterm-viewport {
|
||||
/* xterm defaults this element to pure black. It remains visible around the
|
||||
fitted canvas, so keep it on the same surface as the terminal theme. */
|
||||
background-color: #0b1117;
|
||||
overscroll-behavior: contain;
|
||||
scrollbar-color: color-mix(in oklab, white 28%, transparent) transparent;
|
||||
scrollbar-width: thin;
|
||||
}
|
||||
|
||||
@media (max-width: 640px) {
|
||||
.terminal-page {
|
||||
height: calc(100dvh - 2.5rem);
|
||||
}
|
||||
|
||||
.terminal-controls {
|
||||
width: 100%;
|
||||
justify-content: flex-start;
|
||||
}
|
||||
|
||||
.terminal-surface {
|
||||
min-height: 0;
|
||||
}
|
||||
|
||||
.terminal-surface .xterm {
|
||||
padding: 0.5rem;
|
||||
}
|
||||
}
|
||||
|
||||
/* ── Presence badges ──────────────────────────────────────────────────── */
|
||||
|
||||
.presence-dot {
|
||||
|
||||
@@ -23,6 +23,7 @@ serde_json = "1"
|
||||
tokio = { version = "1", features = [
|
||||
"fs",
|
||||
"macros",
|
||||
"io-util",
|
||||
"rt-multi-thread",
|
||||
"time",
|
||||
"signal",
|
||||
|
||||
@@ -35,6 +35,28 @@ pub struct AgentConfig {
|
||||
pub mac_name_cache_path: PathBuf,
|
||||
#[serde(default = "default_observation_store_path")]
|
||||
pub observation_store_path: PathBuf,
|
||||
#[serde(default)]
|
||||
pub terminal: TerminalConfig,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct TerminalConfig {
|
||||
#[serde(default)]
|
||||
pub enabled: bool,
|
||||
#[serde(default = "default_terminal_shell")]
|
||||
pub shell: PathBuf,
|
||||
#[serde(default = "default_terminal_max_sessions")]
|
||||
pub max_sessions: usize,
|
||||
}
|
||||
|
||||
impl Default for TerminalConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: false,
|
||||
shell: default_terminal_shell(),
|
||||
max_sessions: default_terminal_max_sessions(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub static DEFAULT_CONFIG: LazyLock<AgentConfig> = LazyLock::new(|| AgentConfig {
|
||||
@@ -49,6 +71,7 @@ pub static DEFAULT_CONFIG: LazyLock<AgentConfig> = LazyLock::new(|| AgentConfig
|
||||
dhcp_leases_path: default_dhcp_leases_path(),
|
||||
mac_name_cache_path: default_mac_name_cache_path(),
|
||||
observation_store_path: default_observation_store_path(),
|
||||
terminal: TerminalConfig::default(),
|
||||
});
|
||||
|
||||
impl fmt::Debug for AgentConfig {
|
||||
@@ -71,6 +94,7 @@ impl fmt::Debug for AgentConfig {
|
||||
.field("dhcp_leases_path", &self.dhcp_leases_path)
|
||||
.field("mac_name_cache_path", &self.mac_name_cache_path)
|
||||
.field("observation_store_path", &self.observation_store_path)
|
||||
.field("terminal", &self.terminal)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -107,6 +131,14 @@ fn default_observation_store_path() -> PathBuf {
|
||||
DEFAULT_OBSERVATION_STORE_PATH.into()
|
||||
}
|
||||
|
||||
fn default_terminal_shell() -> PathBuf {
|
||||
"/bin/ash".into()
|
||||
}
|
||||
|
||||
const fn default_terminal_max_sessions() -> usize {
|
||||
2
|
||||
}
|
||||
|
||||
impl AgentConfig {
|
||||
pub fn local_path_envs(&self) -> Vec<(&'static str, &Path)> {
|
||||
vec![
|
||||
@@ -214,6 +246,11 @@ mod tests {
|
||||
dhcp_leases_path: "/tmp/test-dhcp.leases".into(),
|
||||
mac_name_cache_path: "/tmp/test-names.json".into(),
|
||||
observation_store_path: "/tmp/test-observations.json".into(),
|
||||
terminal: TerminalConfig {
|
||||
enabled: true,
|
||||
shell: "/bin/sh".into(),
|
||||
max_sessions: 2,
|
||||
},
|
||||
};
|
||||
|
||||
save_config(&path, &config).expect("save");
|
||||
@@ -245,5 +282,6 @@ agent_token = "secret"
|
||||
config.observation_retention_days,
|
||||
DEFAULT_OBSERVATION_RETENTION_DAYS
|
||||
);
|
||||
assert_eq!(config.terminal, TerminalConfig::default());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -151,6 +151,11 @@ mod tests {
|
||||
dhcp_leases_path: "/tmp/custom-dhcp.leases".into(),
|
||||
mac_name_cache_path: "/tmp/custom-names.json".into(),
|
||||
observation_store_path: "/tmp/custom-observations.json".into(),
|
||||
terminal: crate::config::TerminalConfig {
|
||||
enabled: true,
|
||||
shell: "/bin/ash".into(),
|
||||
max_sessions: 2,
|
||||
},
|
||||
};
|
||||
|
||||
let outcome = enroll(&server_url, "enroll-abc", &path, Some(&base_config))
|
||||
@@ -164,6 +169,7 @@ mod tests {
|
||||
assert_eq!(config.observation_retention_days, 11);
|
||||
assert_eq!(config.pid_file, base_config.pid_file);
|
||||
assert_eq!(config.dhcp_leases_path, base_config.dhcp_leases_path);
|
||||
assert_eq!(config.terminal, base_config.terminal);
|
||||
assert!(outcome.backup_path.is_none());
|
||||
|
||||
let persisted = crate::config::load_config(&path).expect("load persisted config");
|
||||
|
||||
@@ -2,14 +2,15 @@ mod cli;
|
||||
mod config;
|
||||
mod dispatch;
|
||||
mod enroll;
|
||||
mod protocol;
|
||||
mod serve;
|
||||
mod session;
|
||||
mod terminal;
|
||||
mod tracing;
|
||||
|
||||
use anyhow::Result;
|
||||
use clap::Parser;
|
||||
use cli::{Cli, Command, InitConfigArgs, ObserveCommand};
|
||||
use wakey_agent::protocol;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
|
||||
@@ -17,6 +17,61 @@ impl RequestId {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
#[serde(transparent)]
|
||||
pub struct TerminalId(String);
|
||||
|
||||
impl TerminalId {
|
||||
pub fn new(value: impl Into<String>) -> Result<Self, String> {
|
||||
let value = value.into();
|
||||
if value.trim().is_empty() {
|
||||
return Err("terminal_id must not be empty".into());
|
||||
}
|
||||
Ok(Self(value))
|
||||
}
|
||||
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for TerminalId {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
self.0.fmt(f)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum AgentCapability {
|
||||
Terminal,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum TerminalControl {
|
||||
Resize { rows: u16, cols: u16 },
|
||||
Ready,
|
||||
Exited { exit_code: Option<i32> },
|
||||
Error { code: String, message: String },
|
||||
Close,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum TerminalAgentHandshake {
|
||||
Auth {
|
||||
agent_id: String,
|
||||
relay_token: String,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum TerminalOperatorHandshake {
|
||||
Attach { attachment_token: String },
|
||||
}
|
||||
|
||||
impl TryFrom<String> for RequestId {
|
||||
type Error = String;
|
||||
|
||||
@@ -153,6 +208,8 @@ pub enum CommandResult {
|
||||
pub enum ClientMessage {
|
||||
Hello {
|
||||
agent_id: String,
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
capabilities: Vec<AgentCapability>,
|
||||
},
|
||||
Auth {
|
||||
agent_id: String,
|
||||
@@ -173,6 +230,10 @@ pub enum ClientMessage {
|
||||
request_id: RequestId,
|
||||
error: ErrorPayload,
|
||||
},
|
||||
TerminalRejected {
|
||||
terminal_id: TerminalId,
|
||||
error: ErrorPayload,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
@@ -183,6 +244,15 @@ pub enum ServerMessage {
|
||||
command: AgentCommand,
|
||||
},
|
||||
SyncDeviceSnapshot,
|
||||
OpenTerminal {
|
||||
terminal_id: TerminalId,
|
||||
relay_token: String,
|
||||
rows: u16,
|
||||
cols: u16,
|
||||
},
|
||||
CloseTerminal {
|
||||
terminal_id: TerminalId,
|
||||
},
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -232,4 +302,24 @@ mod tests {
|
||||
assert!(json.contains("\"type\":\"device_snapshot\""));
|
||||
assert!(json.contains("\"devices\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn terminal_messages_serialize_with_explicit_frame_types() {
|
||||
let message = ServerMessage::OpenTerminal {
|
||||
terminal_id: TerminalId::new("term-1").expect("terminal id"),
|
||||
relay_token: "secret".into(),
|
||||
rows: 30,
|
||||
cols: 120,
|
||||
};
|
||||
let json = serde_json::to_string(&message).expect("serialize open terminal");
|
||||
assert!(json.contains("\"type\":\"open_terminal\""));
|
||||
assert!(json.contains("\"terminal_id\":\"term-1\""));
|
||||
|
||||
let resize = TerminalControl::Resize {
|
||||
rows: 40,
|
||||
cols: 160,
|
||||
};
|
||||
let json = serde_json::to_string(&resize).expect("serialize resize");
|
||||
assert_eq!(json, r#"{"type":"resize","rows":40,"cols":160}"#);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,7 +8,8 @@ use tracing::{debug, error, info, info_span, warn};
|
||||
|
||||
use crate::config::AgentConfig;
|
||||
use crate::dispatch::{dispatch_command, inventory_for_config};
|
||||
use crate::protocol::{AgentCommand, ClientMessage, ErrorPayload, ServerMessage};
|
||||
use crate::protocol::{AgentCapability, AgentCommand, ClientMessage, ErrorPayload, ServerMessage};
|
||||
use crate::terminal::TerminalManager;
|
||||
|
||||
pub async fn run(config: AgentConfig) -> Result<()> {
|
||||
let mut backoff = config.reconnect_base_ms.max(100);
|
||||
@@ -71,6 +72,7 @@ async fn run_once(config: &AgentConfig) -> Result<()> {
|
||||
&mut sink,
|
||||
&ClientMessage::Hello {
|
||||
agent_id: config.agent_id.clone(),
|
||||
capabilities: agent_capabilities(config),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
@@ -91,9 +93,23 @@ async fn run_once(config: &AgentConfig) -> Result<()> {
|
||||
));
|
||||
snapshot_sync.set_missed_tick_behavior(MissedTickBehavior::Skip);
|
||||
snapshot_sync.reset();
|
||||
let (terminal_manager, mut terminal_events) = TerminalManager::new(config);
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
Some(event) = terminal_events.recv() => {
|
||||
send_json(
|
||||
&mut sink,
|
||||
&ClientMessage::TerminalRejected {
|
||||
terminal_id: event.terminal_id,
|
||||
error: ErrorPayload {
|
||||
code: "terminal_worker_failed".into(),
|
||||
message: event.error,
|
||||
retryable: Some(false),
|
||||
},
|
||||
},
|
||||
).await?;
|
||||
}
|
||||
_ = heartbeat.tick() => {
|
||||
send_json(&mut sink, &ClientMessage::Heartbeat {
|
||||
agent_id: config.agent_id.clone(),
|
||||
@@ -118,7 +134,7 @@ async fn run_once(config: &AgentConfig) -> Result<()> {
|
||||
Message::Text(text) => {
|
||||
match serde_json::from_str::<ServerMessage>(&text) {
|
||||
Ok(message) => {
|
||||
handle_server_message(config, &mut sink, &mut snapshot_sync, message).await?;
|
||||
handle_server_message(config, &terminal_manager, &mut sink, &mut snapshot_sync, message).await?;
|
||||
}
|
||||
Err(err) => {
|
||||
// Allow the server to introduce extra frame types without
|
||||
@@ -181,6 +197,7 @@ pub fn next_backoff_ms(current_ms: u64, max_ms: u64) -> u64 {
|
||||
|
||||
async fn handle_server_message<S>(
|
||||
config: &AgentConfig,
|
||||
terminal_manager: &TerminalManager,
|
||||
sink: &mut S,
|
||||
snapshot_sync: &mut tokio::time::Interval,
|
||||
message: ServerMessage,
|
||||
@@ -223,6 +240,35 @@ where
|
||||
send_device_snapshot_ws(sink, config).await?;
|
||||
snapshot_sync.reset();
|
||||
}
|
||||
ServerMessage::OpenTerminal {
|
||||
terminal_id,
|
||||
relay_token,
|
||||
rows,
|
||||
cols,
|
||||
} => {
|
||||
info!(terminal_id = %terminal_id, rows, cols, "received terminal open request");
|
||||
if let Err(err) =
|
||||
terminal_manager.open(config, terminal_id.clone(), relay_token, rows, cols)
|
||||
{
|
||||
error!(terminal_id = %terminal_id, error = %err, "terminal open request rejected");
|
||||
send_json(
|
||||
sink,
|
||||
&ClientMessage::TerminalRejected {
|
||||
terminal_id,
|
||||
error: ErrorPayload {
|
||||
code: "terminal_open_rejected".into(),
|
||||
message: err.to_string(),
|
||||
retryable: Some(false),
|
||||
},
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
ServerMessage::CloseTerminal { terminal_id } => {
|
||||
info!(terminal_id = %terminal_id, "received terminal close request");
|
||||
terminal_manager.close(&terminal_id);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -301,9 +347,19 @@ fn client_message_kind(message: &ClientMessage) -> &'static str {
|
||||
ClientMessage::DeviceSnapshot { .. } => "device_snapshot",
|
||||
ClientMessage::Result { .. } => "result",
|
||||
ClientMessage::Error { .. } => "error",
|
||||
ClientMessage::TerminalRejected { .. } => "terminal_rejected",
|
||||
}
|
||||
}
|
||||
|
||||
fn agent_capabilities(config: &AgentConfig) -> Vec<AgentCapability> {
|
||||
#[cfg(unix)]
|
||||
if config.terminal.enabled {
|
||||
return vec![AgentCapability::Terminal];
|
||||
}
|
||||
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
pub fn websocket_url(server_url: &str) -> Result<url::Url> {
|
||||
let base = url::Url::parse(server_url).context("invalid server_url")?;
|
||||
let scheme = match base.scheme() {
|
||||
|
||||
@@ -0,0 +1,345 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::{Arc, Mutex, Weak};
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::config::AgentConfig;
|
||||
use crate::protocol::{TerminalAgentHandshake, TerminalControl, TerminalId};
|
||||
use anyhow::{Context, Result};
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use tracing::{info, warn};
|
||||
|
||||
const MAX_TERMINAL_FRAME_BYTES: usize = 64 * 1024;
|
||||
const PROCESS_SIGNAL_GRACE: Duration = Duration::from_secs(1);
|
||||
|
||||
/// Owns cancellation handles for terminal workers started by the control socket.
|
||||
pub struct TerminalManager {
|
||||
active: Arc<Mutex<HashMap<String, oneshot::Sender<()>>>>,
|
||||
max_sessions: usize,
|
||||
events: mpsc::UnboundedSender<TerminalManagerEvent>,
|
||||
}
|
||||
|
||||
pub struct TerminalManagerEvent {
|
||||
pub terminal_id: TerminalId,
|
||||
pub error: String,
|
||||
}
|
||||
|
||||
impl TerminalManager {
|
||||
pub fn new(config: &AgentConfig) -> (Self, mpsc::UnboundedReceiver<TerminalManagerEvent>) {
|
||||
let (events, event_rx) = mpsc::unbounded_channel();
|
||||
(
|
||||
Self {
|
||||
active: Arc::new(Mutex::new(HashMap::new())),
|
||||
max_sessions: config.terminal.max_sessions.max(1),
|
||||
events,
|
||||
},
|
||||
event_rx,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn open(
|
||||
&self,
|
||||
config: &AgentConfig,
|
||||
terminal_id: TerminalId,
|
||||
relay_token: String,
|
||||
rows: u16,
|
||||
cols: u16,
|
||||
) -> Result<()> {
|
||||
if !config.terminal.enabled {
|
||||
anyhow::bail!("terminal capability is disabled");
|
||||
}
|
||||
|
||||
let terminal_key = terminal_id.to_string();
|
||||
let (cancel_tx, cancel_rx) = oneshot::channel();
|
||||
{
|
||||
let mut active = self.active.lock().expect("terminal manager poisoned");
|
||||
if active.contains_key(&terminal_key) {
|
||||
anyhow::bail!("terminal session {terminal_key} is already active");
|
||||
}
|
||||
if active.len() >= self.max_sessions {
|
||||
anyhow::bail!("agent terminal session limit reached");
|
||||
}
|
||||
active.insert(terminal_key.clone(), cancel_tx);
|
||||
}
|
||||
|
||||
let config = config.clone();
|
||||
let active = Arc::downgrade(&self.active);
|
||||
let events = self.events.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(err) =
|
||||
run_terminal(&config, &terminal_id, &relay_token, rows, cols, cancel_rx).await
|
||||
{
|
||||
warn!(terminal_id = %terminal_id, error = %err, "terminal worker failed");
|
||||
let _ = events.send(TerminalManagerEvent {
|
||||
terminal_id: terminal_id.clone(),
|
||||
error: err.to_string(),
|
||||
});
|
||||
}
|
||||
remove_completed(&active, terminal_id.as_str());
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn close(&self, terminal_id: &TerminalId) -> bool {
|
||||
self.active
|
||||
.lock()
|
||||
.expect("terminal manager poisoned")
|
||||
.remove(terminal_id.as_str())
|
||||
.is_some_and(|cancel| cancel.send(()).is_ok())
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for TerminalManager {
|
||||
fn drop(&mut self) {
|
||||
if Arc::strong_count(&self.active) == 1
|
||||
&& let Ok(mut active) = self.active.lock()
|
||||
{
|
||||
for (_, cancel) in active.drain() {
|
||||
let _ = cancel.send(());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn remove_completed(active: &Weak<Mutex<HashMap<String, oneshot::Sender<()>>>>, terminal_id: &str) {
|
||||
if let Some(active) = active.upgrade()
|
||||
&& let Ok(mut active) = active.lock()
|
||||
{
|
||||
active.remove(terminal_id);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
async fn run_terminal(
|
||||
config: &AgentConfig,
|
||||
terminal_id: &TerminalId,
|
||||
relay_token: &str,
|
||||
rows: u16,
|
||||
cols: u16,
|
||||
mut cancel: oneshot::Receiver<()>,
|
||||
) -> Result<()> {
|
||||
let ws_url = terminal_websocket_url(&config.server_url, terminal_id)?;
|
||||
let (stream, _) = tokio::select! {
|
||||
_ = &mut cancel => return Ok(()),
|
||||
result = tokio_tungstenite::connect_async(ws_url.as_str()) => {
|
||||
result.context("failed to connect terminal relay websocket")?
|
||||
}
|
||||
};
|
||||
let (mut sink, mut source) = stream.split();
|
||||
send_json(
|
||||
&mut sink,
|
||||
&TerminalAgentHandshake::Auth {
|
||||
agent_id: config.agent_id.clone(),
|
||||
relay_token: relay_token.to_string(),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
|
||||
let terminal = match wakey::wakey_linux::terminal::TerminalPty::spawn(
|
||||
Path::new(&config.terminal.shell),
|
||||
rows,
|
||||
cols,
|
||||
) {
|
||||
Ok(terminal) => terminal,
|
||||
Err(err) => {
|
||||
let _ = send_json(
|
||||
&mut sink,
|
||||
&TerminalControl::Error {
|
||||
code: "terminal_spawn_failed".into(),
|
||||
message: err.to_string(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let _ = sink.send(Message::Close(None)).await;
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
let wakey::wakey_linux::terminal::TerminalPty {
|
||||
mut reader,
|
||||
mut writer,
|
||||
mut child,
|
||||
} = terminal;
|
||||
let process_group = child.id();
|
||||
send_json(&mut sink, &TerminalControl::Ready).await?;
|
||||
info!(terminal_id = %terminal_id, shell = %config.terminal.shell.display(), "terminal PTY ready");
|
||||
|
||||
let mut output = [0_u8; 16 * 1024];
|
||||
let mut requested_close = false;
|
||||
let mut observed_status = None;
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = &mut cancel => {
|
||||
requested_close = true;
|
||||
break;
|
||||
}
|
||||
status = child.wait() => {
|
||||
observed_status = Some(status.context("failed waiting for terminal child")?);
|
||||
break;
|
||||
}
|
||||
read = reader.read(&mut output) => {
|
||||
match read {
|
||||
Ok(0) => break,
|
||||
Ok(count) => sink
|
||||
.send(Message::Binary(output[..count].to_vec().into()))
|
||||
.await
|
||||
.context("failed to send PTY output")?,
|
||||
// Linux PTY masters commonly report EIO after the slave closes.
|
||||
Err(err) if err.raw_os_error() == Some(5) => break,
|
||||
Err(err) => return Err(err).context("failed to read PTY output"),
|
||||
}
|
||||
}
|
||||
incoming = source.next() => {
|
||||
let Some(message) = incoming else { break; };
|
||||
match message.context("terminal relay websocket receive failed")? {
|
||||
Message::Binary(bytes) => {
|
||||
if bytes.len() > MAX_TERMINAL_FRAME_BYTES {
|
||||
anyhow::bail!("terminal input frame exceeds size limit");
|
||||
}
|
||||
writer.write_all(&bytes).await.context("failed to write PTY input")?;
|
||||
}
|
||||
Message::Text(text) => {
|
||||
match serde_json::from_str::<TerminalControl>(&text)
|
||||
.context("invalid terminal control frame")?
|
||||
{
|
||||
TerminalControl::Resize { rows, cols } => {
|
||||
validate_size(rows, cols)?;
|
||||
wakey::wakey_linux::terminal::resize_terminal(
|
||||
&writer, rows, cols,
|
||||
)?;
|
||||
}
|
||||
TerminalControl::Close => {
|
||||
requested_close = true;
|
||||
break;
|
||||
}
|
||||
_ => anyhow::bail!("terminal control frame has invalid direction"),
|
||||
}
|
||||
}
|
||||
Message::Ping(payload) => sink.send(Message::Pong(payload)).await?,
|
||||
Message::Pong(_) => {}
|
||||
Message::Close(_) => {
|
||||
requested_close = true;
|
||||
break;
|
||||
}
|
||||
Message::Frame(_) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let status = match observed_status {
|
||||
Some(status) => status,
|
||||
None => terminate_process_group(&mut child, process_group).await?,
|
||||
};
|
||||
let _ = send_json(
|
||||
&mut sink,
|
||||
&TerminalControl::Exited {
|
||||
exit_code: status.code(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let _ = sink.send(Message::Close(None)).await;
|
||||
info!(terminal_id = %terminal_id, exit_code = ?status.code(), requested_close, "terminal worker exited");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
async fn run_terminal(
|
||||
_config: &AgentConfig,
|
||||
_terminal_id: &TerminalId,
|
||||
_relay_token: &str,
|
||||
_rows: u16,
|
||||
_cols: u16,
|
||||
_cancel: oneshot::Receiver<()>,
|
||||
) -> Result<()> {
|
||||
anyhow::bail!("terminal sessions are unsupported on this platform")
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
async fn terminate_process_group(
|
||||
child: &mut tokio::process::Child,
|
||||
process_group: Option<u32>,
|
||||
) -> Result<std::process::ExitStatus> {
|
||||
use nix::sys::signal::{Signal, killpg};
|
||||
use nix::unistd::Pid;
|
||||
|
||||
let Some(process_group) = process_group else {
|
||||
child
|
||||
.start_kill()
|
||||
.context("failed to kill terminal child")?;
|
||||
return child
|
||||
.wait()
|
||||
.await
|
||||
.context("failed waiting for terminal child");
|
||||
};
|
||||
let pid = Pid::from_raw(process_group as i32);
|
||||
let _ = killpg(pid, Signal::SIGHUP);
|
||||
if let Ok(status) = tokio::time::timeout(PROCESS_SIGNAL_GRACE, child.wait()).await {
|
||||
return status.context("failed waiting for terminal child after SIGHUP");
|
||||
}
|
||||
let _ = killpg(pid, Signal::SIGTERM);
|
||||
if let Ok(status) = tokio::time::timeout(PROCESS_SIGNAL_GRACE, child.wait()).await {
|
||||
return status.context("failed waiting for terminal child after SIGTERM");
|
||||
}
|
||||
let _ = killpg(pid, Signal::SIGKILL);
|
||||
child
|
||||
.wait()
|
||||
.await
|
||||
.context("failed waiting for terminal child after SIGKILL")
|
||||
}
|
||||
|
||||
async fn send_json<S, T>(sink: &mut S, value: &T) -> Result<()>
|
||||
where
|
||||
S: futures_util::Sink<Message> + Unpin,
|
||||
S::Error: std::error::Error + Send + Sync + 'static,
|
||||
T: serde::Serialize,
|
||||
{
|
||||
let json = serde_json::to_string(value).context("failed to encode terminal frame")?;
|
||||
sink.send(Message::Text(json.into()))
|
||||
.await
|
||||
.context("failed to send terminal frame")
|
||||
}
|
||||
|
||||
fn terminal_websocket_url(server_url: &str, terminal_id: &TerminalId) -> Result<url::Url> {
|
||||
let mut url = url::Url::parse(server_url).context("invalid server_url")?;
|
||||
let scheme = match url.scheme() {
|
||||
"http" => "ws",
|
||||
"https" => "wss",
|
||||
"ws" => "ws",
|
||||
"wss" => "wss",
|
||||
other => anyhow::bail!("unsupported server_url scheme `{other}`"),
|
||||
};
|
||||
url.set_scheme(scheme)
|
||||
.map_err(|_| anyhow::anyhow!("failed to convert server_url scheme"))?;
|
||||
url.set_path(&format!(
|
||||
"/api/v1/agent/terminals/{}/ws",
|
||||
terminal_id.as_str()
|
||||
));
|
||||
url.set_query(None);
|
||||
url.set_fragment(None);
|
||||
Ok(url)
|
||||
}
|
||||
|
||||
fn validate_size(rows: u16, cols: u16) -> Result<()> {
|
||||
if !(1..=300).contains(&rows) || !(1..=500).contains(&cols) {
|
||||
anyhow::bail!("terminal size is outside supported bounds");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn terminal_url_uses_dedicated_agent_path() {
|
||||
let id = TerminalId::new("term-1").expect("terminal id");
|
||||
let url = terminal_websocket_url("https://example.com/base", &id).expect("url");
|
||||
assert_eq!(
|
||||
url.as_str(),
|
||||
"wss://example.com/api/v1/agent/terminals/term-1/ws"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -6,6 +6,7 @@ use serde::{Deserialize, Serialize};
|
||||
use std::time::Instant;
|
||||
use tracing::{info, info_span, warn};
|
||||
use uuid::Uuid;
|
||||
use wakey_agent::protocol::AgentCapability;
|
||||
use wakey_agent::protocol::{AgentCommand, ErrorPayload, RequestId, ServerMessage};
|
||||
|
||||
use crate::api::ApiError;
|
||||
@@ -17,6 +18,7 @@ pub struct AgentStatus {
|
||||
pub agent_id: String,
|
||||
pub connected: bool,
|
||||
pub nickname: Option<String>,
|
||||
pub capabilities: Vec<AgentCapability>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
@@ -43,6 +45,10 @@ pub async fn list_agents(State(state): State<AppState>) -> Result<impl IntoRespo
|
||||
.into_iter()
|
||||
.map(|(agent_id, nickname)| AgentStatus {
|
||||
connected: sessions.contains_key(&agent_id),
|
||||
capabilities: sessions
|
||||
.get(&agent_id)
|
||||
.map(|session| session.capabilities.clone())
|
||||
.unwrap_or_default(),
|
||||
agent_id,
|
||||
nickname,
|
||||
})
|
||||
|
||||
@@ -5,6 +5,7 @@ mod alerts;
|
||||
mod audit;
|
||||
mod commands;
|
||||
mod control;
|
||||
mod terminals;
|
||||
|
||||
pub use alerts::{active_alerts, alert_history, alerts_stream};
|
||||
pub use audit::list_audit_events;
|
||||
@@ -16,6 +17,10 @@ pub use control::{
|
||||
list_fleet_devices, list_known_devices, merge_known_device, refresh_fleet_devices,
|
||||
revoke_agent, revoke_enroll_token, set_agent_nickname, state_stats, wake_fleet_device,
|
||||
};
|
||||
pub use terminals::{
|
||||
agent_terminal_ws, attach_terminal, close_terminal, create_terminal, get_terminal,
|
||||
operator_terminal_ws,
|
||||
};
|
||||
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
@@ -0,0 +1,617 @@
|
||||
use axum::Json;
|
||||
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
|
||||
use axum::extract::{Path, State};
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::Response;
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tracing::{info, warn};
|
||||
use wakey_agent::protocol::{
|
||||
AgentCapability, ServerMessage, TerminalAgentHandshake, TerminalControl, TerminalId,
|
||||
TerminalOperatorHandshake,
|
||||
};
|
||||
|
||||
use crate::api::ApiError;
|
||||
use crate::runtime::terminals::{
|
||||
TERMINAL_ABSOLUTE_TIMEOUT, TERMINAL_ATTACH_TIMEOUT, TERMINAL_DISCONNECT_GRACE,
|
||||
TERMINAL_MAX_FRAME_BYTES, TerminalRelayFrame,
|
||||
};
|
||||
use crate::runtime::{AppState, SessionEvent};
|
||||
use crate::state::AuditEventInput;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct CreateTerminalRequest {
|
||||
pub agent_id: String,
|
||||
#[serde(default = "default_rows")]
|
||||
pub rows: u16,
|
||||
#[serde(default = "default_cols")]
|
||||
pub cols: u16,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct TerminalSessionResponse {
|
||||
pub terminal_id: String,
|
||||
pub agent_id: String,
|
||||
pub created_at_unix: u64,
|
||||
pub agent_attached: bool,
|
||||
pub operator_attached: bool,
|
||||
pub websocket_url: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub attachment_token: Option<String>,
|
||||
}
|
||||
|
||||
pub async fn create_terminal(
|
||||
State(state): State<AppState>,
|
||||
Json(request): Json<CreateTerminalRequest>,
|
||||
) -> Result<(StatusCode, Json<TerminalSessionResponse>), ApiError> {
|
||||
validate_size(request.rows, request.cols)?;
|
||||
let agent_tx = {
|
||||
let sessions = state.sessions.read().await;
|
||||
let session = sessions.get(&request.agent_id).ok_or_else(|| {
|
||||
ApiError::new(
|
||||
StatusCode::NOT_FOUND,
|
||||
"agent_not_connected",
|
||||
"agent is not connected",
|
||||
)
|
||||
})?;
|
||||
if !session.capabilities.contains(&AgentCapability::Terminal) {
|
||||
return Err(ApiError::new(
|
||||
StatusCode::CONFLICT,
|
||||
"terminal_not_supported",
|
||||
"agent has not advertised terminal capability",
|
||||
));
|
||||
}
|
||||
session.tx.clone()
|
||||
};
|
||||
|
||||
let created = state
|
||||
.terminals
|
||||
.create(request.agent_id.clone())
|
||||
.await
|
||||
.map_err(registry_error)?;
|
||||
let terminal_id = TerminalId::new(created.terminal_id.clone()).map_err(|message| {
|
||||
ApiError::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"terminal_id_invalid",
|
||||
message,
|
||||
)
|
||||
})?;
|
||||
if agent_tx
|
||||
.send(SessionEvent::Message(ServerMessage::OpenTerminal {
|
||||
terminal_id: terminal_id.clone(),
|
||||
relay_token: created.relay_token,
|
||||
rows: request.rows,
|
||||
cols: request.cols,
|
||||
}))
|
||||
.is_err()
|
||||
{
|
||||
state.terminals.remove(terminal_id.as_str()).await;
|
||||
return Err(ApiError::new(
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"agent_send_failed",
|
||||
"failed to send terminal request to agent",
|
||||
));
|
||||
}
|
||||
|
||||
spawn_absolute_timeout(state.clone(), terminal_id.clone(), request.agent_id.clone());
|
||||
append_terminal_audit(
|
||||
&state,
|
||||
&request.agent_id,
|
||||
terminal_id.as_str(),
|
||||
TerminalAudit {
|
||||
actor_type: "admin_api",
|
||||
event_type: "terminal_request",
|
||||
outcome: "sent",
|
||||
message: "terminal session requested",
|
||||
metadata: serde_json::json!({ "rows": request.rows, "cols": request.cols }),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
info!(terminal_id = %terminal_id, agent_id = %request.agent_id, "terminal session requested");
|
||||
Ok((
|
||||
StatusCode::CREATED,
|
||||
Json(TerminalSessionResponse {
|
||||
websocket_url: operator_ws_path(terminal_id.as_str()),
|
||||
terminal_id: terminal_id.to_string(),
|
||||
agent_id: request.agent_id,
|
||||
created_at_unix: created.created_at_unix,
|
||||
agent_attached: false,
|
||||
operator_attached: false,
|
||||
attachment_token: Some(created.attachment_token),
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
pub async fn get_terminal(
|
||||
State(state): State<AppState>,
|
||||
Path(terminal_id): Path<String>,
|
||||
) -> Result<Json<TerminalSessionResponse>, ApiError> {
|
||||
let (agent_id, created_at_unix, agent_attached, operator_attached) = state
|
||||
.terminals
|
||||
.summary(&terminal_id)
|
||||
.await
|
||||
.ok_or_else(|| terminal_not_found(&terminal_id))?;
|
||||
Ok(Json(TerminalSessionResponse {
|
||||
websocket_url: operator_ws_path(&terminal_id),
|
||||
terminal_id,
|
||||
agent_id,
|
||||
created_at_unix,
|
||||
agent_attached,
|
||||
operator_attached,
|
||||
attachment_token: None,
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn attach_terminal(
|
||||
State(state): State<AppState>,
|
||||
Path(terminal_id): Path<String>,
|
||||
) -> Result<Json<TerminalSessionResponse>, ApiError> {
|
||||
let attachment_token = state
|
||||
.terminals
|
||||
.issue_attachment_token(&terminal_id)
|
||||
.await
|
||||
.map_err(registry_error)?;
|
||||
let (agent_id, created_at_unix, agent_attached, operator_attached) = state
|
||||
.terminals
|
||||
.summary(&terminal_id)
|
||||
.await
|
||||
.ok_or_else(|| terminal_not_found(&terminal_id))?;
|
||||
Ok(Json(TerminalSessionResponse {
|
||||
websocket_url: operator_ws_path(&terminal_id),
|
||||
terminal_id,
|
||||
agent_id,
|
||||
created_at_unix,
|
||||
agent_attached,
|
||||
operator_attached,
|
||||
attachment_token: Some(attachment_token),
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn close_terminal(
|
||||
State(state): State<AppState>,
|
||||
Path(terminal_id): Path<String>,
|
||||
) -> Result<StatusCode, ApiError> {
|
||||
if let Some(agent_id) = close_registered_terminal(&state, &terminal_id).await {
|
||||
info!(terminal_id, agent_id, "terminal session closed by operator");
|
||||
append_terminal_audit(
|
||||
&state,
|
||||
&agent_id,
|
||||
&terminal_id,
|
||||
TerminalAudit {
|
||||
actor_type: "admin_api",
|
||||
event_type: "terminal_close",
|
||||
outcome: "ok",
|
||||
message: "terminal session closed by operator",
|
||||
metadata: serde_json::json!({}),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
} else {
|
||||
if !state.terminals.was_closed(&terminal_id).await {
|
||||
return Err(terminal_not_found(&terminal_id));
|
||||
}
|
||||
info!(terminal_id, "terminal session was already closed");
|
||||
}
|
||||
Ok(StatusCode::NO_CONTENT)
|
||||
}
|
||||
|
||||
pub async fn agent_terminal_ws(
|
||||
ws: WebSocketUpgrade,
|
||||
State(state): State<AppState>,
|
||||
Path(terminal_id): Path<String>,
|
||||
) -> Response {
|
||||
ws.max_message_size(TERMINAL_MAX_FRAME_BYTES)
|
||||
.on_upgrade(move |socket| handle_agent_terminal_socket(state, terminal_id, socket))
|
||||
}
|
||||
|
||||
pub async fn operator_terminal_ws(
|
||||
ws: WebSocketUpgrade,
|
||||
State(state): State<AppState>,
|
||||
Path(terminal_id): Path<String>,
|
||||
) -> Response {
|
||||
ws.max_message_size(TERMINAL_MAX_FRAME_BYTES)
|
||||
.on_upgrade(move |socket| handle_operator_terminal_socket(state, terminal_id, socket))
|
||||
}
|
||||
|
||||
async fn handle_agent_terminal_socket(state: AppState, terminal_id: String, mut socket: WebSocket) {
|
||||
let auth = match receive_text_handshake(&mut socket).await.and_then(|text| {
|
||||
serde_json::from_str::<TerminalAgentHandshake>(&text).map_err(|_| "invalid handshake")
|
||||
}) {
|
||||
Ok(TerminalAgentHandshake::Auth {
|
||||
agent_id,
|
||||
relay_token,
|
||||
}) => (agent_id, relay_token),
|
||||
Err(code) => {
|
||||
close_socket(&mut socket, code).await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
let (mut outbound, pending) = match state
|
||||
.terminals
|
||||
.attach_agent(&terminal_id, &auth.0, &auth.1)
|
||||
.await
|
||||
{
|
||||
Ok(outbound) => outbound,
|
||||
Err(code) => {
|
||||
close_socket(&mut socket, code).await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
info!(terminal_id, agent_id = %auth.0, "agent terminal socket attached");
|
||||
append_terminal_audit(
|
||||
&state,
|
||||
&auth.0,
|
||||
&terminal_id,
|
||||
TerminalAudit {
|
||||
actor_type: "agent",
|
||||
event_type: "terminal_agent_attach",
|
||||
outcome: "ok",
|
||||
message: "agent terminal transport attached",
|
||||
metadata: serde_json::json!({}),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let (mut write, mut read) = socket.split();
|
||||
for frame in pending {
|
||||
let closes = matches!(frame, TerminalRelayFrame::Close);
|
||||
if send_relay_frame(&mut write, frame).await.is_err() || closes {
|
||||
state.terminals.remove(&terminal_id).await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
loop {
|
||||
tokio::select! {
|
||||
outbound_frame = outbound.recv() => {
|
||||
let Some(frame) = outbound_frame else { break; };
|
||||
if send_relay_frame(&mut write, frame).await.is_err() { break; }
|
||||
}
|
||||
incoming = read.next() => {
|
||||
let Some(Ok(message)) = incoming else { break; };
|
||||
match agent_relay_frame(message) {
|
||||
Ok(Some(frame)) => {
|
||||
let closes = matches!(frame, TerminalRelayFrame::Close);
|
||||
audit_agent_control_frame(&state, &auth.0, &terminal_id, &frame).await;
|
||||
if state.terminals.relay_from_agent(&terminal_id, frame).await.is_err() { break; }
|
||||
if closes { break; }
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(code) => {
|
||||
warn!(terminal_id, code, "invalid agent terminal frame");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
state.terminals.remove(&terminal_id).await;
|
||||
info!(
|
||||
terminal_id,
|
||||
"agent terminal socket detached; session closed"
|
||||
);
|
||||
}
|
||||
|
||||
async fn handle_operator_terminal_socket(
|
||||
state: AppState,
|
||||
terminal_id: String,
|
||||
mut socket: WebSocket,
|
||||
) {
|
||||
let attachment_token = match receive_text_handshake(&mut socket).await.and_then(|text| {
|
||||
serde_json::from_str::<TerminalOperatorHandshake>(&text).map_err(|_| "invalid handshake")
|
||||
}) {
|
||||
Ok(TerminalOperatorHandshake::Attach { attachment_token }) => attachment_token,
|
||||
Err(code) => {
|
||||
close_socket(&mut socket, code).await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
let (mut outbound, replay) = match state
|
||||
.terminals
|
||||
.attach_operator(&terminal_id, &attachment_token)
|
||||
.await
|
||||
{
|
||||
Ok(attached) => attached,
|
||||
Err(code) => {
|
||||
close_socket(&mut socket, code).await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
info!(terminal_id, "operator terminal socket attached");
|
||||
if let Some((agent_id, _, _, _)) = state.terminals.summary(&terminal_id).await {
|
||||
append_terminal_audit(
|
||||
&state,
|
||||
&agent_id,
|
||||
&terminal_id,
|
||||
TerminalAudit {
|
||||
actor_type: "admin_api",
|
||||
event_type: "terminal_operator_attach",
|
||||
outcome: "ok",
|
||||
message: "operator terminal transport attached",
|
||||
metadata: serde_json::json!({}),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let (mut write, mut read) = socket.split();
|
||||
for frame in replay {
|
||||
if send_relay_frame(&mut write, frame).await.is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
let mut explicit_close = false;
|
||||
loop {
|
||||
tokio::select! {
|
||||
outbound_frame = outbound.recv() => {
|
||||
let Some(frame) = outbound_frame else { break; };
|
||||
if send_relay_frame(&mut write, frame).await.is_err() { break; }
|
||||
}
|
||||
incoming = read.next() => {
|
||||
let Some(Ok(message)) = incoming else { break; };
|
||||
match operator_relay_frame(message) {
|
||||
Ok(Some(frame)) => {
|
||||
explicit_close = matches!(frame, TerminalRelayFrame::Close);
|
||||
if state.terminals.relay_to_agent(&terminal_id, frame).await.is_err() { break; }
|
||||
if explicit_close { break; }
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(code) => {
|
||||
warn!(terminal_id, code, "invalid operator terminal frame");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if explicit_close {
|
||||
close_registered_terminal(&state, &terminal_id).await;
|
||||
} else if let Some(detached_at) = state.terminals.detach_operator(&terminal_id).await {
|
||||
let terminals = state.terminals.clone();
|
||||
let terminal_id_for_grace = terminal_id.clone();
|
||||
tokio::spawn(async move {
|
||||
tokio::time::sleep(TERMINAL_DISCONNECT_GRACE).await;
|
||||
terminals
|
||||
.remove_if_still_detached(&terminal_id_for_grace, detached_at)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
info!(
|
||||
terminal_id,
|
||||
explicit_close, "operator terminal socket detached"
|
||||
);
|
||||
}
|
||||
|
||||
fn agent_relay_frame(message: Message) -> Result<Option<TerminalRelayFrame>, &'static str> {
|
||||
match message {
|
||||
Message::Binary(bytes) => Ok(Some(TerminalRelayFrame::Binary(bytes.to_vec()))),
|
||||
Message::Text(text) => {
|
||||
let control: TerminalControl =
|
||||
serde_json::from_str(&text).map_err(|_| "terminal_control_invalid")?;
|
||||
if !matches!(
|
||||
control,
|
||||
TerminalControl::Ready
|
||||
| TerminalControl::Exited { .. }
|
||||
| TerminalControl::Error { .. }
|
||||
| TerminalControl::Close
|
||||
) {
|
||||
return Err("terminal_control_direction_invalid");
|
||||
}
|
||||
Ok(Some(TerminalRelayFrame::Text(
|
||||
serde_json::to_string(&control).map_err(|_| "terminal_control_invalid")?,
|
||||
)))
|
||||
}
|
||||
Message::Ping(_) | Message::Pong(_) => Ok(None),
|
||||
Message::Close(_) => Ok(Some(TerminalRelayFrame::Close)),
|
||||
}
|
||||
}
|
||||
|
||||
fn operator_relay_frame(message: Message) -> Result<Option<TerminalRelayFrame>, &'static str> {
|
||||
match message {
|
||||
Message::Binary(bytes) => Ok(Some(TerminalRelayFrame::Binary(bytes.to_vec()))),
|
||||
Message::Text(text) => {
|
||||
let control: TerminalControl =
|
||||
serde_json::from_str(&text).map_err(|_| "terminal_control_invalid")?;
|
||||
if !matches!(
|
||||
control,
|
||||
TerminalControl::Resize { .. } | TerminalControl::Close
|
||||
) {
|
||||
return Err("terminal_control_direction_invalid");
|
||||
}
|
||||
if let TerminalControl::Resize { rows, cols } = control {
|
||||
validate_size(rows, cols).map_err(|_| "terminal_size_invalid")?;
|
||||
}
|
||||
Ok(Some(TerminalRelayFrame::Text(
|
||||
serde_json::to_string(&control).map_err(|_| "terminal_control_invalid")?,
|
||||
)))
|
||||
}
|
||||
Message::Ping(_) | Message::Pong(_) => Ok(None),
|
||||
Message::Close(_) => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
async fn receive_text_handshake(socket: &mut WebSocket) -> Result<String, &'static str> {
|
||||
match tokio::time::timeout(TERMINAL_ATTACH_TIMEOUT, socket.recv()).await {
|
||||
Ok(Some(Ok(Message::Text(text)))) => Ok(text.to_string()),
|
||||
Ok(_) => Err("terminal_handshake_required"),
|
||||
Err(_) => Err("terminal_handshake_timeout"),
|
||||
}
|
||||
}
|
||||
|
||||
async fn close_socket(socket: &mut WebSocket, code: &str) {
|
||||
let control = TerminalControl::Error {
|
||||
code: code.into(),
|
||||
message: code.replace('_', " "),
|
||||
};
|
||||
if let Ok(json) = serde_json::to_string(&control) {
|
||||
let _ = socket.send(Message::Text(json.into())).await;
|
||||
}
|
||||
let _ = socket.send(Message::Close(None)).await;
|
||||
}
|
||||
|
||||
async fn send_relay_frame<S>(write: &mut S, frame: TerminalRelayFrame) -> Result<(), axum::Error>
|
||||
where
|
||||
S: futures_util::Sink<Message, Error = axum::Error> + Unpin,
|
||||
{
|
||||
let message = match frame {
|
||||
TerminalRelayFrame::Binary(bytes) => Message::Binary(bytes.into()),
|
||||
TerminalRelayFrame::Text(text) => Message::Text(text.into()),
|
||||
TerminalRelayFrame::Close => Message::Close(None),
|
||||
};
|
||||
write.send(message).await
|
||||
}
|
||||
|
||||
async fn close_registered_terminal(state: &AppState, terminal_id: &str) -> Option<String> {
|
||||
let agent_id = state.terminals.remove(terminal_id).await?;
|
||||
if let Some(session) = state.sessions.read().await.get(&agent_id) {
|
||||
let terminal_id = TerminalId::new(terminal_id.to_string()).ok()?;
|
||||
let _ = session
|
||||
.tx
|
||||
.send(SessionEvent::Message(ServerMessage::CloseTerminal {
|
||||
terminal_id,
|
||||
}));
|
||||
}
|
||||
Some(agent_id)
|
||||
}
|
||||
|
||||
fn spawn_absolute_timeout(state: AppState, terminal_id: TerminalId, agent_id: String) {
|
||||
tokio::spawn(async move {
|
||||
tokio::time::sleep(TERMINAL_ABSOLUTE_TIMEOUT).await;
|
||||
if close_registered_terminal(&state, terminal_id.as_str())
|
||||
.await
|
||||
.is_some()
|
||||
{
|
||||
info!(terminal_id = %terminal_id, agent_id, "terminal absolute timeout reached");
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
struct TerminalAudit<'a> {
|
||||
actor_type: &'a str,
|
||||
event_type: &'a str,
|
||||
outcome: &'a str,
|
||||
message: &'a str,
|
||||
metadata: serde_json::Value,
|
||||
}
|
||||
|
||||
async fn audit_agent_control_frame(
|
||||
state: &AppState,
|
||||
agent_id: &str,
|
||||
terminal_id: &str,
|
||||
frame: &TerminalRelayFrame,
|
||||
) {
|
||||
let TerminalRelayFrame::Text(text) = frame else {
|
||||
return;
|
||||
};
|
||||
let Ok(control) = serde_json::from_str::<TerminalControl>(text) else {
|
||||
return;
|
||||
};
|
||||
let audit = match control {
|
||||
TerminalControl::Ready => TerminalAudit {
|
||||
actor_type: "agent",
|
||||
event_type: "terminal_ready",
|
||||
outcome: "ok",
|
||||
message: "terminal PTY ready",
|
||||
metadata: serde_json::json!({}),
|
||||
},
|
||||
TerminalControl::Exited { exit_code } => TerminalAudit {
|
||||
actor_type: "agent",
|
||||
event_type: "terminal_exit",
|
||||
outcome: "exited",
|
||||
message: "terminal process exited",
|
||||
metadata: serde_json::json!({ "exit_code": exit_code }),
|
||||
},
|
||||
TerminalControl::Error { code, .. } => TerminalAudit {
|
||||
actor_type: "agent",
|
||||
event_type: "terminal_error",
|
||||
outcome: "error",
|
||||
message: "terminal worker reported an error",
|
||||
metadata: serde_json::json!({ "code": code }),
|
||||
},
|
||||
_ => return,
|
||||
};
|
||||
append_terminal_audit(state, agent_id, terminal_id, audit).await;
|
||||
}
|
||||
|
||||
async fn append_terminal_audit(
|
||||
state: &AppState,
|
||||
agent_id: &str,
|
||||
terminal_id: &str,
|
||||
audit: TerminalAudit<'_>,
|
||||
) {
|
||||
if let Err(err) = state
|
||||
.store
|
||||
.append_audit_event(AuditEventInput {
|
||||
actor_type: audit.actor_type.into(),
|
||||
actor_id: None,
|
||||
agent_id: Some(agent_id.to_string()),
|
||||
request_id: Some(terminal_id.to_string()),
|
||||
event_type: audit.event_type.into(),
|
||||
outcome: audit.outcome.into(),
|
||||
latency_ms: None,
|
||||
message: audit.message.into(),
|
||||
metadata: audit.metadata,
|
||||
})
|
||||
.await
|
||||
{
|
||||
warn!(terminal_id, event_type = audit.event_type, error = %err, "failed to append terminal audit event");
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_size(rows: u16, cols: u16) -> Result<(), ApiError> {
|
||||
if !(1..=300).contains(&rows) || !(1..=500).contains(&cols) {
|
||||
return Err(ApiError::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
"terminal_size_invalid",
|
||||
"terminal rows must be 1..=300 and columns must be 1..=500",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn registry_error(code: &'static str) -> ApiError {
|
||||
let status = match code {
|
||||
"terminal_not_found" => StatusCode::NOT_FOUND,
|
||||
"terminal_relay_token_invalid" | "terminal_attachment_token_invalid" => {
|
||||
StatusCode::UNAUTHORIZED
|
||||
}
|
||||
_ => StatusCode::CONFLICT,
|
||||
};
|
||||
ApiError::new(status, code, code.replace('_', " "))
|
||||
}
|
||||
|
||||
fn terminal_not_found(terminal_id: &str) -> ApiError {
|
||||
ApiError::new(
|
||||
StatusCode::NOT_FOUND,
|
||||
"terminal_not_found",
|
||||
format!("terminal session {terminal_id} was not found"),
|
||||
)
|
||||
}
|
||||
|
||||
fn operator_ws_path(terminal_id: &str) -> String {
|
||||
format!("/api/v1/control/terminals/{terminal_id}/ws")
|
||||
}
|
||||
|
||||
const fn default_rows() -> u16 {
|
||||
24
|
||||
}
|
||||
|
||||
const fn default_cols() -> u16 {
|
||||
80
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn terminal_control_direction_is_enforced() {
|
||||
let resize = serde_json::to_string(&TerminalControl::Resize { rows: 24, cols: 80 })
|
||||
.expect("resize json");
|
||||
assert!(operator_relay_frame(Message::Text(resize.into())).is_ok());
|
||||
|
||||
let ready = serde_json::to_string(&TerminalControl::Ready).expect("ready json");
|
||||
assert!(operator_relay_frame(Message::Text(ready.into())).is_err());
|
||||
}
|
||||
}
|
||||
@@ -16,7 +16,7 @@ use tower_http::services::ServeFile;
|
||||
use tracing::info;
|
||||
#[cfg(unix)]
|
||||
use tracing::warn;
|
||||
use wakey_agent::protocol::{ErrorPayload, ServerMessage};
|
||||
use wakey_agent::protocol::{AgentCapability, ErrorPayload, ServerMessage};
|
||||
|
||||
use crate::api;
|
||||
use crate::config;
|
||||
@@ -25,6 +25,7 @@ use crate::ws;
|
||||
|
||||
mod admin;
|
||||
mod process;
|
||||
pub mod terminals;
|
||||
pub use admin::revoke_agent;
|
||||
pub use admin::{
|
||||
issue_enroll_token, list_enroll_tokens, migrate_sqlite_state, revoke_enroll_token, state_stats,
|
||||
@@ -41,12 +42,14 @@ pub struct AppState {
|
||||
pub public_url: String,
|
||||
pub command_timeout: Duration,
|
||||
pub enroll_token_ttl: Duration,
|
||||
pub terminals: terminals::TerminalRegistry,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AgentSession {
|
||||
pub connection_id: String,
|
||||
pub tx: mpsc::UnboundedSender<SessionEvent>,
|
||||
pub capabilities: Vec<AgentCapability>,
|
||||
}
|
||||
#[derive(Clone)]
|
||||
pub enum SessionEvent {
|
||||
@@ -72,6 +75,10 @@ fn public_api_routes(ui_dist_dir: std::path::PathBuf) -> Router<AppState> {
|
||||
.route("/healthz", get(api::healthz))
|
||||
.route("/api/v1/agents/enroll", post(api::enroll))
|
||||
.route("/api/v1/agent/ws", get(ws::agent_ws))
|
||||
.route(
|
||||
"/api/v1/agent/terminals/{terminal_id}/ws",
|
||||
get(api::agent_terminal_ws),
|
||||
)
|
||||
}
|
||||
|
||||
fn control_api_routes() -> Router<AppState> {
|
||||
@@ -135,6 +142,19 @@ fn control_api_routes() -> Router<AppState> {
|
||||
"/api/v1/control/agents/{agent_id}/command",
|
||||
post(api::run_command),
|
||||
)
|
||||
.route("/api/v1/control/terminals", post(api::create_terminal))
|
||||
.route(
|
||||
"/api/v1/control/terminals/{terminal_id}",
|
||||
get(api::get_terminal).delete(api::close_terminal),
|
||||
)
|
||||
.route(
|
||||
"/api/v1/control/terminals/{terminal_id}/attach",
|
||||
post(api::attach_terminal),
|
||||
)
|
||||
.route(
|
||||
"/api/v1/control/terminals/{terminal_id}/ws",
|
||||
get(api::operator_terminal_ws),
|
||||
)
|
||||
}
|
||||
|
||||
/// Starts the control-plane HTTP and websocket surfaces and manages daemon lifecycle hooks.
|
||||
@@ -157,6 +177,7 @@ pub async fn serve(daemon: config::DaemonConfig) -> Result<()> {
|
||||
public_url: daemon.public_url.clone(),
|
||||
command_timeout: daemon.command_timeout,
|
||||
enroll_token_ttl: daemon.enroll_token_ttl,
|
||||
terminals: terminals::TerminalRegistry::new(),
|
||||
};
|
||||
|
||||
// Keep route classes explicit so edge policy can map directly:
|
||||
|
||||
@@ -0,0 +1,481 @@
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use tokio::sync::{Mutex, mpsc};
|
||||
use uuid::Uuid;
|
||||
|
||||
pub const TERMINAL_RELAY_QUEUE: usize = 32;
|
||||
pub const TERMINAL_MAX_FRAME_BYTES: usize = 64 * 1024;
|
||||
pub const TERMINAL_REPLAY_BYTES: usize = 256 * 1024;
|
||||
pub const TERMINAL_MAX_SESSIONS_PER_AGENT: usize = 2;
|
||||
pub const TERMINAL_ATTACH_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
pub const TERMINAL_DISCONNECT_GRACE: Duration = Duration::from_secs(15);
|
||||
pub const TERMINAL_ABSOLUTE_TIMEOUT: Duration = Duration::from_secs(30 * 60);
|
||||
const TERMINAL_TOMBSTONE_TTL: Duration = Duration::from_secs(5 * 60);
|
||||
const TERMINAL_MAX_TOMBSTONES: usize = 1024;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum TerminalRelayFrame {
|
||||
Binary(Vec<u8>),
|
||||
Text(String),
|
||||
Close,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct TerminalRegistry {
|
||||
inner: Arc<Mutex<HashMap<String, TerminalSession>>>,
|
||||
closed: Arc<Mutex<HashMap<String, Instant>>>,
|
||||
}
|
||||
|
||||
struct TerminalSession {
|
||||
agent_id: String,
|
||||
created_at_unix: u64,
|
||||
expires_at: Instant,
|
||||
relay_token: Option<String>,
|
||||
attachment_token: Option<String>,
|
||||
agent_tx: Option<mpsc::Sender<TerminalRelayFrame>>,
|
||||
pending_agent: VecDeque<TerminalRelayFrame>,
|
||||
pending_agent_bytes: usize,
|
||||
operator_tx: Option<mpsc::Sender<TerminalRelayFrame>>,
|
||||
operator_detached_at: Option<Instant>,
|
||||
replay: VecDeque<TerminalRelayFrame>,
|
||||
replay_bytes: usize,
|
||||
}
|
||||
|
||||
pub struct CreatedTerminal {
|
||||
pub terminal_id: String,
|
||||
pub relay_token: String,
|
||||
pub attachment_token: String,
|
||||
pub created_at_unix: u64,
|
||||
}
|
||||
|
||||
impl Default for TerminalRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl TerminalRegistry {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
inner: Arc::new(Mutex::new(HashMap::new())),
|
||||
closed: Arc::new(Mutex::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn create(&self, agent_id: String) -> Result<CreatedTerminal, &'static str> {
|
||||
let mut sessions = self.inner.lock().await;
|
||||
if sessions
|
||||
.values()
|
||||
.filter(|session| session.agent_id == agent_id)
|
||||
.count()
|
||||
>= TERMINAL_MAX_SESSIONS_PER_AGENT
|
||||
{
|
||||
return Err("agent_terminal_limit_reached");
|
||||
}
|
||||
|
||||
let terminal_id = Uuid::new_v4().to_string();
|
||||
let relay_token = new_token();
|
||||
let attachment_token = new_token();
|
||||
let created_at_unix = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
sessions.insert(
|
||||
terminal_id.clone(),
|
||||
TerminalSession {
|
||||
agent_id,
|
||||
created_at_unix,
|
||||
expires_at: Instant::now() + TERMINAL_ABSOLUTE_TIMEOUT,
|
||||
relay_token: Some(relay_token.clone()),
|
||||
attachment_token: Some(attachment_token.clone()),
|
||||
agent_tx: None,
|
||||
pending_agent: VecDeque::new(),
|
||||
pending_agent_bytes: 0,
|
||||
operator_tx: None,
|
||||
operator_detached_at: None,
|
||||
replay: VecDeque::new(),
|
||||
replay_bytes: 0,
|
||||
},
|
||||
);
|
||||
|
||||
Ok(CreatedTerminal {
|
||||
terminal_id,
|
||||
relay_token,
|
||||
attachment_token,
|
||||
created_at_unix,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn remove(&self, terminal_id: &str) -> Option<String> {
|
||||
let session = self.inner.lock().await.remove(terminal_id)?;
|
||||
self.remember_closed(terminal_id).await;
|
||||
if let Some(tx) = session.agent_tx {
|
||||
let _ =
|
||||
tokio::time::timeout(Duration::from_secs(1), tx.send(TerminalRelayFrame::Close))
|
||||
.await;
|
||||
}
|
||||
if let Some(tx) = session.operator_tx {
|
||||
let _ =
|
||||
tokio::time::timeout(Duration::from_secs(1), tx.send(TerminalRelayFrame::Close))
|
||||
.await;
|
||||
}
|
||||
Some(session.agent_id)
|
||||
}
|
||||
|
||||
/// Reports whether a session ID was recently removed. Tombstones make
|
||||
/// idempotent DELETE distinguishable from a completely unknown ID.
|
||||
pub async fn was_closed(&self, terminal_id: &str) -> bool {
|
||||
let mut closed = self.closed.lock().await;
|
||||
prune_tombstones(&mut closed);
|
||||
closed.contains_key(terminal_id)
|
||||
}
|
||||
|
||||
async fn remember_closed(&self, terminal_id: &str) {
|
||||
let mut closed = self.closed.lock().await;
|
||||
prune_tombstones(&mut closed);
|
||||
closed.insert(terminal_id.to_string(), Instant::now());
|
||||
if closed.len() > TERMINAL_MAX_TOMBSTONES
|
||||
&& let Some(oldest) = closed
|
||||
.iter()
|
||||
.min_by_key(|(_, closed_at)| **closed_at)
|
||||
.map(|(terminal_id, _)| terminal_id.clone())
|
||||
{
|
||||
closed.remove(&oldest);
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn remove_agent(&self, agent_id: &str) {
|
||||
let terminal_ids = {
|
||||
let sessions = self.inner.lock().await;
|
||||
sessions
|
||||
.iter()
|
||||
.filter(|(_, session)| session.agent_id == agent_id)
|
||||
.map(|(terminal_id, _)| terminal_id.clone())
|
||||
.collect::<Vec<_>>()
|
||||
};
|
||||
for terminal_id in terminal_ids {
|
||||
self.remove(&terminal_id).await;
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn issue_attachment_token(&self, terminal_id: &str) -> Result<String, &'static str> {
|
||||
let mut sessions = self.inner.lock().await;
|
||||
let session = active_session(&mut sessions, terminal_id)?;
|
||||
if session.operator_tx.is_some() {
|
||||
return Err("terminal_operator_already_attached");
|
||||
}
|
||||
let token = new_token();
|
||||
session.attachment_token = Some(token.clone());
|
||||
Ok(token)
|
||||
}
|
||||
|
||||
pub async fn attach_agent(
|
||||
&self,
|
||||
terminal_id: &str,
|
||||
agent_id: &str,
|
||||
relay_token: &str,
|
||||
) -> Result<(mpsc::Receiver<TerminalRelayFrame>, Vec<TerminalRelayFrame>), &'static str> {
|
||||
let mut sessions = self.inner.lock().await;
|
||||
let session = active_session(&mut sessions, terminal_id)?;
|
||||
if session.agent_id != agent_id {
|
||||
return Err("terminal_agent_mismatch");
|
||||
}
|
||||
if session.agent_tx.is_some() {
|
||||
return Err("terminal_agent_already_attached");
|
||||
}
|
||||
if session.relay_token.as_deref() != Some(relay_token) {
|
||||
return Err("terminal_relay_token_invalid");
|
||||
}
|
||||
session.relay_token = None;
|
||||
let (tx, rx) = mpsc::channel(TERMINAL_RELAY_QUEUE);
|
||||
session.agent_tx = Some(tx);
|
||||
let pending = session.pending_agent.drain(..).collect();
|
||||
session.pending_agent_bytes = 0;
|
||||
Ok((rx, pending))
|
||||
}
|
||||
|
||||
pub async fn attach_operator(
|
||||
&self,
|
||||
terminal_id: &str,
|
||||
attachment_token: &str,
|
||||
) -> Result<(mpsc::Receiver<TerminalRelayFrame>, Vec<TerminalRelayFrame>), &'static str> {
|
||||
let mut sessions = self.inner.lock().await;
|
||||
let session = active_session(&mut sessions, terminal_id)?;
|
||||
if session.operator_tx.is_some() {
|
||||
return Err("terminal_operator_already_attached");
|
||||
}
|
||||
if session.attachment_token.as_deref() != Some(attachment_token) {
|
||||
return Err("terminal_attachment_token_invalid");
|
||||
}
|
||||
session.attachment_token = None;
|
||||
session.operator_detached_at = None;
|
||||
let replay = session.replay.drain(..).collect();
|
||||
session.replay_bytes = 0;
|
||||
let (tx, rx) = mpsc::channel(TERMINAL_RELAY_QUEUE);
|
||||
session.operator_tx = Some(tx);
|
||||
Ok((rx, replay))
|
||||
}
|
||||
|
||||
pub async fn relay_to_agent(
|
||||
&self,
|
||||
terminal_id: &str,
|
||||
frame: TerminalRelayFrame,
|
||||
) -> Result<(), &'static str> {
|
||||
let tx = self
|
||||
.inner
|
||||
.lock()
|
||||
.await
|
||||
.get(terminal_id)
|
||||
.and_then(|session| session.agent_tx.clone());
|
||||
if let Some(tx) = tx {
|
||||
return tx
|
||||
.send(frame)
|
||||
.await
|
||||
.map_err(|_| "terminal_agent_disconnected");
|
||||
}
|
||||
|
||||
let mut sessions = self.inner.lock().await;
|
||||
let session = active_session(&mut sessions, terminal_id)?;
|
||||
session.pending_agent_bytes += relay_frame_size(&frame);
|
||||
session.pending_agent.push_back(frame);
|
||||
while session.pending_agent_bytes > TERMINAL_REPLAY_BYTES {
|
||||
if let Some(dropped) = session.pending_agent.pop_front() {
|
||||
session.pending_agent_bytes -= relay_frame_size(&dropped);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn relay_from_agent(
|
||||
&self,
|
||||
terminal_id: &str,
|
||||
frame: TerminalRelayFrame,
|
||||
) -> Result<(), &'static str> {
|
||||
let operator_tx = self
|
||||
.inner
|
||||
.lock()
|
||||
.await
|
||||
.get(terminal_id)
|
||||
.and_then(|session| session.operator_tx.clone());
|
||||
|
||||
if let Some(tx) = operator_tx {
|
||||
match tx.send(frame).await {
|
||||
Ok(()) => return Ok(()),
|
||||
Err(err) => {
|
||||
// The browser task may not have marked itself detached yet.
|
||||
// Preserve this frame so that race does not kill the PTY.
|
||||
let mut sessions = self.inner.lock().await;
|
||||
let session = active_session(&mut sessions, terminal_id)?;
|
||||
session.operator_tx = None;
|
||||
session
|
||||
.operator_detached_at
|
||||
.get_or_insert_with(Instant::now);
|
||||
push_replay(session, err.0);
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut sessions = self.inner.lock().await;
|
||||
let session = active_session(&mut sessions, terminal_id)?;
|
||||
push_replay(session, frame);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn reject(
|
||||
&self,
|
||||
terminal_id: &str,
|
||||
agent_id: &str,
|
||||
error_json: String,
|
||||
) -> Result<(), &'static str> {
|
||||
let matches_agent = self
|
||||
.inner
|
||||
.lock()
|
||||
.await
|
||||
.get(terminal_id)
|
||||
.is_some_and(|session| session.agent_id == agent_id);
|
||||
if !matches_agent {
|
||||
return Err("terminal_agent_mismatch");
|
||||
}
|
||||
self.relay_from_agent(terminal_id, TerminalRelayFrame::Text(error_json))
|
||||
.await?;
|
||||
self.relay_from_agent(terminal_id, TerminalRelayFrame::Close)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn detach_operator(&self, terminal_id: &str) -> Option<Instant> {
|
||||
let mut sessions = self.inner.lock().await;
|
||||
let session = sessions.get_mut(terminal_id)?;
|
||||
session.operator_tx = None;
|
||||
let detached_at = Instant::now();
|
||||
session.operator_detached_at = Some(detached_at);
|
||||
Some(detached_at)
|
||||
}
|
||||
|
||||
pub async fn remove_if_still_detached(&self, terminal_id: &str, detached_at: Instant) -> bool {
|
||||
let should_remove = self
|
||||
.inner
|
||||
.lock()
|
||||
.await
|
||||
.get(terminal_id)
|
||||
.is_some_and(|session| session.operator_detached_at == Some(detached_at));
|
||||
if should_remove {
|
||||
self.remove(terminal_id).await;
|
||||
}
|
||||
should_remove
|
||||
}
|
||||
|
||||
pub async fn summary(&self, terminal_id: &str) -> Option<(String, u64, bool, bool)> {
|
||||
self.inner.lock().await.get(terminal_id).map(|session| {
|
||||
(
|
||||
session.agent_id.clone(),
|
||||
session.created_at_unix,
|
||||
session.agent_tx.is_some(),
|
||||
session.operator_tx.is_some(),
|
||||
)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn active_session<'a>(
|
||||
sessions: &'a mut HashMap<String, TerminalSession>,
|
||||
terminal_id: &str,
|
||||
) -> Result<&'a mut TerminalSession, &'static str> {
|
||||
let session = sessions.get_mut(terminal_id).ok_or("terminal_not_found")?;
|
||||
if session.expires_at <= Instant::now() {
|
||||
return Err("terminal_expired");
|
||||
}
|
||||
Ok(session)
|
||||
}
|
||||
|
||||
fn new_token() -> String {
|
||||
format!("{}{}", Uuid::new_v4().simple(), Uuid::new_v4().simple())
|
||||
}
|
||||
|
||||
fn prune_tombstones(closed: &mut HashMap<String, Instant>) {
|
||||
closed.retain(|_, closed_at| closed_at.elapsed() < TERMINAL_TOMBSTONE_TTL);
|
||||
}
|
||||
|
||||
fn relay_frame_size(frame: &TerminalRelayFrame) -> usize {
|
||||
match frame {
|
||||
TerminalRelayFrame::Binary(bytes) => bytes.len(),
|
||||
TerminalRelayFrame::Text(text) => text.len(),
|
||||
TerminalRelayFrame::Close => 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn push_replay(session: &mut TerminalSession, frame: TerminalRelayFrame) {
|
||||
session.replay_bytes += relay_frame_size(&frame);
|
||||
session.replay.push_back(frame);
|
||||
while session.replay_bytes > TERMINAL_REPLAY_BYTES {
|
||||
if let Some(dropped) = session.replay.pop_front() {
|
||||
session.replay_bytes -= relay_frame_size(&dropped);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn credentials_are_scoped_and_single_use() {
|
||||
let registry = TerminalRegistry::new();
|
||||
let created = registry.create("router".into()).await.expect("create");
|
||||
|
||||
assert!(
|
||||
registry
|
||||
.attach_agent(&created.terminal_id, "other", &created.relay_token)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
let _ = registry
|
||||
.attach_agent(&created.terminal_id, "router", &created.relay_token)
|
||||
.await
|
||||
.expect("attach agent");
|
||||
assert_eq!(
|
||||
registry
|
||||
.attach_agent(&created.terminal_id, "router", &created.relay_token)
|
||||
.await
|
||||
.expect_err("relay token is single use"),
|
||||
"terminal_agent_already_attached"
|
||||
);
|
||||
|
||||
registry
|
||||
.attach_operator(&created.terminal_id, &created.attachment_token)
|
||||
.await
|
||||
.expect("attach operator");
|
||||
assert!(
|
||||
registry
|
||||
.attach_operator(&created.terminal_id, &created.attachment_token)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn detached_output_replay_is_bounded() {
|
||||
let registry = TerminalRegistry::new();
|
||||
let created = registry.create("router".into()).await.expect("create");
|
||||
for _ in 0..10 {
|
||||
registry
|
||||
.relay_from_agent(
|
||||
&created.terminal_id,
|
||||
TerminalRelayFrame::Binary(vec![0; TERMINAL_REPLAY_BYTES / 4]),
|
||||
)
|
||||
.await
|
||||
.expect("buffer output");
|
||||
}
|
||||
|
||||
let (_, replay) = registry
|
||||
.attach_operator(&created.terminal_id, &created.attachment_token)
|
||||
.await
|
||||
.expect("attach operator");
|
||||
assert!(replay.iter().map(relay_frame_size).sum::<usize>() <= TERMINAL_REPLAY_BYTES);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn operator_input_waits_for_agent_attachment() {
|
||||
let registry = TerminalRegistry::new();
|
||||
let created = registry.create("router".into()).await.expect("create");
|
||||
let resize = TerminalRelayFrame::Text(r#"{"type":"resize","rows":30,"cols":120}"#.into());
|
||||
registry
|
||||
.relay_to_agent(&created.terminal_id, resize)
|
||||
.await
|
||||
.expect("queue resize before agent attachment");
|
||||
|
||||
let (_, pending) = registry
|
||||
.attach_agent(&created.terminal_id, "router", &created.relay_token)
|
||||
.await
|
||||
.expect("attach agent");
|
||||
assert_eq!(pending.len(), 1);
|
||||
assert!(matches!(pending[0], TerminalRelayFrame::Text(_)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn agent_can_hold_two_terminal_sessions() {
|
||||
let registry = TerminalRegistry::new();
|
||||
registry.create("router".into()).await.expect("first");
|
||||
registry.create("router".into()).await.expect("second");
|
||||
let third = registry.create("router".into()).await;
|
||||
assert_eq!(third.err(), Some("agent_terminal_limit_reached"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn removed_session_is_distinct_from_unknown_session() {
|
||||
let registry = TerminalRegistry::new();
|
||||
let created = registry.create("router".into()).await.expect("create");
|
||||
|
||||
assert!(!registry.was_closed(&created.terminal_id).await);
|
||||
assert!(!registry.was_closed("never-existed").await);
|
||||
registry.remove(&created.terminal_id).await.expect("remove");
|
||||
assert!(registry.was_closed(&created.terminal_id).await);
|
||||
assert!(!registry.was_closed("never-existed").await);
|
||||
}
|
||||
}
|
||||
@@ -8,7 +8,9 @@ use std::time::Instant;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::{debug, info, info_span, warn};
|
||||
use uuid::Uuid;
|
||||
use wakey_agent::protocol::{ErrorPayload, RequestId, ServerMessage};
|
||||
use wakey_agent::protocol::{
|
||||
AgentCapability, ErrorPayload, RequestId, ServerMessage, TerminalControl, TerminalId,
|
||||
};
|
||||
use wakey_core::Device;
|
||||
|
||||
use crate::runtime::{AgentReply, AgentSession, AppState, SessionEvent};
|
||||
@@ -19,6 +21,8 @@ use crate::state::AuditEventInput;
|
||||
enum IncomingClientMessage {
|
||||
Hello {
|
||||
agent_id: String,
|
||||
#[serde(default)]
|
||||
capabilities: Vec<AgentCapability>,
|
||||
},
|
||||
Auth {
|
||||
agent_id: String,
|
||||
@@ -39,6 +43,17 @@ enum IncomingClientMessage {
|
||||
request_id: RequestId,
|
||||
error: ErrorPayload,
|
||||
},
|
||||
TerminalRejected {
|
||||
terminal_id: TerminalId,
|
||||
error: ErrorPayload,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct AgentConnectionState {
|
||||
authed_agent_id: Option<String>,
|
||||
hello_at: Option<Instant>,
|
||||
capabilities: Vec<AgentCapability>,
|
||||
}
|
||||
|
||||
pub async fn agent_ws(ws: WebSocketUpgrade, State(state): State<AppState>) -> impl IntoResponse {
|
||||
@@ -80,8 +95,7 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
|
||||
debug!("websocket writer loop ended");
|
||||
});
|
||||
|
||||
let mut authed_agent_id: Option<String> = None;
|
||||
let mut hello_at: Option<Instant> = None;
|
||||
let mut connection = AgentConnectionState::default();
|
||||
|
||||
loop {
|
||||
let frame = read.next().await;
|
||||
@@ -103,8 +117,7 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
|
||||
&state,
|
||||
&tx,
|
||||
&connection_id,
|
||||
&mut authed_agent_id,
|
||||
&mut hello_at,
|
||||
&mut connection,
|
||||
connected_at,
|
||||
&text,
|
||||
)
|
||||
@@ -126,7 +139,7 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(agent_id) = authed_agent_id {
|
||||
if let Some(agent_id) = connection.authed_agent_id {
|
||||
info!(agent_id = %agent_id, "agent disconnected");
|
||||
let mut sessions = state.sessions.write().await;
|
||||
let should_remove = sessions
|
||||
@@ -136,6 +149,8 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
|
||||
if should_remove {
|
||||
sessions.remove(&agent_id);
|
||||
}
|
||||
drop(sessions);
|
||||
state.terminals.remove_agent(&agent_id).await;
|
||||
if let Err(err) = state
|
||||
.store
|
||||
.append_audit_event(AuditEventInput {
|
||||
@@ -163,8 +178,7 @@ async fn process_agent_text(
|
||||
state: &AppState,
|
||||
tx: &mpsc::UnboundedSender<SessionEvent>,
|
||||
connection_id: &str,
|
||||
authed_agent_id: &mut Option<String>,
|
||||
hello_at: &mut Option<Instant>,
|
||||
connection: &mut AgentConnectionState,
|
||||
connected_at: Instant,
|
||||
text: &str,
|
||||
) -> Result<()> {
|
||||
@@ -172,20 +186,26 @@ async fn process_agent_text(
|
||||
serde_json::from_str(text).context("invalid client websocket payload")?;
|
||||
|
||||
match message {
|
||||
IncomingClientMessage::Hello { agent_id } => {
|
||||
IncomingClientMessage::Hello {
|
||||
agent_id,
|
||||
capabilities,
|
||||
} => {
|
||||
let now = Instant::now();
|
||||
if hello_at.is_none() {
|
||||
*hello_at = Some(now);
|
||||
if connection.hello_at.is_none() {
|
||||
connection.hello_at = Some(now);
|
||||
}
|
||||
let connect_to_hello_ms = connected_at.elapsed().as_millis() as u64;
|
||||
info!(agent_id = %agent_id, connect_to_hello_ms, "agent hello received");
|
||||
connection.capabilities = capabilities;
|
||||
}
|
||||
IncomingClientMessage::Auth {
|
||||
agent_id,
|
||||
agent_token,
|
||||
} => {
|
||||
let connect_to_auth_ms = connected_at.elapsed().as_millis() as u64;
|
||||
let hello_to_auth_ms = hello_at.map(|t| now_duration_ms(t.elapsed()));
|
||||
let hello_to_auth_ms = connection
|
||||
.hello_at
|
||||
.map(|time| now_duration_ms(time.elapsed()));
|
||||
if !state
|
||||
.store
|
||||
.verify_agent_token(&agent_id, &agent_token)
|
||||
@@ -219,9 +239,10 @@ async fn process_agent_text(
|
||||
AgentSession {
|
||||
connection_id: connection_id.to_string(),
|
||||
tx: tx.clone(),
|
||||
capabilities: connection.capabilities.clone(),
|
||||
},
|
||||
);
|
||||
*authed_agent_id = Some(agent_id.clone());
|
||||
connection.authed_agent_id = Some(agent_id.clone());
|
||||
info!(agent_id = %agent_id, connect_to_auth_ms, hello_to_auth_ms = hello_to_auth_ms.unwrap_or(0), "agent authenticated");
|
||||
if let Err(err) = state
|
||||
.store
|
||||
@@ -246,14 +267,14 @@ async fn process_agent_text(
|
||||
let _ = tx.send(SessionEvent::Message(ServerMessage::SyncDeviceSnapshot));
|
||||
}
|
||||
IncomingClientMessage::Heartbeat { agent_id } => {
|
||||
if authed_agent_id.as_deref() != Some(agent_id.as_str()) {
|
||||
if connection.authed_agent_id.as_deref() != Some(agent_id.as_str()) {
|
||||
anyhow::bail!("heartbeat for unauthenticated or mismatched agent");
|
||||
}
|
||||
ensure_current_session(state, &agent_id, connection_id).await?;
|
||||
debug!(agent_id = %agent_id, "heartbeat received");
|
||||
}
|
||||
IncomingClientMessage::DeviceSnapshot { agent_id, devices } => {
|
||||
if authed_agent_id.as_deref() != Some(agent_id.as_str()) {
|
||||
if connection.authed_agent_id.as_deref() != Some(agent_id.as_str()) {
|
||||
anyhow::bail!("device_snapshot for unauthenticated or mismatched agent");
|
||||
}
|
||||
ensure_current_session(state, &agent_id, connection_id).await?;
|
||||
@@ -272,7 +293,8 @@ async fn process_agent_text(
|
||||
}
|
||||
}
|
||||
IncomingClientMessage::Result { request_id, result } => {
|
||||
let agent_id = authed_agent_id
|
||||
let agent_id = connection
|
||||
.authed_agent_id
|
||||
.as_deref()
|
||||
.ok_or_else(|| anyhow::anyhow!("result before auth"))?;
|
||||
ensure_current_session(state, agent_id, connection_id).await?;
|
||||
@@ -284,7 +306,8 @@ async fn process_agent_text(
|
||||
}
|
||||
}
|
||||
IncomingClientMessage::Error { request_id, error } => {
|
||||
let agent_id = authed_agent_id
|
||||
let agent_id = connection
|
||||
.authed_agent_id
|
||||
.as_deref()
|
||||
.ok_or_else(|| anyhow::anyhow!("error before auth"))?;
|
||||
ensure_current_session(state, agent_id, connection_id).await?;
|
||||
@@ -295,6 +318,26 @@ async fn process_agent_text(
|
||||
debug!(request_id = %key, "dropping unsolicited error from agent");
|
||||
}
|
||||
}
|
||||
IncomingClientMessage::TerminalRejected { terminal_id, error } => {
|
||||
let agent_id = connection
|
||||
.authed_agent_id
|
||||
.as_deref()
|
||||
.ok_or_else(|| anyhow::anyhow!("terminal rejection before auth"))?;
|
||||
ensure_current_session(state, agent_id, connection_id).await?;
|
||||
let error_json = serde_json::to_string(&TerminalControl::Error {
|
||||
code: error.code,
|
||||
message: error.message,
|
||||
})?;
|
||||
if let Err(code) = state
|
||||
.terminals
|
||||
.reject(terminal_id.as_str(), agent_id, error_json)
|
||||
.await
|
||||
{
|
||||
debug!(terminal_id = %terminal_id, agent_id, code, "dropping rejection for inactive terminal");
|
||||
return Ok(());
|
||||
}
|
||||
warn!(terminal_id = %terminal_id, agent_id, "agent rejected terminal request");
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -349,6 +392,7 @@ mod tests {
|
||||
AgentSession {
|
||||
connection_id: "conn-new".to_string(),
|
||||
tx,
|
||||
capabilities: Vec::new(),
|
||||
},
|
||||
);
|
||||
|
||||
|
||||
@@ -23,7 +23,7 @@ impl TerminalPty {
|
||||
pty.resize(Size::new(rows, cols))
|
||||
.context("failed to set initial PTY size")?;
|
||||
|
||||
let command = Command::new(program);
|
||||
let command = Command::new(program).kill_on_drop(true);
|
||||
let child = command
|
||||
.spawn(pts)
|
||||
.with_context(|| format!("failed to spawn {} in PTY", program.display()))?;
|
||||
@@ -37,6 +37,14 @@ impl TerminalPty {
|
||||
}
|
||||
}
|
||||
|
||||
/// Resizes an owned PTY writer without exposing `pty-process` protocol types
|
||||
/// to higher-level crates.
|
||||
pub fn resize_terminal(writer: &OwnedWritePty, rows: u16, cols: u16) -> Result<()> {
|
||||
writer
|
||||
.resize(Size::new(rows, cols))
|
||||
.context("failed to resize PTY")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
Reference in New Issue
Block a user