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 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,
  options: { enableHalfOpen?: boolean } = {}
): { 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;
    
    // Default behavior: close both sockets when one closes (required for proxy chains)
    if (!serverClosed && !options.enableHalfOpen) {
      serverSocket.destroy();
    }
    
    // Half-open support (opt-in only)
    if (!serverClosed && serverSocket.writable && options.enableHalfOpen) {
      // 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;
    
    // Default behavior: close both sockets when one closes (required for proxy chains)
    if (!clientClosed && !options.enableHalfOpen) {
      clientSocket.destroy();
    }
    
    // Half-open support (opt-in only)
    if (!clientClosed && clientSocket.writable && options.enableHalfOpen) {
      // 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'}`);
    }
  });
}

/**
 * Setup bidirectional data forwarding between two sockets with proper cleanup
 * @param clientSocket The client/incoming socket
 * @param serverSocket The server/outgoing socket
 * @param handlers Object containing optional handlers for data and cleanup
 * @returns Cleanup functions for both sockets
 */
export function setupBidirectionalForwarding(
  clientSocket: plugins.net.Socket | plugins.tls.TLSSocket,
  serverSocket: plugins.net.Socket | plugins.tls.TLSSocket,
  handlers: {
    onClientData?: (chunk: Buffer) => void;
    onServerData?: (chunk: Buffer) => void;
    onCleanup: (reason: string) => void;
    enableHalfOpen?: boolean;
  }
): { cleanupClient: (reason: string) => Promise<void>, cleanupServer: (reason: string) => Promise<void> } {
  // Set up cleanup handlers
  const { cleanupClient, cleanupServer } = createIndependentSocketHandlers(
    clientSocket,
    serverSocket,
    handlers.onCleanup,
    { enableHalfOpen: handlers.enableHalfOpen }
  );
  
  // Set up error and close handlers
  setupSocketHandlers(clientSocket, cleanupClient, undefined, 'client');
  setupSocketHandlers(serverSocket, cleanupServer, undefined, 'server');
  
  // Set up data forwarding with backpressure handling
  clientSocket.on('data', (chunk: Buffer) => {
    if (handlers.onClientData) {
      handlers.onClientData(chunk);
    }
    
    if (serverSocket.writable) {
      const flushed = serverSocket.write(chunk);
      
      // Handle backpressure
      if (!flushed) {
        clientSocket.pause();
        serverSocket.once('drain', () => {
          if (!clientSocket.destroyed) {
            clientSocket.resume();
          }
        });
      }
    }
  });
  
  serverSocket.on('data', (chunk: Buffer) => {
    if (handlers.onServerData) {
      handlers.onServerData(chunk);
    }
    
    if (clientSocket.writable) {
      const flushed = clientSocket.write(chunk);
      
      // Handle backpressure
      if (!flushed) {
        serverSocket.pause();
        clientSocket.once('drain', () => {
          if (!serverSocket.destroyed) {
            serverSocket.resume();
          }
        });
      }
    }
  });
  
  return { cleanupClient, cleanupServer };
}

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