MP-Server/macropad-relay/src/server.ts

// Express + WebSocket Server Setup

import express from 'express';
import fs from 'fs';
import http from 'http';
import { WebSocketServer } from 'ws';
import cors from 'cors';
import helmet from 'helmet';
import rateLimit from 'express-rate-limit';
import path from 'path';

import { config } from './config';
import { logger } from './utils/logger';
import { SessionManager } from './services/SessionManager';
import { ConnectionManager } from './services/ConnectionManager';
import { handleDesktopConnection } from './handlers/desktopHandler';
import { handleWebClientConnection } from './handlers/webClientHandler';
import { createApiProxy } from './handlers/apiProxy';

// Resolve the client IP for an upgrade request, honouring the proxy's
// X-Forwarded-For first hop when we trust the proxy.
function getClientIp(request: http.IncomingMessage): string {
  if (config.trustProxy) {
    const xff = request.headers['x-forwarded-for'];
    if (typeof xff === 'string' && xff.length > 0) {
      return xff.split(',')[0].trim();
    }
    if (Array.isArray(xff) && xff.length > 0) {
      return xff[0].split(',')[0].trim();
    }
  }
  return request.socket.remoteAddress || 'unknown';
}

// Simple in-memory sliding-window rate limiter for WS upgrades, keyed by IP.
function createUpgradeRateLimiter() {
  const hits = new Map<string, number[]>();

  const cleanup = setInterval(() => {
    const cutoff = Date.now() - config.wsUpgradeWindowMs;
    for (const [ip, times] of hits) {
      const recent = times.filter((t) => t > cutoff);
      if (recent.length === 0) {
        hits.delete(ip);
      } else {
        hits.set(ip, recent);
      }
    }
  }, config.wsUpgradeWindowMs);
  if (typeof cleanup.unref === 'function') cleanup.unref();

  return function allow(ip: string): boolean {
    const now = Date.now();
    const cutoff = now - config.wsUpgradeWindowMs;
    const times = (hits.get(ip) || []).filter((t) => t > cutoff);
    times.push(now);
    hits.set(ip, times);
    return times.length <= config.wsUpgradeMax;
  };
}

// Validate the Origin header for a WS upgrade against the allowlist.
// Native (desktop) clients send no Origin and are permitted.
function isOriginAllowed(origin: string | undefined, host: string | undefined): boolean {
  if (!origin) return true; // non-browser clients (e.g. desktop app)

  if (config.allowedOrigins.length > 0) {
    return config.allowedOrigins.includes(origin);
  }

  // Default: allow only same-host origins.
  try {
    return new URL(origin).host === host;
  } catch {
    return false;
  }
}

export function createServer() {
  const app = express();
  const server = http.createServer(app);

  // Trust the reverse proxy so client IPs (rate limiting) are accurate.
  app.set('trust proxy', config.trustProxy);

  // Initialize managers
  const sessionManager = new SessionManager();
  const connectionManager = new ConnectionManager();

  // Middleware
  app.use(helmet({
    // The bundled login/app pages rely on inline <script> and inline event
    // handlers, so script/style inline is permitted; images may be blobs.
    contentSecurityPolicy: {
      directives: {
        defaultSrc: ["'self'"],
        scriptSrc: ["'self'", "'unsafe-inline'"],
        styleSrc: ["'self'", "'unsafe-inline'"],
        imgSrc: ["'self'", 'data:', 'blob:'],
        connectSrc: ["'self'", 'ws:', 'wss:'],
        manifestSrc: ["'self'", 'blob:'],
        objectSrc: ["'none'"],
        frameAncestors: ["'self'"],
      },
    },
  }));

  // Restrict CORS to the configured origins. When none are configured we
  // disallow cross-origin requests (the first-party app is same-origin).
  const corsOptions: cors.CorsOptions =
    config.allowedOrigins.length > 0
      ? { origin: config.allowedOrigins }
      : { origin: false };
  app.use(cors(corsOptions));
  app.use(express.json());

  // Rate limiting
  const limiter = rateLimit({
    windowMs: config.rateLimitWindowMs,
    max: config.rateLimitMax,
    message: { error: 'Too many requests, please try again later' }
  });
  app.use(limiter);

  // Static files - check both locations (dev vs production)
  const publicPath = fs.existsSync(path.join(__dirname, 'public'))
    ? path.join(__dirname, 'public')
    : path.join(__dirname, '..', 'public');
  app.use('/static', express.static(publicPath));

  // Health check - counts only, never leak session ids.
  app.get('/health', (req, res) => {
    const stats = connectionManager.getStats();
    res.json({
      status: 'ok',
      desktops: stats.desktopConnections,
      webClients: stats.webClients,
      uptime: process.uptime()
    });
  });

  // Ping endpoint for container health checks
  app.get('/ping', (req, res) => {
    res.json({ status: 'ok' });
  });

  // Index page
  app.get('/', (req, res) => {
    res.sendFile(path.join(publicPath, 'index.html'));
  });

  // Login page for session. Served uniformly regardless of whether the
  // session exists so validity cannot be probed via this endpoint.
  app.get('/:sessionId', (req, res) => {
    res.sendFile(path.join(publicPath, 'login.html'));
  });

  // PWA app page (served uniformly; auth is enforced by the API/WS layer).
  app.get('/:sessionId/app', (req, res) => {
    res.sendFile(path.join(publicPath, 'app.html'));
  });

  // API proxy routes
  const apiProxy = createApiProxy(connectionManager, sessionManager);
  app.all('/:sessionId/api/*', apiProxy);

  // WebSocket server
  const wss = new WebSocketServer({ noServer: true });

  const allowUpgrade = createUpgradeRateLimiter();

  // Handle HTTP upgrade for WebSocket
  server.on('upgrade', (request, socket, head) => {
    // IP-based upgrade rate limiting (WS upgrades bypass Express middleware).
    const ip = getClientIp(request);
    if (!allowUpgrade(ip)) {
      logger.warn(`WS upgrade rate limit exceeded for ${ip}`);
      socket.destroy();
      return;
    }

    // Origin validation.
    if (!isOriginAllowed(request.headers.origin, request.headers.host)) {
      logger.warn(`WS upgrade rejected: bad origin ${request.headers.origin}`);
      socket.destroy();
      return;
    }

    const url = new URL(request.url || '', `http://${request.headers.host}`);
    const pathname = url.pathname;

    // Desktop connection: /desktop
    if (pathname === '/desktop') {
      wss.handleUpgrade(request, socket, head, (ws) => {
        handleDesktopConnection(ws, connectionManager, sessionManager);
      });
      return;
    }

    // Web client connection: /:sessionId/ws
    const webClientMatch = pathname.match(/^\/([a-zA-Z0-9]+)\/ws$/);
    if (webClientMatch) {
      const sessionId = webClientMatch[1];
      wss.handleUpgrade(request, socket, head, (ws) => {
        handleWebClientConnection(ws, sessionId, connectionManager, sessionManager);
      });
      return;
    }

    // Invalid path
    socket.destroy();
  });

  // Graceful shutdown
  const shutdown = () => {
    logger.info('Shutting down...');
    connectionManager.shutdown();
    sessionManager.shutdown();
    server.close(() => {
      logger.info('Server closed');
      process.exit(0);
    });
  };

  process.on('SIGTERM', shutdown);
  process.on('SIGINT', shutdown);

  return { app, server, sessionManager, connectionManager };
}