diff --git a/server/src/__tests__/broadcast-service.test.ts b/server/src/__tests__/broadcast-service.test.ts index 4a168e3e..7843bb56 100644 --- a/server/src/__tests__/broadcast-service.test.ts +++ b/server/src/__tests__/broadcast-service.test.ts @@ -2,7 +2,9 @@ * BroadcastService Tests * Tests WebSocket broadcast functions for task changes and telemetry. */ -import { describe, it, expect, beforeEach } from 'vitest'; +import { describe, it, expect } from 'vitest'; +import type { AnyTelemetryEvent } from '@veritas-kanban/shared'; +import type { WebSocketServer } from 'ws'; import { initBroadcast, broadcastTaskChange, @@ -12,13 +14,21 @@ import { // Minimal mock WebSocket server function createMockWss() { const sentMessages: string[] = []; - const clients = new Set<{ readyState: number; send: (data: string) => void }>(); + const clients = new Set<{ + readyState: number; + auth?: { role: 'admin' | 'agent' | 'read-only'; isLocalhost: boolean; workspaceId?: string }; + send: (data: string) => void; + }>(); return { clients, - addClient(readyState = 1) { + addClient( + readyState = 1, + auth = { role: 'admin' as const, isLocalhost: false, workspaceId: 'local' } + ) { const client = { readyState, + auth, send: (data: string) => sentMessages.push(data), }; clients.add(client); @@ -28,13 +38,26 @@ function createMockWss() { }; } +function asWebSocketServer(wss: ReturnType): WebSocketServer { + return wss as unknown as WebSocketServer; +} + +function telemetryEvent(): AnyTelemetryEvent { + return { + type: 'run.started', + taskId: 'task_789', + agent: 'claude-code', + timestamp: '2024-01-01T00:00:00Z', + } as unknown as AnyTelemetryEvent; +} + describe('BroadcastService', () => { describe('broadcastTaskChange()', () => { it('should broadcast to all connected clients', () => { const wss = createMockWss(); wss.addClient(1); // OPEN wss.addClient(1); // OPEN - initBroadcast(wss as any); + initBroadcast(asWebSocketServer(wss)); broadcastTaskChange('created', 'task_123'); @@ -51,7 +74,7 @@ describe('BroadcastService', () => { wss.addClient(1); // OPEN wss.addClient(0); // CONNECTING wss.addClient(3); // CLOSED - initBroadcast(wss as any); + initBroadcast(asWebSocketServer(wss)); broadcastTaskChange('updated', 'task_456'); @@ -60,7 +83,7 @@ describe('BroadcastService', () => { it('should handle no connected clients gracefully', () => { const wss = createMockWss(); - initBroadcast(wss as any); + initBroadcast(asWebSocketServer(wss)); // Should not throw broadcastTaskChange('deleted'); @@ -70,7 +93,7 @@ describe('BroadcastService', () => { it('should support all change types', () => { const wss = createMockWss(); wss.addClient(1); - initBroadcast(wss as any); + initBroadcast(asWebSocketServer(wss)); const types = ['created', 'updated', 'deleted', 'archived', 'restored', 'reordered'] as const; for (const type of types) { @@ -79,22 +102,26 @@ describe('BroadcastService', () => { expect(wss.sentMessages).toHaveLength(6); }); + + it('should filter task events by client workspace', () => { + const wss = createMockWss(); + wss.addClient(1, { role: 'read-only', isLocalhost: false, workspaceId: 'local' }); + wss.addClient(1, { role: 'read-only', isLocalhost: false, workspaceId: 'other' }); + initBroadcast(asWebSocketServer(wss)); + + broadcastTaskChange('updated', 'task_456'); + + expect(wss.sentMessages).toHaveLength(1); + }); }); describe('broadcastTelemetryEvent()', () => { it('should broadcast telemetry events to all connected clients', () => { const wss = createMockWss(); wss.addClient(1); - initBroadcast(wss as any); + initBroadcast(asWebSocketServer(wss)); - const event = { - type: 'run.started', - taskId: 'task_789', - agent: 'claude-code', - timestamp: '2024-01-01T00:00:00Z', - } as any; - - broadcastTelemetryEvent(event); + broadcastTelemetryEvent(telemetryEvent()); expect(wss.sentMessages).toHaveLength(1); const msg = JSON.parse(wss.sentMessages[0]); @@ -102,17 +129,28 @@ describe('BroadcastService', () => { expect(msg.event.taskId).toBe('task_789'); }); + it('should filter telemetry events by read permission', () => { + const wss = createMockWss(); + wss.addClient(1, { role: 'read-only', isLocalhost: false, workspaceId: 'local' }); + wss.addClient(1, { role: 'agent', isLocalhost: false, workspaceId: 'local' }); + initBroadcast(asWebSocketServer(wss)); + + broadcastTelemetryEvent(telemetryEvent()); + + expect(wss.sentMessages).toHaveLength(1); + }); + it('should do nothing when wss is not initialized', () => { - initBroadcast(null as any); + initBroadcast(null as unknown as WebSocketServer); // Should not throw - broadcastTelemetryEvent({ type: 'run.started' } as any); + broadcastTelemetryEvent(telemetryEvent()); }); }); describe('initBroadcast()', () => { it('should accept a WebSocket server', () => { const wss = createMockWss(); - expect(() => initBroadcast(wss as any)).not.toThrow(); + expect(() => initBroadcast(asWebSocketServer(wss))).not.toThrow(); }); }); }); diff --git a/server/src/index.ts b/server/src/index.ts index 18fd0cba..ba9ffdc9 100644 --- a/server/src/index.ts +++ b/server/src/index.ts @@ -40,6 +40,7 @@ import { authorize, authorizeWrite, authenticateWebSocket, + type AuthPermission, validateWebSocketOrigin, getAuthStatus, checkAdminKeyStrength, @@ -59,6 +60,7 @@ import { healthRouter, apiHealthRouter, setHealthWss } from './routes/health.js' import { getPrometheusCollector } from './services/metrics/prometheus.js'; import { metricsCollector } from './middleware/metrics-collector.js'; import { getStorageTypeFromEnv, initStorage, shutdownStorage } from './storage/index.js'; +import { canReceiveWebSocketEvent } from './services/websocket-permissions.js'; const log = createLogger('server'); @@ -628,6 +630,32 @@ const agentSubscriptions = new Map>(); // Track chat subscriptions: sessionId -> Set of WebSocket clients const chatSubscriptions = new Map>(); +function getMessageWorkspaceId(message: Record): string { + return typeof message.workspaceId === 'string' && message.workspaceId.trim() + ? message.workspaceId + : 'local'; +} + +function sendWebSocketForbidden( + ws: HeartbeatWebSocket, + requestType: string, + permissions: AuthPermission[], + workspaceId: string +): void { + if (ws.readyState !== WebSocket.OPEN) return; + + ws.send( + JSON.stringify({ + type: 'error', + requestType, + code: 'FORBIDDEN', + message: 'Insufficient permissions', + required: permissions, + workspaceId, + }) + ); +} + // ---- Heartbeat: server pings every WS_HEARTBEAT_INTERVAL_MS ---- const heartbeatInterval = setInterval(() => { for (const client of wss.clients) { @@ -758,6 +786,13 @@ wss.on('connection', (ws: HeartbeatWebSocket, req) => { // Handle subscription to chat session if (message.type === 'chat:subscribe' && message.sessionId) { + const workspaceId = getMessageWorkspaceId(message); + const permissions: AuthPermission[] = ['task:read']; + if (!canReceiveWebSocketEvent(ws, { workspaceId, permissions })) { + sendWebSocketForbidden(ws, 'chat:subscribe', permissions, workspaceId); + return; + } + // Unsubscribe from previous chat session if (subscribedChatSession) { const subs = chatSubscriptions.get(subscribedChatSession); @@ -792,6 +827,13 @@ wss.on('connection', (ws: HeartbeatWebSocket, req) => { // Handle subscription to agent output if (message.type === 'subscribe' && message.taskId) { + const workspaceId = getMessageWorkspaceId(message); + const permissions: AuthPermission[] = ['task:read']; + if (!canReceiveWebSocketEvent(ws, { workspaceId, permissions })) { + sendWebSocketForbidden(ws, 'subscribe', permissions, workspaceId); + return; + } + // Unsubscribe from previous task if (subscribedTaskId) { const subs = agentSubscriptions.get(subscribedTaskId); diff --git a/server/src/routes/agent-status.ts b/server/src/routes/agent-status.ts index b0527c77..a44fb5a1 100644 --- a/server/src/routes/agent-status.ts +++ b/server/src/routes/agent-status.ts @@ -9,6 +9,7 @@ import { statusHistoryService, type AgentStatusState as HistoryStatusState, } from '../services/status-history-service.js'; +import { sendWebSocketEvent } from '../services/websocket-permissions.js'; import { createLogger } from '../lib/logger.js'; const log = createLogger('agent-status'); @@ -115,10 +116,7 @@ function broadcastAgentStatusChange(): void { const payload = JSON.stringify(message); wssRef.clients.forEach((client: WebSocket) => { - if (client.readyState === 1) { - // WebSocket.OPEN = 1 - client.send(payload); - } + sendWebSocketEvent(client, payload, { permissions: ['agent:read'] }); }); } diff --git a/server/src/services/broadcast-service.ts b/server/src/services/broadcast-service.ts index f976312c..3159903b 100644 --- a/server/src/services/broadcast-service.ts +++ b/server/src/services/broadcast-service.ts @@ -5,6 +5,10 @@ import { notifyChatMessage, type TaskContext, } from './clawdbot-webhook-service.js'; +import { + canReceiveWebSocketEvent, + type WebSocketDeliveryOptions, +} from './websocket-permissions.js'; /** * Simple broadcast service that sends task change events to all connected WebSocket clients. @@ -26,11 +30,13 @@ export function initBroadcast(wss: WebSocketServer): void { * * @param payload - Pre-serialized JSON string */ -function broadcastToClients(payload: string): void { +function broadcastToClients(payload: string, options: WebSocketDeliveryOptions = {}): void { if (!wssRef) return; const clients = Array.from(wssRef.clients); - const openClients = clients.filter((c) => c.readyState === 1); + const openClients = clients.filter( + (client) => client.readyState === 1 && canReceiveWebSocketEvent(client, options) + ); // For small client counts, send synchronously if (openClients.length <= BROADCAST_BATCH_SIZE) { @@ -97,7 +103,7 @@ export function broadcastTaskChange( const payload = JSON.stringify(message); - broadcastToClients(payload); + broadcastToClients(payload, { permissions: ['task:read'] }); // Also notify via webhook (fire-and-forget) notifyTaskChange(changeType, taskId, taskContext); @@ -119,7 +125,7 @@ export function broadcastChatMessage(sessionId: string, event: ChatBroadcastEven const payload = JSON.stringify(event); - broadcastToClients(payload); + broadcastToClients(payload, { permissions: ['task:read'] }); // Also notify via webhook (fire-and-forget) notifyChatMessage( @@ -147,7 +153,7 @@ export function broadcastSquadMessage(message: SquadMessage): void { const payload = JSON.stringify(event); - broadcastToClients(payload); + broadcastToClients(payload, { permissions: ['agent:read'] }); } /** @@ -164,7 +170,7 @@ export function broadcastTelemetryEvent(event: AnyTelemetryEvent): void { const payload = JSON.stringify(message); - broadcastToClients(payload); + broadcastToClients(payload, { permissions: ['telemetry:read'] }); } export interface BroadcastMessageEvent { @@ -194,7 +200,7 @@ export function broadcastNewMessage(broadcast: BroadcastMessageEvent['broadcast' const payload = JSON.stringify(message); - broadcastToClients(payload); + broadcastToClients(payload, { permissions: ['agent:read'] }); } export interface WorkflowStatusEvent { @@ -282,5 +288,5 @@ export function broadcastWorkflowStatus(run: { const payload = JSON.stringify(message); - broadcastToClients(payload); + broadcastToClients(payload, { permissions: ['workflow:read'] }); } diff --git a/server/src/services/websocket-permissions.ts b/server/src/services/websocket-permissions.ts new file mode 100644 index 00000000..111aaf7c --- /dev/null +++ b/server/src/services/websocket-permissions.ts @@ -0,0 +1,44 @@ +import type { WebSocket } from 'ws'; +import { + hasPermission, + type AuthenticatedWebSocket, + type AuthPermission, +} from '../middleware/auth.js'; + +const DEFAULT_WORKSPACE_ID = 'local'; +const WEBSOCKET_OPEN = 1; + +export interface WebSocketDeliveryOptions { + workspaceId?: string; + permissions?: AuthPermission[]; +} + +export function canReceiveWebSocketEvent( + client: WebSocket, + options: WebSocketDeliveryOptions = {} +): boolean { + const auth = (client as AuthenticatedWebSocket).auth; + if (!auth) return false; + + const eventWorkspaceId = options.workspaceId ?? DEFAULT_WORKSPACE_ID; + if (auth.role !== 'admin' && auth.workspaceId !== eventWorkspaceId) { + return false; + } + + const permissions = options.permissions ?? []; + if (permissions.length === 0) return true; + + return permissions.some((permission) => hasPermission(auth, permission)); +} + +export function sendWebSocketEvent( + client: WebSocket, + payload: string, + options: WebSocketDeliveryOptions = {} +): boolean { + if (client.readyState !== WEBSOCKET_OPEN) return false; + if (!canReceiveWebSocketEvent(client, options)) return false; + + client.send(payload); + return true; +}