ShieldAI/server/webrtc/signaling-server.ts

import { WebSocketServer, WebSocket, Data } from 'ws';
import { randomBytes } from 'crypto';
import { IncomingMessage } from 'http';
import jwt from 'jsonwebtoken';

/**
 * WebRTC Signaling Server
 * Handles offer/answer/ICE candidate exchange for P2P connections.
 *
 * Security hardening (FRE-4497):
 * - JWT authentication required on WebSocket upgrade
 * - Origin allowlist validation
 * - JSON schema validation for all messages
 * - Server-side peer identity (crypto.randomBytes)
 * - Message size limits to prevent DoS
 * - Connection timeout for idle peers
 */

// ── Types ────────────────────────────────────────────────────────────────────

export interface SignalingServerConfig {
  port: number;
  host: string;
  allowedOrigins: string[];
  jwtSecret: string;
  maxMessageSize: number;
  idleTimeoutMs: number;
  maxConnectionsPerPeer: number;
}

export interface SignalingMessage {
  type: 'offer' | 'answer' | 'ice-candidate' | 'ping' | 'pong' | 'close';
  payload?: Record<string, unknown>;
  targetPeerId?: string;
}

export interface PeerConnection {
  ws: WebSocket;
  peerId: string;
  authenticatedUserId: string;
  connections: Map<string, PeerSession>;
  lastActivity: number;
  iceCandidates: Array<Record<string, unknown>>;
}

export interface PeerSession {
  targetPeerId: string;
  dataChannelReady: boolean;
  bufferedCandidates: Array<Record<string, unknown>>;
}

// ── Constants ────────────────────────────────────────────────────────────────

const DEFAULT_CONFIG: SignalingServerConfig = {
  port: parseInt(process.env.SIGNALING_PORT || '3001', 10),
  host: process.env.SIGNALING_HOST || '0.0.0.0',
  allowedOrigins: (process.env.ALLOWED_ORIGINS || '').split(',').filter(Boolean),
  jwtSecret: process.env.JWT_SECRET || randomBytes(32).toString('hex'),
  maxMessageSize: 65536,
  idleTimeoutMs: 300_000,
  maxConnectionsPerPeer: 10,
};

// Message schema validators
const MESSAGE_TYPES = new Set(['offer', 'answer', 'ice-candidate', 'ping', 'pong', 'close']);

function validateMessage(raw: unknown): raw is SignalingMessage {
  if (typeof raw !== 'object' || raw === null) return false;
  const msg = raw as Record<string, unknown>;
  if (!MESSAGE_TYPES.has(msg.type as string)) return false;
  if (msg.payload && typeof msg.payload !== 'object') return false;
  if (msg.targetPeerId !== undefined && typeof msg.targetPeerId !== 'string') return false;
  if (msg.targetPeerId && msg.targetPeerId.length > 64) return false;
  return true;
}

// ── JWT Helper ───────────────────────────────────────────────────────────────

/**
 * Extract JWT from Authorization header (preferred) or URL query (fallback).
 * Header-first avoids token exposure in access logs.
 */
function extractJwt(req: IncomingMessage): string | null {
  const auth = req.headers['authorization'];
  if (auth?.startsWith('Bearer ')) return auth.slice(7);
  const match = req.url?.match(/[?&]token=([^&]+)/);
  return match ? decodeURIComponent(match[1]) : null;
}

function verifyJwt(token: string, secret: string): { sub: string; exp?: number } | null {
  try {
    const decoded = jwt.verify(token, secret, { algorithms: ['HS256'] });
    if (typeof decoded !== 'object' || !decoded.sub) return null;
    return {
      sub: String(decoded.sub),
      exp: decoded.exp ? Number(decoded.exp) : undefined,
    };
  } catch {
    return null;
  }
}

// ── Server ───────────────────────────────────────────────────────────────────

export class SignalingServer {
  private wss: WebSocketServer;
  private peers: Map<string, PeerConnection> = new Map();
  private config: SignalingServerConfig;
  private idleTimers: Map<string, NodeJS.Timeout> = new Map();

  constructor(config: Partial<SignalingServerConfig> = {}) {
    this.config = { ...DEFAULT_CONFIG, ...config };
    this.wss = new WebSocketServer({
      port: this.config.port,
      host: this.config.host,
      maxPayload: this.config.maxMessageSize,
      verifyClient: this.verifyClient.bind(this),
    });
    this.wss.on('connection', this.handleConnection.bind(this));
    console.log(`[Signaling] Server listening on ${this.config.host}:${this.config.port}`);
  }

  /**
   * Verify incoming WebSocket connection: origin + auth
   */
  private verifyClient(info: { req: IncomingMessage; origin: string }, cb: (result: boolean, status?: number, reason?: string) => void) {
    // Origin validation
    if (this.config.allowedOrigins.length > 0) {
      const origin = info.origin || info.req.headers['origin'] || '';
      const allowed = this.config.allowedOrigins.some(
        allowedOrigin => origin === allowedOrigin || origin.startsWith(allowedOrigin)
      );
      if (!allowed) {
        cb(false, 403, `Origin "${origin}" not in allowlist`);
        return;
      }
    }

    // JWT authentication
    const token = extractJwt(info.req);
    if (!token) {
      cb(false, 401, 'Missing JWT token');
      return;
    }

    const payload = verifyJwt(token, this.config.jwtSecret);
    if (!payload) {
      cb(false, 401, 'Invalid or expired JWT');
      return;
    }

    cb(true);
  }

  /**
   * Handle new WebSocket connection
   */
  private handleConnection(ws: WebSocket, req: IncomingMessage) {
    const token = extractJwt(req);
    const payload = token ? verifyJwt(token, this.config.jwtSecret) : null;
    const authenticatedUserId = payload?.sub || '';

    // Server-side peer identity (crypto random)
    const peerId = `peer_${randomBytes(8).toString('hex')}`;

    const peer: PeerConnection = {
      ws,
      peerId,
      authenticatedUserId,
      connections: new Map(),
      lastActivity: Date.now(),
      iceCandidates: [],
    };

    this.peers.set(peerId, peer);

    // Send handshake with assigned peer ID
    ws.send(JSON.stringify({
      type: 'handshake',
      payload: { peerId, message: 'Connected' },
    }));

    // Idle timeout
    const timer = setTimeout(() => {
      if (Date.now() - peer.lastActivity > this.config.idleTimeoutMs) {
        ws.close(1001, 'Idle timeout');
      }
    }, this.config.idleTimeoutMs);
    this.idleTimers.set(peerId, timer);

    ws.on('message', this.handleMessage(peer).bind(this));
    ws.on('close', () => this.handleDisconnect(peer));
    ws.on('error', (err) => {
      console.error(`[Signaling] Peer ${peerId} error:`, err.message);
      this.handleDisconnect(peer);
    });
  }

  /**
   * Handle incoming message from peer
   */
  private handleMessage(peer: PeerConnection) {
    return (data: Data) => {
      peer.lastActivity = Date.now();

      // Parse with size guard
      let raw: unknown;
      try {
        const str = data.toString();
        if (str.length > this.config.maxMessageSize) {
          peer.ws.send(JSON.stringify({ type: 'error', payload: { message: 'Message too large' } }));
          return;
        }
        raw = JSON.parse(str);
      } catch {
        peer.ws.send(JSON.stringify({ type: 'error', payload: { message: 'Invalid JSON' } }));
        return;
      }

      // Schema validation
      if (!validateMessage(raw)) {
        peer.ws.send(JSON.stringify({ type: 'error', payload: { message: 'Invalid message schema' } }));
        return;
      }

      const msg = raw as SignalingMessage;

      switch (msg.type) {
        case 'ping':
          peer.ws.send(JSON.stringify({ type: 'pong', payload: { timestamp: Date.now() } }));
          break;
        case 'offer':
          this.handleOffer(peer, msg);
          break;
        case 'answer':
          this.handleAnswer(peer, msg);
          break;
        case 'ice-candidate':
          this.handleIceCandidate(peer, msg);
          break;
        case 'close':
          peer.ws.close(1000, 'Peer requested close');
          break;
      }
    };
  }

  /**
   * Route offer to target peer
   */
  private handleOffer(source: PeerConnection, msg: SignalingMessage) {
    const targetId = msg.targetPeerId;
    if (!targetId) {
      source.ws.send(JSON.stringify({ type: 'error', payload: { message: 'Missing targetPeerId' } }));
      return;
    }

    // Enforce max connections
    if (source.connections.size >= this.config.maxConnectionsPerPeer) {
      source.ws.send(JSON.stringify({ type: 'error', payload: { message: 'Max connections reached' } }));
      return;
    }

    const target = this.peers.get(targetId);
    if (!target) {
      source.ws.send(JSON.stringify({ type: 'error', payload: { message: `Peer ${targetId} not found` } }));
      return;
    }

    // Register session
    const session: PeerSession = {
      targetPeerId: targetId,
      dataChannelReady: false,
      bufferedCandidates: [...source.iceCandidates],
    };
    source.connections.set(targetId, session);
    target.connections.set(source.peerId, {
      targetPeerId: source.peerId,
      dataChannelReady: false,
      bufferedCandidates: [],
    });

    // Forward offer to target
    target.ws.send(JSON.stringify({
      type: 'offer',
      payload: msg.payload,
      targetPeerId: source.peerId,
    }));

    // Send buffered ICE candidates if data channel is ready
    if (session.dataChannelReady) {
      for (const candidate of session.bufferedCandidates) {
        target.ws.send(JSON.stringify({
          type: 'ice-candidate',
          payload: candidate,
          targetPeerId: source.peerId,
        }));
      }
    }
  }

  /**
   * Route answer to target peer
   */
  private handleAnswer(source: PeerConnection, msg: SignalingMessage) {
    const targetId = msg.targetPeerId;
    if (!targetId) {
      source.ws.send(JSON.stringify({ type: 'error', payload: { message: 'Missing targetPeerId' } }));
      return;
    }

    const target = this.peers.get(targetId);
    if (!target) {
      source.ws.send(JSON.stringify({ type: 'error', payload: { message: `Peer ${targetId} not found` } }));
      return;
    }

    // Mark data channel as ready for buffered candidate delivery
    const session = source.connections.get(targetId);
    if (session) {
      session.dataChannelReady = true;
    }

    target.ws.send(JSON.stringify({
      type: 'answer',
      payload: msg.payload,
      targetPeerId: source.peerId,
    }));
  }

  /**
   * Route ICE candidate to target peer
   */
  private handleIceCandidate(source: PeerConnection, msg: SignalingMessage) {
    const targetId = msg.targetPeerId;
    if (!targetId) {
      source.ws.send(JSON.stringify({ type: 'error', payload: { message: 'Missing targetPeerId' } }));
      return;
    }

    const candidate = msg.payload as Record<string, unknown> | undefined;
    if (!candidate) return;

    // Buffer candidate if target session not ready yet
    const session = source.connections.get(targetId);
    if (session && !session.dataChannelReady) {
      source.iceCandidates.push(candidate);
      return;
    }

    const target = this.peers.get(targetId);
    if (!target) return;

    target.ws.send(JSON.stringify({
      type: 'ice-candidate',
      payload: candidate,
      targetPeerId: source.peerId,
    }));
  }

  /**
   * Handle peer disconnect
   */
  private handleDisconnect(peer: PeerConnection) {
    // Notify connected peers
    for (const [targetId, session] of peer.connections) {
      const target = this.peers.get(targetId);
      if (target) {
        target.ws.send(JSON.stringify({
          type: 'close',
          payload: { peerId: peer.peerId, reason: 'Remote peer disconnected' },
          targetPeerId: peer.peerId,
        }));
        target.connections.delete(peer.peerId);
      }
    }

    // Clear idle timer
    const timer = this.idleTimers.get(peer.peerId);
    if (timer) clearTimeout(timer);
    this.idleTimers.delete(peer.peerId);

    this.peers.delete(peer.peerId);
    console.log(`[Signaling] Peer ${peer.peerId} disconnected`);
  }

  /**
   * Graceful shutdown with timeout
   */
  async stop(timeoutMs = 5000): Promise<void> {
    return new Promise((resolve) => {
      const deadline = Date.now() + timeoutMs;

      for (const [peerId, peer] of this.peers) {
        const remaining = Math.max(100, deadline - Date.now());
        setTimeout(() => {
          peer.ws.close(1001, 'Server shutting down');
        }, remaining);
      }

      const serverTimer = setTimeout(() => {
        this.wss.close();
        resolve();
      }, timeoutMs);

      this.wss.close(() => {
        clearTimeout(serverTimer);
        resolve();
      });
    });
  }

  /**
   * Get server stats
   */
  getStats() {
    return {
      connectedPeers: this.peers.size,
      totalConnections: Array.from(this.peers.values()).reduce((sum, p) => sum + p.connections.size, 0),
    };
  }
}

export function createSignalingServer(config?: Partial<SignalingServerConfig>): SignalingServer {
  return new SignalingServer(config);
}