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';

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);

  // 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':
          await handleAuth(socket, client, message, sessionId, sessionManager);
          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);
  });
}

async function handleAuth(
  socket: WebSocket,
  client: WebClientConnection,
  message: AuthMessage,
  sessionId: string,
  sessionManager: SessionManager
): Promise<void> {
  const valid = await sessionManager.validatePassword(sessionId, message.password);

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

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