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

// Desktop WebSocket Handler

import WebSocket from 'ws';
import { ConnectionManager } 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';
  sessionId: string | null;
  password: string;
}

interface ApiResponseMessage {
  type: 'api_response';
  requestId: string;
  status: number;
  body: any;
}

interface WsBroadcastMessage {
  type: 'ws_broadcast';
  data: any;
}

interface PongMessage {
  type: 'pong';
}

type DesktopMessage = AuthMessage | ApiResponseMessage | WsBroadcastMessage | PongMessage;

export function handleDesktopConnection(
  socket: WebSocket,
  connectionManager: ConnectionManager,
  sessionManager: SessionManager,
  clientIp: string
): void {
  let authenticatedSessionId: string | null = null;
  // Track failed auth attempts on this individual socket.
  let socketAuthFailures = 0;

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

      switch (message.type) {
        case 'auth': {
          const failed = await handleAuth(socket, message, sessionManager, connectionManager, clientIp, (sessionId) => {
            authenticatedSessionId = sessionId;
          });
          if (failed) {
            socketAuthFailures++;
            if (socketAuthFailures >= config.authMaxSocketFailures) {
              logger.warn('Desktop exceeded auth attempts, closing socket');
              socket.close(WS_POLICY_VIOLATION, 'Too many failed attempts');
            }
          } else {
            socketAuthFailures = 0;
          }
          break;
        }

        case 'api_response':
          if (authenticatedSessionId) {
            handleApiResponse(message, authenticatedSessionId, connectionManager);
          }
          break;

        case 'ws_broadcast':
          if (authenticatedSessionId) {
            handleWsBroadcast(message, authenticatedSessionId, connectionManager);
          }
          break;

        case 'pong':
          if (authenticatedSessionId) {
            connectionManager.updateDesktopPing(authenticatedSessionId);
          }
          break;

        default:
          logger.warn('Unknown message type from desktop:', (message as any).type);
      }
    } catch (error) {
      logger.error('Error handling desktop message:', error);
    }
  });

  socket.on('close', () => {
    if (authenticatedSessionId) {
      connectionManager.disconnectDesktop(authenticatedSessionId);
    }
  });

  socket.on('error', (error) => {
    logger.error('Desktop WebSocket error:', error);
    if (authenticatedSessionId) {
      connectionManager.disconnectDesktop(authenticatedSessionId);
    }
  });
}

// Returns true if the auth attempt failed (so the caller can count it).
async function handleAuth(
  socket: WebSocket,
  message: AuthMessage,
  sessionManager: SessionManager,
  connectionManager: ConnectionManager,
  clientIp: string,
  setSessionId: (id: string) => void
): Promise<boolean> {
  try {
    let sessionId = message.sessionId;
    let session;

    if (sessionId) {
      // Namespace the brute-force lockout key to the desktop keyspace and
      // include the client IP. This keeps the desktop's lockout state entirely
      // separate from the web-client keyspace, so a web-auth flood can never
      // lock the desktop out of (re)authenticating.
      const throttleKey = `desktop:${sessionId}:${clientIp}`;

      // 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;
      }

      // Validate existing session
      const valid = await sessionManager.validatePassword(sessionId, message.password);
      if (!valid) {
        recordFailure(throttleKey);
        socket.send(JSON.stringify({
          type: 'auth_response',
          success: false,
          error: 'Invalid session ID or password'
        }));
        return true;
      }
      recordSuccess(throttleKey);
      session = sessionManager.getSession(sessionId);
    } else {
      // Create new session (may throw on limit / short password)
      session = await sessionManager.createSession(message.password);
      sessionId = session.id;
    }

    if (!session || !sessionId) {
      socket.send(JSON.stringify({
        type: 'auth_response',
        success: false,
        error: 'Failed to create session'
      }));
      return true;
    }

    // Add to connection manager
    connectionManager.addDesktopConnection(sessionId, socket);
    setSessionId(sessionId);

    // Send success response
    socket.send(JSON.stringify({
      type: 'auth_response',
      success: true,
      sessionId: sessionId
    }));

    logger.info(`Desktop authenticated: ${sessionId}`);
    return false;
  } catch (error: any) {
    logger.error('Desktop auth error:', error);
    socket.send(JSON.stringify({
      type: 'auth_response',
      success: false,
      error: error?.message || 'Authentication failed'
    }));
    return true;
  }
}

function handleApiResponse(
  message: ApiResponseMessage,
  sessionId: string,
  connectionManager: ConnectionManager
): void {
  connectionManager.handleApiResponse(sessionId, message.requestId, message.status, message.body);
}

function handleWsBroadcast(
  message: WsBroadcastMessage,
  sessionId: string,
  connectionManager: ConnectionManager
): void {
  // Forward the broadcast to all web clients for this session
  connectionManager.broadcastToWebClients(sessionId, message.data);
}