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

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
): void {
  let authenticatedSessionId: string | null = null;

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

      switch (message.type) {
        case 'auth':
          await handleAuth(socket, message, sessionManager, connectionManager, (sessionId) => {
            authenticatedSessionId = sessionId;
          });
          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);
    }
  });
}

async function handleAuth(
  socket: WebSocket,
  message: AuthMessage,
  sessionManager: SessionManager,
  connectionManager: ConnectionManager,
  setSessionId: (id: string) => void
): Promise<void> {
  try {
    let sessionId = message.sessionId;
    let session;

    if (sessionId) {
      // Validate existing session
      const valid = await sessionManager.validatePassword(sessionId, message.password);
      if (!valid) {
        socket.send(JSON.stringify({
          type: 'auth_response',
          success: false,
          error: 'Invalid session ID or password'
        }));
        socket.close();
        return;
      }
      session = sessionManager.getSession(sessionId);
    } else {
      // Create new session
      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'
      }));
      socket.close();
      return;
    }

    // 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}`);
  } catch (error) {
    logger.error('Desktop auth error:', error);
    socket.send(JSON.stringify({
      type: 'auth_response',
      success: false,
      error: 'Authentication failed'
    }));
    socket.close();
  }
}

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