BetterDesk/web-nodejs/services/remoteRelay.js

/**
 * BetterDesk Console — Remote Desktop WebSocket Relay
 *
 * Bridges the BetterDesk desktop agent (JPEG screen stream) and the operator
 * browser viewer (Canvas element).
 *
 * Endpoints:
 *   WS /ws/remote-agent/<device_id>    — agent standby + stream source
 *   WS /ws/remote-viewer/<device_id>   — operator browser viewer
 *
 * Flow:
 *   1. Agent connects to /ws/remote-agent/<id> and sends { "type": "agent-ready" }
 *   2. Operator opens /remote/<id> in the web console (browser)
 *   3. Browser connects to /ws/remote-viewer/<id>
 *   4. Server sends { "type": "start" } to the agent
 *   5. Agent starts capturing and sends binary JPEG frames
 *   6. Server forwards binary frames to all connected viewers
 *   7. Viewer sends JSON input events → server → agent
 *   8. Session ends when viewer disconnects or sends { "type": "stop" }
 *
 * Security:
 *   - Agent identified by device_id only (no token — relies on network isolation
 *     and the fact that only registered devices can reach the server)
 *   - Viewer requires valid session cookie (admin or operator role)
 *   - Input events are forwarded verbatim — agent validates the whitelist
 *   - Max binary frame: 2 MB (covers 1920×1080 JPEG at high quality)
 *   - Max concurrent viewers per device: 5
 */

'use strict';

const WebSocket = require('ws');

const MAX_BINARY_FRAME  = 2 * 1024 * 1024; // 2 MB
const MAX_VIEWERS       = 5;
const PING_INTERVAL     = 20000; // ms
const AGENT_IDLE_TTL    = 90000; // close idle agent after 90 s of no viewer

const log = {
    info:  (...a) => console.log('[RemoteRelay]', ...a),
    warn:  (...a) => console.warn('[RemoteRelay]', ...a),
    error: (...a) => console.error('[RemoteRelay]', ...a),
};

// ---------------------------------------------------------------------------
//  State
// ---------------------------------------------------------------------------

// device_id → { agentWs, viewers: Set<WebSocket>, streaming: bool, idleTimer }
const sessions = new Map();

// ---------------------------------------------------------------------------
//  Helpers
// ---------------------------------------------------------------------------

function getSession(deviceId) {
    if (!sessions.has(deviceId)) {
        sessions.set(deviceId, { agentWs: null, viewers: new Set(), streaming: false, idleTimer: null });
    }
    return sessions.get(deviceId);
}

function sendJson(ws, data) {
    if (ws && ws.readyState === WebSocket.OPEN) {
        ws.send(JSON.stringify(data));
    }
}

function broadcast(viewers, data, isBinary = false) {
    viewers.forEach(ws => {
        if (ws.readyState === WebSocket.OPEN) {
            ws.send(data, { binary: isBinary });
        }
    });
}

function setupPing(ws) {
    const t = setInterval(() => {
        if (ws.readyState === WebSocket.OPEN) ws.ping();
        else clearInterval(t);
    }, PING_INTERVAL);
    ws.on('close', () => clearInterval(t));
}

function startAgentStream(session, deviceId) {
    if (!session.agentWs || session.streaming) return;
    session.streaming = true;
    clearTimeout(session.idleTimer);
    session.idleTimer = null;
    log.info(`Session ${deviceId}: sending START to agent`);
    sendJson(session.agentWs, { type: 'start' });
    // Notify viewers
    broadcast(session.viewers, JSON.stringify({ type: 'stream-started' }));
}

function stopAgentStream(session, deviceId, reason = 'no viewers') {
    if (!session.streaming) return;
    session.streaming = false;
    log.info(`Session ${deviceId}: sending STOP to agent (${reason})`);
    if (session.agentWs && session.agentWs.readyState === WebSocket.OPEN) {
        sendJson(session.agentWs, { type: 'stop' });
    }
    broadcast(session.viewers, JSON.stringify({ type: 'stream-stopped', reason }));
}

function scheduleIdleClose(session, deviceId) {
    clearTimeout(session.idleTimer);
    session.idleTimer = setTimeout(() => {
        if (session.viewers.size === 0 && session.agentWs) {
            log.info(`Session ${deviceId}: idle timeout — closing agent`);
            session.agentWs.close(1000, 'Idle timeout');
        }
    }, AGENT_IDLE_TTL);
}

// ---------------------------------------------------------------------------
//  Agent connection handler (/ws/remote-agent/:device_id)
// ---------------------------------------------------------------------------

function handleAgentConnection(ws, deviceId) {
    const session = getSession(deviceId);

    // Replace any previous stale agent connection
    if (session.agentWs && session.agentWs.readyState !== WebSocket.CLOSED) {
        session.agentWs.close(1001, 'New agent connected');
    }
    session.streaming = false;
    session.agentWs = ws;

    log.info(`Agent connected: ${deviceId}`);

    // Notify waiting viewers that the agent is ready
    broadcast(session.viewers, JSON.stringify({ type: 'agent-ready', device_id: deviceId }));

    // If viewers are already waiting, start immediately
    if (session.viewers.size > 0) {
        startAgentStream(session, deviceId);
    } else {
        scheduleIdleClose(session, deviceId);
    }

    setupPing(ws);

    ws.on('message', (data, isBinary) => {
        if (isBinary) {
            // Binary = JPEG frame — forward to all viewers
            if (data.length > MAX_BINARY_FRAME) return;
            broadcast(session.viewers, data, true);
        } else {
            // JSON control frame from agent
            let frame;
            try { frame = JSON.parse(data.toString()); } catch { return; }
            // Forward status events to viewers
            if (frame.type === 'pong' || frame.type === 'status') {
                broadcast(session.viewers, JSON.stringify(frame));
            }
        }
    });

    ws.on('close', () => {
        if (session.agentWs === ws) {
            session.agentWs = null;
            session.streaming = false;
            clearTimeout(session.idleTimer);
            log.info(`Agent disconnected: ${deviceId}`);
            broadcast(session.viewers, JSON.stringify({ type: 'agent-disconnected' }));
        }
    });

    ws.on('error', (err) => {
        log.warn(`Agent WS error for ${deviceId}: ${err.message}`);
    });
}

// ---------------------------------------------------------------------------
//  Viewer connection handler (/ws/remote-viewer/:device_id)
// ---------------------------------------------------------------------------

function handleViewerConnection(ws, deviceId, operatorName) {
    const session = getSession(deviceId);

    if (session.viewers.size >= MAX_VIEWERS) {
        ws.close(1008, 'Max viewers reached');
        return;
    }

    session.viewers.add(ws);
    log.info(`Viewer ${operatorName} joined session ${deviceId} (${session.viewers.size} total)`);

    // Notify viewer of current state
    const agentReady = !!session.agentWs && session.agentWs.readyState === WebSocket.OPEN;
    sendJson(ws, { type: 'session-info', agent_ready: agentReady, streaming: session.streaming });

    // Start agent stream if agent is ready and not yet streaming
    if (agentReady && !session.streaming) {
        startAgentStream(session, deviceId);
    }

    setupPing(ws);

    ws.on('message', (data, isBinary) => {
        if (isBinary) return; // Viewers only send JSON

        let frame;
        try { frame = JSON.parse(data.toString()); } catch { return; }

        switch (frame.type) {
            case 'stop':
                stopAgentStream(session, deviceId, 'operator requested stop');
                break;

            case 'input':
                // Transform viewer event into agent InputEvent format.
                // Viewer sends { type: 'input', event_type: 'mouse_move', x, y, ... }
                // Agent expects { type: 'mouse_move', x, y, ... } (Rust InputEvent struct)
                if (session.agentWs && session.agentWs.readyState === WebSocket.OPEN) {
                    const agentEvent = { ...frame, type: frame.event_type || frame.kind || 'unknown' };
                    delete agentEvent.event_type;
                    sendJson(session.agentWs, agentEvent);
                }
                break;

            default:
                break;
        }
    });

    ws.on('close', () => {
        session.viewers.delete(ws);
        log.info(`Viewer ${operatorName} left session ${deviceId} (${session.viewers.size} remaining)`);

        if (session.viewers.size === 0) {
            stopAgentStream(session, deviceId, 'last viewer left');
            scheduleIdleClose(session, deviceId);
        }
    });

    ws.on('error', (err) => {
        log.warn(`Viewer WS error for ${deviceId}: ${err.message}`);
    });
}

// ---------------------------------------------------------------------------
//  Init — attach to existing HTTP server
// ---------------------------------------------------------------------------

function initRemoteRelay(server, sessionMiddleware) {
    const wss = new WebSocket.Server({ noServer: true });
    const { enforceOrigin } = require('../middleware/wsOrigin');

    server.on('upgrade', (req, socket, head) => {
        const url    = new URL(req.url, `http://${req.headers.host}`);
        const path   = url.pathname;

        // Agent: /ws/remote-agent/<device_id>
        const agentMatch = path.match(/^\/ws\/remote-agent\/([^/]+)$/);
        if (agentMatch) {
            if (!enforceOrigin(req, socket, `remote-agent ${path}`)) return;
            const deviceId = decodeURIComponent(agentMatch[1]);
            // Validate device ID format (reject path traversal etc.)
            if (!/^[A-Za-z0-9_-]{3,64}$/.test(deviceId)) {
                socket.write('HTTP/1.1 400 Bad Request\r\n\r\n');
                socket.destroy();
                return;
            }
            wss.handleUpgrade(req, socket, head, (ws) => {
                handleAgentConnection(ws, deviceId);
            });
            return;
        }

        // Viewer (operator): /ws/remote-viewer/<device_id>
        const viewerMatch = path.match(/^\/ws\/remote-viewer\/([^/]+)$/);
        if (viewerMatch) {
            if (!enforceOrigin(req, socket, `remote-viewer ${path}`)) return;
            sessionMiddleware(req, {}, () => {
                if (!req.session || !req.session.userId) {
                    socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
                    socket.destroy();
                    return;
                }
                wss.handleUpgrade(req, socket, head, (ws) => {
                    const opName = req.session.username || 'operator';
                    handleViewerConnection(ws, viewerMatch[1], opName);
                });
            });
            return;
        }
    });

    log.info('Remote relay initialized (/ws/remote-agent/:id, /ws/remote-viewer/:id)');
    return wss;
}

// ---------------------------------------------------------------------------
//  Exports
// ---------------------------------------------------------------------------

module.exports = {
    initRemoteRelay,
    /** Get session state (for admin REST API) */
    getSessionState(deviceId) {
        const s = sessions.get(deviceId);
        if (!s) return null;
        return {
            agentConnected: !!s.agentWs && s.agentWs.readyState === WebSocket.OPEN,
            viewerCount: s.viewers.size,
            streaming: s.streaming,
        };
    },
    getActiveSessions: () => [...sessions.keys()],
};