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
): void {
  // Check if session exists
  const session = sessionManager.getSession(sessionId);
  if (!session) {
    socket.send(JSON.stringify({
      type: 'error',
      error: 'Session not found'
    }));
    socket.close();
    return;
  }

  // Add client (not authenticated yet)
  const client = connectionManager.addWebClient(sessionId, socket, false);

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

  // Check if desktop is connected
  const desktop = connectionManager.getDesktopBySessionId(sessionId);
  socket.send(JSON.stringify({
    type: 'desktop_status',
    status: desktop ? 'connected' : 'disconnected'
  }));

  // 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);
          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;
          }
          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
): Promise<boolean> {
  // Reject early if this session id is currently locked out.
  const lockRemaining = getLockoutRemaining(sessionId);
  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(sessionId);
    client.authenticated = true;
    socket.send(JSON.stringify({
      type: 'auth_response',
      success: true
    }));
    logger.debug(`Web client authenticated for session: ${sessionId}`);
    return false;
  } else {
    recordFailure(sessionId);
    socket.send(JSON.stringify({
      type: 'auth_response',
      success: false,
      error: 'Invalid password'
    }));
    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
    }));
  }
}