smartproxy/ts/core/utils/socket-utils.ts

import * as plugins from '../../plugins.js';

export interface CleanupOptions {
  immediate?: boolean;      // Force immediate destruction
  allowDrain?: boolean;     // Allow write buffer to drain
  gracePeriod?: number;     // Ms to wait before force close
}

export interface SafeSocketOptions {
  port: number;
  host: string;
  onError?: (error: Error) => void;
  onConnect?: () => void;
  timeout?: number;
}

/**
 * Safely cleanup a socket by removing all listeners and destroying it
 * @param socket The socket to cleanup
 * @param socketName Optional name for logging
 * @param options Cleanup options
 */
export function cleanupSocket(
  socket: plugins.net.Socket | plugins.tls.TLSSocket | null, 
  socketName?: string,
  options: CleanupOptions = {}
): Promise<void> {
  if (!socket || socket.destroyed) return Promise.resolve();
  
  return new Promise<void>((resolve) => {
    const cleanup = () => {
      try {
        // Remove all event listeners
        socket.removeAllListeners();
        
        // Destroy if not already destroyed
        if (!socket.destroyed) {
          socket.destroy();
        }
      } catch (err) {
        console.error(`Error cleaning up socket${socketName ? ` (${socketName})` : ''}: ${err}`);
      }
      resolve();
    };
    
    if (options.immediate) {
      // Immediate cleanup (old behavior)
      socket.unpipe();
      cleanup();
    } else if (options.allowDrain && socket.writable) {
      // Allow pending writes to complete
      socket.end(() => cleanup());
      
      // Force cleanup after grace period
      if (options.gracePeriod) {
        setTimeout(() => {
          if (!socket.destroyed) {
            cleanup();
          }
        }, options.gracePeriod);
      }
    } else {
      // Default: immediate cleanup
      socket.unpipe();
      cleanup();
    }
  });
}

/**
 * Create a cleanup handler for paired sockets (client and server)
 * @param clientSocket The client socket
 * @param serverSocket The server socket (optional)
 * @param onCleanup Optional callback when cleanup is done
 * @returns A cleanup function that can be called multiple times safely
 * @deprecated Use createIndependentSocketHandlers for better half-open support
 */
export function createSocketCleanupHandler(
  clientSocket: plugins.net.Socket | plugins.tls.TLSSocket,
  serverSocket?: plugins.net.Socket | plugins.tls.TLSSocket | null,
  onCleanup?: (reason: string) => void
): (reason: string) => void {
  let cleanedUp = false;
  
  return (reason: string) => {
    if (cleanedUp) return;
    cleanedUp = true;
    
    // Cleanup both sockets (old behavior - too aggressive)
    cleanupSocket(clientSocket, 'client', { immediate: true });
    if (serverSocket) {
      cleanupSocket(serverSocket, 'server', { immediate: true });
    }
    
    // Call cleanup callback if provided
    if (onCleanup) {
      onCleanup(reason);
    }
  };
}

/**
 * Create independent cleanup handlers for paired sockets that support half-open connections
 * @param clientSocket The client socket
 * @param serverSocket The server socket
 * @param onBothClosed Callback when both sockets are closed
 * @returns Independent cleanup functions for each socket
 */
export function createIndependentSocketHandlers(
  clientSocket: plugins.net.Socket | plugins.tls.TLSSocket,
  serverSocket: plugins.net.Socket | plugins.tls.TLSSocket,
  onBothClosed: (reason: string) => void
): { cleanupClient: (reason: string) => Promise<void>, cleanupServer: (reason: string) => Promise<void> } {
  let clientClosed = false;
  let serverClosed = false;
  let clientReason = '';
  let serverReason = '';
  
  const checkBothClosed = () => {
    if (clientClosed && serverClosed) {
      onBothClosed(`client: ${clientReason}, server: ${serverReason}`);
    }
  };
  
  const cleanupClient = async (reason: string) => {
    if (clientClosed) return;
    clientClosed = true;
    clientReason = reason;
    
    // Allow server to continue if still active
    if (!serverClosed && serverSocket.writable) {
      // Half-close: stop reading from client, let server finish
      clientSocket.pause();
      clientSocket.unpipe(serverSocket);
      await cleanupSocket(clientSocket, 'client', { allowDrain: true, gracePeriod: 5000 });
    } else {
      await cleanupSocket(clientSocket, 'client', { immediate: true });
    }
    
    checkBothClosed();
  };
  
  const cleanupServer = async (reason: string) => {
    if (serverClosed) return;
    serverClosed = true;
    serverReason = reason;
    
    // Allow client to continue if still active
    if (!clientClosed && clientSocket.writable) {
      // Half-close: stop reading from server, let client finish
      serverSocket.pause();
      serverSocket.unpipe(clientSocket);
      await cleanupSocket(serverSocket, 'server', { allowDrain: true, gracePeriod: 5000 });
    } else {
      await cleanupSocket(serverSocket, 'server', { immediate: true });
    }
    
    checkBothClosed();
  };
  
  return { cleanupClient, cleanupServer };
}

/**
 * Setup socket error and close handlers with proper cleanup
 * @param socket The socket to setup handlers for
 * @param handleClose The cleanup function to call
 * @param handleTimeout Optional custom timeout handler
 * @param errorPrefix Optional prefix for error messages
 */
export function setupSocketHandlers(
  socket: plugins.net.Socket | plugins.tls.TLSSocket,
  handleClose: (reason: string) => void,
  handleTimeout?: (socket: plugins.net.Socket | plugins.tls.TLSSocket) => void,
  errorPrefix?: string
): void {
  socket.on('error', (error) => {
    const prefix = errorPrefix || 'Socket';
    handleClose(`${prefix}_error: ${error.message}`);
  });
  
  socket.on('close', () => {
    const prefix = errorPrefix || 'socket';
    handleClose(`${prefix}_closed`);
  });
  
  socket.on('timeout', () => {
    if (handleTimeout) {
      handleTimeout(socket);  // Custom timeout handling
    } else {
      // Default: just log, don't close
      console.warn(`Socket timeout: ${errorPrefix || 'socket'}`);
    }
  });
}

/**
 * Pipe two sockets together with proper cleanup on either end
 * @param socket1 First socket
 * @param socket2 Second socket
 */
export function pipeSockets(
  socket1: plugins.net.Socket | plugins.tls.TLSSocket,
  socket2: plugins.net.Socket | plugins.tls.TLSSocket
): void {
  socket1.pipe(socket2);
  socket2.pipe(socket1);
}

/**
 * Create a socket with immediate error handling to prevent crashes
 * @param options Socket creation options
 * @returns The created socket
 */
export function createSocketWithErrorHandler(options: SafeSocketOptions): plugins.net.Socket {
  const { port, host, onError, onConnect, timeout } = options;
  
  // Create socket with immediate error handler attachment
  const socket = new plugins.net.Socket();
  
  // Attach error handler BEFORE connecting to catch immediate errors
  socket.on('error', (error) => {
    console.error(`Socket connection error to ${host}:${port}: ${error.message}`);
    if (onError) {
      onError(error);
    }
  });
  
  // Attach connect handler if provided
  if (onConnect) {
    socket.on('connect', onConnect);
  }
  
  // Set timeout if provided
  if (timeout) {
    socket.setTimeout(timeout);
  }
  
  // Now attempt to connect - any immediate errors will be caught
  socket.connect(port, host);
  
  return socket;
}