cowork-os 0.3.21

package/dist/electron/electron/control-plane/server.js

@@ -0,0 +1,640 @@
+"use strict";
+/**
+ * WebSocket Control Plane Server
+ *
+ * The main WebSocket server that handles client connections, authentication,
+ * and message routing for the control plane.
+ */
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.ControlPlaneServer = void 0;
+const http_1 = __importDefault(require("http"));
+const ws_1 = require("ws");
+const crypto_1 = __importDefault(require("crypto"));
+const protocol_1 = require("./protocol");
+const client_1 = require("./client");
+const settings_1 = require("./settings");
+const tailscale_1 = require("../tailscale");
+/**
+ * WebSocket Control Plane Server
+ */
+class ControlPlaneServer {
+    constructor(config) {
+        this.httpServer = null;
+        this.wss = null;
+        this.methods = new Map();
+        this.heartbeatInterval = null;
+        this.cleanupInterval = null;
+        this.tailscaleCleanup = null;
+        // Rate limiting for auth attempts: Map<remoteAddress, { attempts: number, bannedUntil?: number }>
+        this.authAttempts = new Map();
+        this.config = {
+            port: config.port ?? 18789,
+            host: config.host ?? '127.0.0.1',
+            token: config.token,
+            handshakeTimeoutMs: config.handshakeTimeoutMs ?? 10000,
+            heartbeatIntervalMs: config.heartbeatIntervalMs ?? 30000,
+            cleanupIntervalMs: config.cleanupIntervalMs ?? 60000,
+            maxPayloadBytes: config.maxPayloadBytes ?? 10 * 1024 * 1024,
+            maxAuthAttempts: config.maxAuthAttempts ?? 5,
+            authBanDurationMs: config.authBanDurationMs ?? 5 * 60 * 1000, // 5 minutes
+            onEvent: config.onEvent ?? (() => { }),
+        };
+        this.clients = new client_1.ClientRegistry();
+        this.registerBuiltinMethods();
+    }
+    /**
+     * Check if the server is running
+     */
+    get isRunning() {
+        return this.httpServer !== null && this.wss !== null;
+    }
+    /**
+     * Get server address
+     */
+    getAddress() {
+        if (!this.httpServer)
+            return null;
+        const addr = this.httpServer.address();
+        if (typeof addr === 'string' || !addr)
+            return null;
+        return {
+            host: addr.address,
+            port: addr.port,
+            wsUrl: `ws://${addr.address}:${addr.port}`,
+        };
+    }
+    /**
+     * Start the control plane server
+     */
+    async start() {
+        if (this.isRunning) {
+            console.info('[ControlPlane] Server already running');
+            return;
+        }
+        return new Promise((resolve, reject) => {
+            // Create HTTP server for WebSocket upgrade
+            this.httpServer = http_1.default.createServer((req, res) => {
+                // Health check endpoint
+                if (req.url === '/health' && req.method === 'GET') {
+                    res.writeHead(200, { 'Content-Type': 'application/json' });
+                    res.end(JSON.stringify({
+                        status: 'ok',
+                        timestamp: Date.now(),
+                        clients: this.clients.count,
+                    }));
+                    return;
+                }
+                // Return 404 for other HTTP requests
+                res.writeHead(404, { 'Content-Type': 'application/json' });
+                res.end(JSON.stringify({ error: 'Not found' }));
+            });
+            // Create WebSocket server
+            this.wss = new ws_1.WebSocketServer({
+                server: this.httpServer,
+                maxPayload: this.config.maxPayloadBytes,
+            });
+            // Handle new connections
+            this.wss.on('connection', (socket, request) => {
+                this.handleConnection(socket, request);
+            });
+            this.wss.on('error', (error) => {
+                console.error('[ControlPlane] WebSocket server error:', error);
+                this.emitEvent({ action: 'error', timestamp: Date.now(), error: String(error) });
+            });
+            this.httpServer.on('error', (error) => {
+                console.error('[ControlPlane] HTTP server error:', error);
+                reject(error);
+            });
+            // Start listening
+            this.httpServer.listen(this.config.port, this.config.host, () => {
+                console.info(`[ControlPlane] Server listening on ws://${this.config.host}:${this.config.port}`);
+                this.emitEvent({ action: 'started', timestamp: Date.now() });
+                // Start heartbeat interval
+                this.startHeartbeat();
+                // Start cleanup interval
+                this.startCleanup();
+                resolve();
+            });
+        });
+    }
+    /**
+     * Start with Tailscale exposure
+     */
+    async startWithTailscale() {
+        const settings = settings_1.ControlPlaneSettingsManager.loadSettings();
+        // Start the WebSocket server first
+        await this.start();
+        // If Tailscale is configured, start exposure
+        if (settings.tailscale.mode !== 'off') {
+            const result = await (0, tailscale_1.startTailscaleExposure)({
+                mode: settings.tailscale.mode,
+                port: this.config.port,
+                resetOnExit: settings.tailscale.resetOnExit,
+                log: (msg) => console.log(msg),
+                warn: (msg) => console.warn(msg),
+            });
+            if (result.cleanup) {
+                this.tailscaleCleanup = result.cleanup;
+            }
+            return result;
+        }
+        return null;
+    }
+    /**
+     * Stop the control plane server
+     */
+    async stop() {
+        // Stop heartbeat
+        if (this.heartbeatInterval) {
+            clearInterval(this.heartbeatInterval);
+            this.heartbeatInterval = null;
+        }
+        // Stop cleanup
+        if (this.cleanupInterval) {
+            clearInterval(this.cleanupInterval);
+            this.cleanupInterval = null;
+        }
+        // Cleanup Tailscale
+        if (this.tailscaleCleanup) {
+            await this.tailscaleCleanup();
+            this.tailscaleCleanup = null;
+        }
+        // Broadcast shutdown event
+        this.clients.broadcast(protocol_1.Events.SHUTDOWN, { reason: 'Server stopping' });
+        // Close all client connections
+        this.clients.closeAll(1001, 'Server shutting down');
+        // Close WebSocket server
+        if (this.wss) {
+            this.wss.close();
+            this.wss = null;
+        }
+        // Close HTTP server
+        if (this.httpServer) {
+            return new Promise((resolve) => {
+                this.httpServer.close(() => {
+                    console.info('[ControlPlane] Server stopped');
+                    this.emitEvent({ action: 'stopped', timestamp: Date.now() });
+                    this.httpServer = null;
+                    resolve();
+                });
+            });
+        }
+    }
+    /**
+     * Register a method handler
+     */
+    registerMethod(method, handler) {
+        this.methods.set(method, handler);
+    }
+    /**
+     * Get server status
+     */
+    getStatus() {
+        return {
+            running: this.isRunning,
+            address: this.getAddress(),
+            clients: this.clients.getStatus(),
+            tailscale: (0, tailscale_1.getExposureStatus)(),
+        };
+    }
+    /**
+     * Broadcast an event to all authenticated clients
+     */
+    broadcast(event, payload) {
+        return this.clients.broadcast(event, payload);
+    }
+    // ===== Private Methods =====
+    /**
+     * Handle a new WebSocket connection
+     */
+    handleConnection(socket, request) {
+        const remoteAddress = request.headers['x-forwarded-for']?.split(',')[0]?.trim() ||
+            request.socket.remoteAddress ||
+            'unknown';
+        const userAgent = request.headers['user-agent'];
+        const origin = request.headers['origin'];
+        const client = new client_1.ControlPlaneClient(socket, remoteAddress, userAgent, origin);
+        this.clients.add(client);
+        console.info(`[ControlPlane] Client connected: ${client.id} from ${remoteAddress}`);
+        this.emitEvent({
+            action: 'client_connected',
+            timestamp: Date.now(),
+            clientId: client.id,
+        });
+        // Send challenge
+        client.sendChallenge();
+        // Set handshake timeout
+        const handshakeTimeout = setTimeout(() => {
+            if (!client.isAuthenticated) {
+                console.warn(`[ControlPlane] Handshake timeout for client ${client.id}`);
+                client.close(4008, 'Handshake timeout');
+            }
+        }, this.config.handshakeTimeoutMs);
+        // Handle messages
+        socket.on('message', async (data) => {
+            try {
+                const message = data.toString();
+                await this.handleMessage(client, message);
+            }
+            catch (error) {
+                console.error(`[ControlPlane] Message handling error:`, error);
+            }
+        });
+        // Handle close
+        socket.on('close', (code, reason) => {
+            clearTimeout(handshakeTimeout);
+            // If this was a node, broadcast disconnection event to operators
+            if (client.isNode) {
+                const nodeInfo = client.getNodeInfo();
+                this.clients.broadcastToOperators(protocol_1.Events.NODE_DISCONNECTED, {
+                    nodeId: client.id,
+                    node: nodeInfo,
+                });
+                console.info(`[ControlPlane] Node disconnected: ${client.id} (${nodeInfo?.displayName || 'unnamed'}) (code: ${code})`);
+            }
+            else {
+                console.info(`[ControlPlane] Client disconnected: ${client.id} (code: ${code})`);
+            }
+            this.clients.remove(client.id);
+            this.emitEvent({
+                action: 'client_disconnected',
+                timestamp: Date.now(),
+                clientId: client.id,
+                details: { code, reason: reason.toString(), wasNode: client.isNode },
+            });
+        });
+        // Handle error
+        socket.on('error', (error) => {
+            console.error(`[ControlPlane] Client error (${client.id}):`, error);
+        });
+    }
+    /**
+     * Handle an incoming message from a client
+     */
+    async handleMessage(client, message) {
+        const frame = (0, protocol_1.parseFrame)(message);
+        if (!frame) {
+            console.warn(`[ControlPlane] Invalid frame from ${client.id}`);
+            return;
+        }
+        client.updateActivity();
+        // Only handle request frames
+        if (frame.type !== protocol_1.FrameType.Request) {
+            return;
+        }
+        const request = frame;
+        // Handle connect method (authentication)
+        if (request.method === protocol_1.Methods.CONNECT) {
+            await this.handleConnect(client, request);
+            return;
+        }
+        // All other methods require authentication
+        if (!client.isAuthenticated) {
+            client.send((0, protocol_1.createErrorResponse)(request.id, protocol_1.ErrorCodes.UNAUTHORIZED, 'Authentication required'));
+            return;
+        }
+        // Route to method handler
+        await this.handleRequest(client, request);
+    }
+    /**
+     * Handle connect/authentication request
+     */
+    async handleConnect(client, request) {
+        const remoteAddress = client.info.remoteAddress;
+        // Check if IP is banned due to too many failed attempts
+        const authRecord = this.authAttempts.get(remoteAddress);
+        if (authRecord?.bannedUntil && authRecord.bannedUntil > Date.now()) {
+            const remainingMs = authRecord.bannedUntil - Date.now();
+            console.warn(`[ControlPlane] Auth blocked for ${remoteAddress}: banned for ${Math.ceil(remainingMs / 1000)}s`);
+            client.send((0, protocol_1.createErrorResponse)(request.id, protocol_1.ErrorCodes.UNAUTHORIZED, `Too many failed attempts. Try again in ${Math.ceil(remainingMs / 1000)} seconds.`));
+            client.close(4029, 'Rate limited');
+            return;
+        }
+        const params = request.params;
+        // Verify token
+        const providedToken = params?.token || '';
+        if (!this.verifyToken(providedToken)) {
+            // Track failed attempt
+            this.recordFailedAuth(remoteAddress);
+            client.reject();
+            client.send((0, protocol_1.createErrorResponse)(request.id, protocol_1.ErrorCodes.UNAUTHORIZED, 'Invalid token'));
+            client.close(4001, 'Authentication failed');
+            return;
+        }
+        // Clear auth attempts on success
+        this.authAttempts.delete(remoteAddress);
+        // Check if this is a node (mobile companion) connection
+        const isNode = params?.role === 'node';
+        if (isNode) {
+            // Authenticate as a node
+            const platform = (params?.client?.platform || 'ios');
+            const capabilities = (params?.capabilities || []);
+            const commands = params?.commands || [];
+            const permissions = params?.permissions || {};
+            client.authenticateAsNode({
+                deviceName: params?.client?.displayName || params?.deviceName,
+                platform,
+                version: params?.client?.version || '0.0.0',
+                deviceId: params?.client?.id,
+                modelIdentifier: params?.client?.modelIdentifier,
+                capabilities,
+                commands,
+                permissions,
+            });
+            console.info(`[ControlPlane] Node authenticated: ${client.id} (${params?.client?.displayName || 'unnamed'}) [${platform}]`);
+            this.emitEvent({
+                action: 'client_authenticated',
+                timestamp: Date.now(),
+                clientId: client.id,
+                details: {
+                    deviceName: params?.client?.displayName,
+                    role: 'node',
+                    platform,
+                    capabilities,
+                },
+            });
+            // Broadcast node connected event to operators
+            this.clients.broadcastToOperators(protocol_1.Events.NODE_CONNECTED, {
+                nodeId: client.id,
+                node: client.getNodeInfo(),
+            });
+            // Send success response
+            client.send((0, protocol_1.createResponseFrame)(request.id, {
+                clientId: client.id,
+                role: 'node',
+                scopes: ['read'],
+            }));
+        }
+        else {
+            // Authenticate as operator with admin scope
+            const scopes = ['admin'];
+            client.authenticate(scopes, params?.deviceName);
+            console.info(`[ControlPlane] Client authenticated: ${client.id} (${params?.deviceName || 'unnamed'})`);
+            this.emitEvent({
+                action: 'client_authenticated',
+                timestamp: Date.now(),
+                clientId: client.id,
+                details: { deviceName: params?.deviceName, role: 'operator' },
+            });
+            // Send success response
+            client.send((0, protocol_1.createResponseFrame)(request.id, {
+                clientId: client.id,
+                role: 'operator',
+                scopes,
+            }));
+        }
+        // Send connect success event
+        client.sendEvent(protocol_1.Events.CONNECT_SUCCESS, {
+            clientId: client.id,
+            serverVersion: '1.0.0',
+        });
+    }
+    /**
+     * Record a failed authentication attempt for rate limiting
+     */
+    recordFailedAuth(remoteAddress) {
+        const record = this.authAttempts.get(remoteAddress) || { attempts: 0 };
+        record.attempts++;
+        if (record.attempts >= this.config.maxAuthAttempts) {
+            record.bannedUntil = Date.now() + this.config.authBanDurationMs;
+            console.warn(`[ControlPlane] IP ${remoteAddress} banned for ${this.config.authBanDurationMs / 1000}s after ${record.attempts} failed attempts`);
+        }
+        this.authAttempts.set(remoteAddress, record);
+    }
+    /**
+     * Handle an authenticated request
+     */
+    async handleRequest(client, request) {
+        const handler = this.methods.get(request.method);
+        this.emitEvent({
+            action: 'request',
+            timestamp: Date.now(),
+            clientId: client.id,
+            method: request.method,
+        });
+        if (!handler) {
+            client.send((0, protocol_1.createErrorResponse)(request.id, protocol_1.ErrorCodes.UNKNOWN_METHOD, `Unknown method: ${request.method}`));
+            return;
+        }
+        try {
+            const result = await handler(client, request.params);
+            client.send((0, protocol_1.createResponseFrame)(request.id, result));
+        }
+        catch (error) {
+            console.error(`[ControlPlane] Method error (${request.method}):`, error);
+            client.send((0, protocol_1.createErrorResponse)(request.id, protocol_1.ErrorCodes.METHOD_FAILED, error.message || 'Method execution failed', error.details));
+        }
+    }
+    /**
+     * Verify authentication token
+     */
+    verifyToken(provided) {
+        if (!this.config.token || !provided)
+            return false;
+        const expected = Buffer.from(this.config.token);
+        const actual = Buffer.from(provided);
+        if (expected.length !== actual.length)
+            return false;
+        return crypto_1.default.timingSafeEqual(expected, actual);
+    }
+    /**
+     * Register built-in method handlers
+     */
+    registerBuiltinMethods() {
+        // Ping/health check
+        this.registerMethod(protocol_1.Methods.PING, async () => ({
+            pong: true,
+            timestamp: Date.now(),
+        }));
+        this.registerMethod(protocol_1.Methods.HEALTH, async () => ({
+            status: 'ok',
+            timestamp: Date.now(),
+            uptime: process.uptime(),
+        }));
+        // Status
+        this.registerMethod(protocol_1.Methods.STATUS, async () => this.getStatus());
+        // ===== Node (Mobile Companion) Methods =====
+        // List connected nodes
+        this.registerMethod(protocol_1.Methods.NODE_LIST, async () => {
+            return {
+                nodes: this.clients.getNodeInfoList(),
+            };
+        });
+        // Describe a specific node
+        this.registerMethod(protocol_1.Methods.NODE_DESCRIBE, async (client, params) => {
+            const { nodeId } = params;
+            if (!nodeId) {
+                throw { code: protocol_1.ErrorCodes.INVALID_PARAMS, message: 'nodeId is required' };
+            }
+            const node = this.clients.getNodeByIdOrName(nodeId);
+            if (!node) {
+                throw { code: protocol_1.ErrorCodes.NODE_NOT_FOUND, message: `Node not found: ${nodeId}` };
+            }
+            return {
+                node: node.getNodeInfo(),
+            };
+        });
+        // Invoke a command on a node
+        this.registerMethod(protocol_1.Methods.NODE_INVOKE, async (client, params) => {
+            const { nodeId, command, params: commandParams, timeoutMs = 30000 } = params;
+            if (!nodeId) {
+                throw { code: protocol_1.ErrorCodes.INVALID_PARAMS, message: 'nodeId is required' };
+            }
+            if (!command) {
+                throw { code: protocol_1.ErrorCodes.INVALID_PARAMS, message: 'command is required' };
+            }
+            const node = this.clients.getNodeByIdOrName(nodeId);
+            if (!node) {
+                throw { code: protocol_1.ErrorCodes.NODE_NOT_FOUND, message: `Node not found: ${nodeId}` };
+            }
+            // Check if node supports this command
+            const nodeInfo = node.getNodeInfo();
+            if (!nodeInfo?.commands.includes(command)) {
+                throw {
+                    code: protocol_1.ErrorCodes.NODE_COMMAND_FAILED,
+                    message: `Node does not support command: ${command}`,
+                };
+            }
+            // Check if node is in foreground (required for most commands)
+            if (!nodeInfo.isForeground && ['camera.snap', 'camera.clip', 'screen.record'].includes(command)) {
+                throw {
+                    code: protocol_1.ErrorCodes.NODE_BACKGROUND_UNAVAILABLE,
+                    message: 'Node app must be in foreground for this command',
+                };
+            }
+            // Forward the command to the node
+            return await this.invokeNodeCommand(node, command, commandParams, timeoutMs);
+        });
+        // Handle node events (from nodes to gateway)
+        this.registerMethod(protocol_1.Methods.NODE_EVENT, async (client, params) => {
+            if (!client.isNode) {
+                throw { code: protocol_1.ErrorCodes.UNAUTHORIZED, message: 'Only nodes can send node events' };
+            }
+            const { event, payload } = params;
+            if (!event) {
+                throw { code: protocol_1.ErrorCodes.INVALID_PARAMS, message: 'event is required' };
+            }
+            // Handle specific node events
+            if (event === 'foreground_changed') {
+                const isForeground = payload?.isForeground ?? true;
+                client.setForeground(isForeground);
+                this.clients.broadcastToOperators(protocol_1.Events.NODE_EVENT, {
+                    nodeId: client.id,
+                    event: 'foreground_changed',
+                    isForeground,
+                });
+            }
+            else if (event === 'capabilities_changed') {
+                const { capabilities, commands, permissions } = payload;
+                if (capabilities && commands && permissions) {
+                    client.updateCapabilities(capabilities, commands, permissions);
+                    this.clients.broadcastToOperators(protocol_1.Events.NODE_CAPABILITIES_CHANGED, {
+                        nodeId: client.id,
+                        node: client.getNodeInfo(),
+                    });
+                }
+            }
+            return { ok: true };
+        });
+    }
+    /**
+     * Invoke a command on a node and wait for response
+     */
+    async invokeNodeCommand(node, command, params, timeoutMs) {
+        return new Promise((resolve) => {
+            const requestId = crypto_1.default.randomUUID();
+            let timeoutHandle;
+            // Set up one-time response handler
+            const handleResponse = (data) => {
+                try {
+                    const message = data.toString();
+                    const frame = (0, protocol_1.parseFrame)(message);
+                    if (frame && frame.type === protocol_1.FrameType.Response && frame.id === requestId) {
+                        clearTimeout(timeoutHandle);
+                        node.info.socket.removeListener('message', handleResponse);
+                        const response = frame;
+                        if (response.ok) {
+                            resolve({ ok: true, payload: response.payload });
+                        }
+                        else {
+                            resolve({
+                                ok: false,
+                                error: response.error || { code: 'UNKNOWN', message: 'Command failed' },
+                            });
+                        }
+                    }
+                }
+                catch {
+                    // Ignore parse errors
+                }
+            };
+            node.info.socket.on('message', handleResponse);
+            // Set timeout
+            timeoutHandle = setTimeout(() => {
+                node.info.socket.removeListener('message', handleResponse);
+                resolve({
+                    ok: false,
+                    error: { code: protocol_1.ErrorCodes.NODE_TIMEOUT, message: `Command timed out after ${timeoutMs}ms` },
+                });
+            }, timeoutMs);
+            // Send command to node
+            const requestFrame = {
+                type: protocol_1.FrameType.Request,
+                id: requestId,
+                method: 'node.invoke',
+                params: { command, params },
+            };
+            node.info.socket.send(JSON.stringify(requestFrame));
+        });
+    }
+    /**
+     * Start heartbeat interval
+     */
+    startHeartbeat() {
+        this.heartbeatInterval = setInterval(() => {
+            const event = (0, protocol_1.createEventFrame)(protocol_1.Events.HEARTBEAT, {
+                timestamp: Date.now(),
+                clients: this.clients.count,
+            });
+            for (const client of this.clients.getAuthenticated()) {
+                client.send(event);
+                client.updateHeartbeat();
+            }
+        }, this.config.heartbeatIntervalMs);
+    }
+    /**
+     * Start cleanup interval
+     */
+    startCleanup() {
+        this.cleanupInterval = setInterval(() => {
+            const removed = this.clients.cleanup();
+            if (removed > 0) {
+                console.info(`[ControlPlane] Cleaned up ${removed} disconnected clients`);
+            }
+            // Also clean up expired auth bans
+            const now = Date.now();
+            for (const [ip, record] of this.authAttempts) {
+                if (record.bannedUntil && record.bannedUntil < now) {
+                    this.authAttempts.delete(ip);
+                }
+            }
+        }, this.config.cleanupIntervalMs);
+    }
+    /**
+     * Emit a server event
+     */
+    emitEvent(event) {
+        if (this.config.onEvent) {
+            try {
+                this.config.onEvent(event);
+            }
+            catch (error) {
+                console.error('[ControlPlane] Event handler error:', error);
+            }
+        }
+    }
+}
+exports.ControlPlaneServer = ControlPlaneServer;