From b2c0711676616c2db69c721dfd1c3f781d058083 Mon Sep 17 00:00:00 2001 From: yushen Date: Wed, 4 Feb 2026 13:27:29 +0800 Subject: [PATCH] feat(hub): add device verification with token + whitelist + user confirmation Add DeviceStore for managing one-time tokens and persistent device whitelist. Create async verify RPC handler that validates tokens and awaits Desktop user confirmation for first-time connections. Whitelisted devices (reconnections) pass through instantly. Add message guard to reject unverified device traffic. Co-Authored-By: Claude Opus 4.5 --- src/hub/device-store.ts | 106 +++++++++++++++++++++++++++++++++ src/hub/hub.ts | 48 ++++++++++++++- src/hub/rpc/dispatcher.ts | 6 +- src/hub/rpc/handlers/verify.ts | 45 ++++++++++++++ 4 files changed, 201 insertions(+), 4 deletions(-) create mode 100644 src/hub/device-store.ts create mode 100644 src/hub/rpc/handlers/verify.ts diff --git a/src/hub/device-store.ts b/src/hub/device-store.ts new file mode 100644 index 000000000..e015b15c8 --- /dev/null +++ b/src/hub/device-store.ts @@ -0,0 +1,106 @@ +import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs"; +import { join } from "node:path"; +import { DATA_DIR } from "../shared/index.js"; + +// ============ Types ============ + +interface TokenEntry { + token: string; + agentId: string; + expiresAt: number; +} + +export interface DeviceEntry { + deviceId: string; + agentId: string; + addedAt: number; +} + +// ============ Persistence ============ + +const DEVICES_DIR = join(DATA_DIR, "devices"); +const DEVICES_FILE = join(DEVICES_DIR, "whitelist.json"); + +function ensureDir(): void { + if (!existsSync(DEVICES_DIR)) { + mkdirSync(DEVICES_DIR, { recursive: true }); + } +} + +function loadDevices(): DeviceEntry[] { + if (!existsSync(DEVICES_FILE)) return []; + try { + return JSON.parse(readFileSync(DEVICES_FILE, "utf-8")) as DeviceEntry[]; + } catch { + return []; + } +} + +function saveDevices(devices: DeviceEntry[]): void { + ensureDir(); + writeFileSync(DEVICES_FILE, JSON.stringify(devices, null, 2), "utf-8"); +} + +// ============ DeviceStore ============ + +export class DeviceStore { + /** One-time tokens (in-memory only, not persisted) */ + private readonly tokens = new Map(); + /** Allowed device IDs (persisted to disk) */ + private readonly allowedDevices = new Map(); + + constructor() { + // Restore from persistent storage + for (const entry of loadDevices()) { + this.allowedDevices.set(entry.deviceId, entry); + } + } + + // ---- Token management ---- + + /** Register a one-time token (called when QR code is generated) */ + registerToken(token: string, agentId: string, expiresAt: number): void { + this.tokens.set(token, { token, agentId, expiresAt }); + } + + /** Validate and consume a token (one-time use). Returns agentId if valid, null otherwise. */ + consumeToken(token: string): { agentId: string } | null { + const entry = this.tokens.get(token); + if (!entry) return null; + // Always delete — consumed or expired + this.tokens.delete(token); + if (Date.now() > entry.expiresAt) return null; + return { agentId: entry.agentId }; + } + + // ---- Device whitelist ---- + + /** Add a device to the whitelist (called after token verification + user confirmation) */ + allowDevice(deviceId: string, agentId: string): void { + const entry: DeviceEntry = { deviceId, agentId, addedAt: Date.now() }; + this.allowedDevices.set(deviceId, entry); + this.persist(); + } + + /** Check if a device is in the whitelist */ + isAllowed(deviceId: string): { agentId: string } | null { + const entry = this.allowedDevices.get(deviceId); + return entry ? { agentId: entry.agentId } : null; + } + + /** Remove a device from the whitelist */ + revokeDevice(deviceId: string): boolean { + const deleted = this.allowedDevices.delete(deviceId); + if (deleted) this.persist(); + return deleted; + } + + /** List all whitelisted devices */ + listDevices(): DeviceEntry[] { + return Array.from(this.allowedDevices.values()); + } + + private persist(): void { + saveDevices(Array.from(this.allowedDevices.values())); + } +} diff --git a/src/hub/hub.ts b/src/hub/hub.ts index bd5241bfa..f18dbc486 100644 --- a/src/hub/hub.ts +++ b/src/hub/hub.ts @@ -21,6 +21,8 @@ import { createListAgentsHandler } from "./rpc/handlers/list-agents.js"; import { createCreateAgentHandler } from "./rpc/handlers/create-agent.js"; import { createDeleteAgentHandler } from "./rpc/handlers/delete-agent.js"; import { createUpdateGatewayHandler } from "./rpc/handlers/update-gateway.js"; +import { DeviceStore } from "./device-store.js"; +import { createVerifyHandler } from "./rpc/handlers/verify.js"; export class Hub { private readonly agents = new Map(); @@ -29,6 +31,8 @@ export class Hub { private readonly agentStreamCounters = new Map(); private readonly rpc: RpcDispatcher; private client: GatewayClient; + readonly deviceStore: DeviceStore; + private _onConfirmDevice: ((deviceId: string, agentId: string) => Promise) | null = null; url: string; readonly path: string; readonly hubId: string; @@ -42,8 +46,20 @@ export class Hub { this.url = url; this.path = path ?? "/ws"; this.hubId = getHubId(); + this.deviceStore = new DeviceStore(); this.rpc = new RpcDispatcher(); + this.rpc.register("verify", createVerifyHandler({ + hubId: this.hubId, + deviceStore: this.deviceStore, + onConfirmDevice: (deviceId, agentId) => { + if (!this._onConfirmDevice) { + // No UI confirm handler registered (CLI mode etc.) — auto-approve + return Promise.resolve(true); + } + return this._onConfirmDevice(deviceId, agentId); + }, + })); this.rpc.register("getAgentMessages", createGetAgentMessagesHandler()); this.rpc.register("getHubInfo", createGetHubInfoHandler(this)); this.rpc.register("listAgents", createListAgentsHandler(this)); @@ -101,10 +117,30 @@ export class Hub { // RPC request if (msg.action === RequestAction) { const payload = msg.payload as RequestPayload; + // verify RPC is always allowed (it IS the verification step) + if (payload.method === "verify") { + void this.handleRpc(msg.from, payload); + return; + } + // Other RPCs require verified device + if (!this.deviceStore.isAllowed(msg.from)) { + this.client.send(msg.from, ResponseAction, { + requestId: payload.requestId, + ok: false, + error: { code: "UNAUTHORIZED", message: "Device not verified" }, + }); + return; + } void this.handleRpc(msg.from, payload); return; } + // Non-RPC messages also require verified device + if (!this.deviceStore.isAllowed(msg.from)) { + console.warn(`[Hub] Rejected message from unverified device: ${msg.from}`); + return; + } + // Regular chat message const payload = msg.payload as { agentId?: string; content?: string } | undefined; const agentId = payload?.agentId; @@ -129,6 +165,16 @@ export class Hub { return client; } + /** Register a confirmation handler for new device connections (called by Desktop UI) */ + setConfirmHandler(handler: ((deviceId: string, agentId: string) => Promise) | null): void { + this._onConfirmDevice = handler; + } + + /** Register a one-time token for device verification (called when QR code is generated) */ + registerToken(token: string, agentId: string, expiresAt: number): void { + this.deviceStore.registerToken(token, agentId, expiresAt); + } + /** 重连到新的 Gateway 地址 */ reconnect(url: string): void { console.log(`[Hub] Reconnecting to ${url}`); @@ -234,7 +280,7 @@ export class Hub { private async handleRpc(from: string, request: RequestPayload): Promise { const { requestId, method } = request; try { - const result = await this.rpc.dispatch(method, request.params); + const result = await this.rpc.dispatch(method, request.params, from); this.client.send(from, ResponseAction, { requestId, ok: true, diff --git a/src/hub/rpc/dispatcher.ts b/src/hub/rpc/dispatcher.ts index 1484568db..f1ae07a7c 100644 --- a/src/hub/rpc/dispatcher.ts +++ b/src/hub/rpc/dispatcher.ts @@ -1,4 +1,4 @@ -export type RpcHandler = (params: unknown) => unknown | Promise; +export type RpcHandler = (params: unknown, from: string) => unknown | Promise; export class RpcError extends Error { constructor( @@ -22,11 +22,11 @@ export class RpcDispatcher { } /** Dispatch an RPC request to its handler */ - async dispatch(method: string, params: unknown): Promise { + async dispatch(method: string, params: unknown, from: string): Promise { const handler = this.handlers.get(method); if (!handler) { throw new RpcError("METHOD_NOT_FOUND", `Unknown RPC method: ${method}`); } - return handler(params); + return handler(params, from); } } diff --git a/src/hub/rpc/handlers/verify.ts b/src/hub/rpc/handlers/verify.ts new file mode 100644 index 000000000..bef35ca86 --- /dev/null +++ b/src/hub/rpc/handlers/verify.ts @@ -0,0 +1,45 @@ +import type { RpcHandler } from "../dispatcher.js"; +import { RpcError } from "../dispatcher.js"; +import type { DeviceStore } from "../../device-store.js"; + +interface VerifyContext { + hubId: string; + deviceStore: DeviceStore; + /** Called for first-time connections. Returns true if user approves, false if rejected. */ + onConfirmDevice: (deviceId: string, agentId: string) => Promise; +} + +interface VerifyParams { + token?: string; +} + +export function createVerifyHandler(ctx: VerifyContext): RpcHandler { + return async (params: unknown, from: string) => { + // 1. Already in whitelist → pass through (reconnection, no confirmation needed) + const allowed = ctx.deviceStore.isAllowed(from); + if (allowed) { + return { hubId: ctx.hubId, agentId: allowed.agentId }; + } + + // 2. Validate token + const { token } = (params ?? {}) as VerifyParams; + if (!token) { + throw new RpcError("UNAUTHORIZED", "Device not authorized"); + } + + const result = ctx.deviceStore.consumeToken(token); + if (!result) { + throw new RpcError("UNAUTHORIZED", "Invalid or expired token"); + } + + // 3. Token valid → await Desktop user confirmation + const confirmed = await ctx.onConfirmDevice(from, result.agentId); + if (!confirmed) { + throw new RpcError("REJECTED", "Connection rejected by user"); + } + + // 4. User confirmed → add to whitelist + ctx.deviceStore.allowDevice(from, result.agentId); + return { hubId: ctx.hubId, agentId: result.agentId }; + }; +}