MP-Server/macropad-relay/src/handlers/webClientHandler.ts

// Web Client WebSocket Handler

import WebSocket from 'ws';
import { ConnectionManager, WebClientConnection } from '../services/ConnectionManager';
import { SessionManager } from '../services/SessionManager';
import { logger } from '../utils/logger';
import { config } from '../config';
import { getLockoutRemaining, recordFailure, recordSuccess } from '../utils/authThrottle';

// Close code 1008 = policy violation (used for auth abuse).
const WS_POLICY_VIOLATION = 1008;

interface AuthMessage {
  type: 'auth';
  password: string;
}

interface PingMessage {
  type: 'ping';
}

type WebClientMessage = AuthMessage | PingMessage | any;

export function handleWebClientConnection(
  socket: WebSocket,
  sessionId: string,
  connectionManager: ConnectionManager,
  sessionManager: SessionManager,
  clientIp: string
): void {
  // Namespace the brute-force lockout key to the web keyspace and include the
  // client IP. This prevents one abusive IP (which knows the shareable session
  // id) from locking out other web clients or the desktop for this session.
  const throttleKey = `web:${sessionId}:${clientIp}`;

  // Do NOT reveal whether the session exists. A non-existent session id and
  // an existing one must be indistinguishable to an unauthenticated client so
  // session ids cannot be enumerated: we always run the same handshake and let
  // the (uniform) auth response be the only signal. validatePassword() safely
  // returns false for an unknown session id, and desktop connectivity is only
  // disclosed AFTER a successful auth (see the 'auth' case below).
  const client = connectionManager.addWebClient(sessionId, socket, false);

  // Track failed auth attempts on this individual socket.
  let socketAuthFailures = 0;

  // Request authentication.
  socket.send(JSON.stringify({
    type: 'auth_required'
  }));

  socket.on('message', async (data) => {
    try {
      const message: WebClientMessage = JSON.parse(data.toString());

      switch (message.type) {
        case 'auth': {
          const failed = await handleAuth(socket, client, message, sessionId, sessionManager, throttleKey);
          if (failed) {
            socketAuthFailures++;
            if (socketAuthFailures >= config.authMaxSocketFailures) {
              logger.warn(`Web client exceeded auth attempts for session ${sessionId}, closing socket`);
              socket.close(WS_POLICY_VIOLATION, 'Too many failed attempts');
            }
          } else {
            socketAuthFailures = 0;
            // Authenticated: only now do we disclose desktop connectivity.
            const desktop = connectionManager.getDesktopBySessionId(sessionId);
            socket.send(JSON.stringify({
              type: 'desktop_status',
              status: desktop ? 'connected' : 'disconnected'
            }));
          }
          break;
        }

        case 'ping':
          socket.send(JSON.stringify({ type: 'pong' }));
          break;

        default:
          // Forward other messages to desktop if authenticated
          if (client.authenticated) {
            forwardToDesktop(message, sessionId, connectionManager);
          }
      }
    } catch (error) {
      logger.error('Error handling web client message:', error);
    }
  });

  socket.on('close', () => {
    connectionManager.removeWebClient(sessionId, client);
  });

  socket.on('error', (error) => {
    logger.error('Web client WebSocket error:', error);
    connectionManager.removeWebClient(sessionId, client);
  });
}

// Returns true if the auth attempt failed (so the caller can count it).
async function handleAuth(
  socket: WebSocket,
  client: WebClientConnection,
  message: AuthMessage,
  sessionId: string,
  sessionManager: SessionManager,
  throttleKey: string
): Promise<boolean> {
  // Reject early if this key is currently locked out.
  const lockRemaining = getLockoutRemaining(throttleKey);
  if (lockRemaining > 0) {
    socket.send(JSON.stringify({
      type: 'auth_response',
      success: false,
      error: 'Too many failed attempts, try again later',
      retryAfterMs: lockRemaining
    }));
    return true;
  }

  const valid = await sessionManager.validatePassword(sessionId, message.password);

  if (valid) {
    recordSuccess(throttleKey);
    client.authenticated = true;
    socket.send(JSON.stringify({
      type: 'auth_response',
      success: true
    }));
    logger.debug(`Web client authenticated for session: ${sessionId}`);
    return false;
  } else {
    recordFailure(throttleKey);
    socket.send(JSON.stringify({
      type: 'auth_response',
      success: false,
      error: 'Authentication failed'
    }));
    return true;
  }
}

function forwardToDesktop(
  message: any,
  sessionId: string,
  connectionManager: ConnectionManager
): void {
  const desktop = connectionManager.getDesktopBySessionId(sessionId);
  if (desktop && desktop.socket.readyState === WebSocket.OPEN) {
    desktop.socket.send(JSON.stringify({
      type: 'ws_message',
      data: message
    }));
  }
}