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

// Express + WebSocket Server Setup

import express from 'express';
import fs from 'fs';
import http from 'http';
import WebSocket, { 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';

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

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

  // Middleware
  app.use(helmet({
    contentSecurityPolicy: false // Allow inline scripts for login page
  }));
  app.use(cors());
  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
  app.get('/health', (req, res) => {
    const stats = connectionManager.getStats();
    res.json({
      status: 'ok',
      ...stats
    });
  });

  // 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
  app.get('/:sessionId', (req, res) => {
    const session = sessionManager.getSession(req.params.sessionId);
    if (!session) {
      return res.status(404).send('Session not found');
    }
    res.sendFile(path.join(publicPath, 'login.html'));
  });

  // PWA app page (after authentication)
  app.get('/:sessionId/app', (req, res) => {
    const session = sessionManager.getSession(req.params.sessionId);
    if (!session) {
      return res.status(404).send('Session not found');
    }
    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 });

  // Handle HTTP upgrade for WebSocket
  server.on('upgrade', (request, socket, head) => {
    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();
    server.close(() => {
      logger.info('Server closed');
      process.exit(0);
    });
  };

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

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