smartproxy/ts/proxies/smart-proxy/connection-manager.ts

import * as plugins from '../../plugins.js';
import type { IConnectionRecord, ISmartProxyOptions } from './models/interfaces.js';
import { SecurityManager } from './security-manager.js';
import { TimeoutManager } from './timeout-manager.js';
import { logger } from '../../core/utils/logger.js';

/**
 * Manages connection lifecycle, tracking, and cleanup
 */
export class ConnectionManager {
  private connectionRecords: Map<string, IConnectionRecord> = new Map();
  private terminationStats: {
    incoming: Record<string, number>;
    outgoing: Record<string, number>;
  } = { incoming: {}, outgoing: {} };

  constructor(
    private settings: ISmartProxyOptions,
    private securityManager: SecurityManager,
    private timeoutManager: TimeoutManager
  ) {}
  
  /**
   * Generate a unique connection ID
   */
  public generateConnectionId(): string {
    return Math.random().toString(36).substring(2, 15) + 
           Math.random().toString(36).substring(2, 15);
  }
  
  /**
   * Create and track a new connection
   */
  public createConnection(socket: plugins.net.Socket): IConnectionRecord {
    const connectionId = this.generateConnectionId();
    const remoteIP = socket.remoteAddress || '';
    const localPort = socket.localPort || 0;

    const record: IConnectionRecord = {
      id: connectionId,
      incoming: socket,
      outgoing: null,
      incomingStartTime: Date.now(),
      lastActivity: Date.now(),
      connectionClosed: false,
      pendingData: [],
      pendingDataSize: 0,
      bytesReceived: 0,
      bytesSent: 0,
      remoteIP,
      localPort,
      isTLS: false,
      tlsHandshakeComplete: false,
      hasReceivedInitialData: false,
      hasKeepAlive: false,
      incomingTerminationReason: null,
      outgoingTerminationReason: null,
      usingNetworkProxy: false,
      isBrowserConnection: false,
      domainSwitches: 0
    };
    
    this.trackConnection(connectionId, record);
    return record;
  }
  
  /**
   * Track an existing connection
   */
  public trackConnection(connectionId: string, record: IConnectionRecord): void {
    this.connectionRecords.set(connectionId, record);
    this.securityManager.trackConnectionByIP(record.remoteIP, connectionId);
  }

  /**
   * Get a connection by ID
   */
  public getConnection(connectionId: string): IConnectionRecord | undefined {
    return this.connectionRecords.get(connectionId);
  }

  /**
   * Get all active connections
   */
  public getConnections(): Map<string, IConnectionRecord> {
    return this.connectionRecords;
  }
  
  /**
   * Get count of active connections
   */
  public getConnectionCount(): number {
    return this.connectionRecords.size;
  }
  
  /**
   * Initiates cleanup once for a connection
   */
  public initiateCleanupOnce(record: IConnectionRecord, reason: string = 'normal'): void {
    if (this.settings.enableDetailedLogging) {
      logger.log('info', `Connection cleanup initiated`, { connectionId: record.id, remoteIP: record.remoteIP, reason, component: 'connection-manager' });
    }

    if (
      record.incomingTerminationReason === null ||
      record.incomingTerminationReason === undefined
    ) {
      record.incomingTerminationReason = reason;
      this.incrementTerminationStat('incoming', reason);
    }

    this.cleanupConnection(record, reason);
  }

  /**
   * Clean up a connection record
   */
  public cleanupConnection(record: IConnectionRecord, reason: string = 'normal'): void {
    if (!record.connectionClosed) {
      record.connectionClosed = true;

      // Track connection termination
      this.securityManager.removeConnectionByIP(record.remoteIP, record.id);

      if (record.cleanupTimer) {
        clearTimeout(record.cleanupTimer);
        record.cleanupTimer = undefined;
      }

      // Detailed logging data
      const duration = Date.now() - record.incomingStartTime;
      const bytesReceived = record.bytesReceived;
      const bytesSent = record.bytesSent;

      // Remove all data handlers to make sure we clean up properly
      if (record.incoming) {
        try {
          // Remove our safe data handler
          record.incoming.removeAllListeners('data');
          // Reset the handler references
          record.renegotiationHandler = undefined;
        } catch (err) {
          logger.log('error', `Error removing data handlers for connection ${record.id}: ${err}`, { connectionId: record.id, error: err, component: 'connection-manager' });
        }
      }

      // Handle incoming socket
      this.cleanupSocket(record, 'incoming', record.incoming);
      
      // Handle outgoing socket
      if (record.outgoing) {
        this.cleanupSocket(record, 'outgoing', record.outgoing);
      }

      // Clear pendingData to avoid memory leaks
      record.pendingData = [];
      record.pendingDataSize = 0;

      // Remove the record from the tracking map
      this.connectionRecords.delete(record.id);

      // Log connection details
      if (this.settings.enableDetailedLogging) {
        logger.log('info', 
          `Connection from ${record.remoteIP} on port ${record.localPort} terminated (${reason}). ` +
          `Duration: ${plugins.prettyMs(duration)}, Bytes IN: ${bytesReceived}, OUT: ${bytesSent}, ` +
          `TLS: ${record.isTLS ? 'Yes' : 'No'}, Keep-Alive: ${record.hasKeepAlive ? 'Yes' : 'No'}` +
          `${record.usingNetworkProxy ? ', Using NetworkProxy' : ''}` +
          `${record.domainSwitches ? `, Domain switches: ${record.domainSwitches}` : ''}`, 
          {
            connectionId: record.id,
            remoteIP: record.remoteIP,
            localPort: record.localPort,
            reason,
            duration: plugins.prettyMs(duration),
            bytes: { in: bytesReceived, out: bytesSent },
            tls: record.isTLS,
            keepAlive: record.hasKeepAlive,
            usingNetworkProxy: record.usingNetworkProxy,
            domainSwitches: record.domainSwitches || 0,
            component: 'connection-manager'
          }
        );
      } else {
        logger.log('info', 
          `Connection from ${record.remoteIP} terminated (${reason}). Active connections: ${this.connectionRecords.size}`, 
          {
            connectionId: record.id,
            remoteIP: record.remoteIP,
            reason,
            activeConnections: this.connectionRecords.size,
            component: 'connection-manager'
          }
        );
      }
    }
  }
  
  /**
   * Helper method to clean up a socket
   */
  private cleanupSocket(record: IConnectionRecord, side: 'incoming' | 'outgoing', socket: plugins.net.Socket): void {
    try {
      if (!socket.destroyed) {
        // Try graceful shutdown first, then force destroy after a short timeout
        socket.end();
        const socketTimeout = setTimeout(() => {
          try {
            if (!socket.destroyed) {
              socket.destroy();
            }
          } catch (err) {
            logger.log('error', `Error destroying ${side} socket for connection ${record.id}: ${err}`, { connectionId: record.id, side, error: err, component: 'connection-manager' });
          }
        }, 1000);

        // Ensure the timeout doesn't block Node from exiting
        if (socketTimeout.unref) {
          socketTimeout.unref();
        }
      }
    } catch (err) {
      logger.log('error', `Error closing ${side} socket for connection ${record.id}: ${err}`, { connectionId: record.id, side, error: err, component: 'connection-manager' });
      try {
        if (!socket.destroyed) {
          socket.destroy();
        }
      } catch (destroyErr) {
        logger.log('error', `Error destroying ${side} socket for connection ${record.id}: ${destroyErr}`, { connectionId: record.id, side, error: destroyErr, component: 'connection-manager' });
      }
    }
  }
  
  /**
   * Creates a generic error handler for incoming or outgoing sockets
   */
  public handleError(side: 'incoming' | 'outgoing', record: IConnectionRecord) {
    return (err: Error) => {
      const code = (err as any).code;
      let reason = 'error';

      const now = Date.now();
      const connectionDuration = now - record.incomingStartTime;
      const lastActivityAge = now - record.lastActivity;

      if (code === 'ECONNRESET') {
        reason = 'econnreset';
        logger.log('warn', `ECONNRESET on ${side} connection from ${record.remoteIP}. Error: ${err.message}. Duration: ${plugins.prettyMs(connectionDuration)}, Last activity: ${plugins.prettyMs(lastActivityAge)}`, {
          connectionId: record.id,
          side,
          remoteIP: record.remoteIP,
          error: err.message,
          duration: plugins.prettyMs(connectionDuration),
          lastActivity: plugins.prettyMs(lastActivityAge),
          component: 'connection-manager'
        });
      } else if (code === 'ETIMEDOUT') {
        reason = 'etimedout';
        logger.log('warn', `ETIMEDOUT on ${side} connection from ${record.remoteIP}. Error: ${err.message}. Duration: ${plugins.prettyMs(connectionDuration)}, Last activity: ${plugins.prettyMs(lastActivityAge)}`, {
          connectionId: record.id,
          side,
          remoteIP: record.remoteIP,
          error: err.message,
          duration: plugins.prettyMs(connectionDuration),
          lastActivity: plugins.prettyMs(lastActivityAge),
          component: 'connection-manager'
        });
      } else {
        logger.log('error', `Error on ${side} connection from ${record.remoteIP}: ${err.message}. Duration: ${plugins.prettyMs(connectionDuration)}, Last activity: ${plugins.prettyMs(lastActivityAge)}`, {
          connectionId: record.id,
          side,
          remoteIP: record.remoteIP,
          error: err.message,
          duration: plugins.prettyMs(connectionDuration),
          lastActivity: plugins.prettyMs(lastActivityAge),
          component: 'connection-manager'
        });
      }

      if (side === 'incoming' && record.incomingTerminationReason === null) {
        record.incomingTerminationReason = reason;
        this.incrementTerminationStat('incoming', reason);
      } else if (side === 'outgoing' && record.outgoingTerminationReason === null) {
        record.outgoingTerminationReason = reason;
        this.incrementTerminationStat('outgoing', reason);
      }

      this.initiateCleanupOnce(record, reason);
    };
  }
  
  /**
   * Creates a generic close handler for incoming or outgoing sockets
   */
  public handleClose(side: 'incoming' | 'outgoing', record: IConnectionRecord) {
    return () => {
      if (this.settings.enableDetailedLogging) {
        logger.log('info', `Connection closed on ${side} side`, {
          connectionId: record.id,
          side,
          remoteIP: record.remoteIP,
          component: 'connection-manager'
        });
      }

      if (side === 'incoming' && record.incomingTerminationReason === null) {
        record.incomingTerminationReason = 'normal';
        this.incrementTerminationStat('incoming', 'normal');
      } else if (side === 'outgoing' && record.outgoingTerminationReason === null) {
        record.outgoingTerminationReason = 'normal';
        this.incrementTerminationStat('outgoing', 'normal');
        // Record the time when outgoing socket closed.
        record.outgoingClosedTime = Date.now();
      }

      this.initiateCleanupOnce(record, 'closed_' + side);
    };
  }
  
  /**
   * Increment termination statistics
   */
  public incrementTerminationStat(side: 'incoming' | 'outgoing', reason: string): void {
    this.terminationStats[side][reason] = (this.terminationStats[side][reason] || 0) + 1;
  }
  
  /**
   * Get termination statistics
   */
  public getTerminationStats(): { incoming: Record<string, number>; outgoing: Record<string, number> } {
    return this.terminationStats;
  }
  
  /**
   * Check for stalled/inactive connections
   */
  public performInactivityCheck(): void {
    const now = Date.now();
    const connectionIds = [...this.connectionRecords.keys()];
    
    for (const id of connectionIds) {
      const record = this.connectionRecords.get(id);
      if (!record) continue;

      // Skip inactivity check if disabled or for immortal keep-alive connections
      if (
        this.settings.disableInactivityCheck ||
        (record.hasKeepAlive && this.settings.keepAliveTreatment === 'immortal')
      ) {
        continue;
      }
      
      const inactivityTime = now - record.lastActivity;

      // Use extended timeout for extended-treatment keep-alive connections
      let effectiveTimeout = this.settings.inactivityTimeout!;
      if (record.hasKeepAlive && this.settings.keepAliveTreatment === 'extended') {
        const multiplier = this.settings.keepAliveInactivityMultiplier || 6;
        effectiveTimeout = effectiveTimeout * multiplier;
      }

      if (inactivityTime > effectiveTimeout && !record.connectionClosed) {
        // For keep-alive connections, issue a warning first
        if (record.hasKeepAlive && !record.inactivityWarningIssued) {
          logger.log('warn', `Keep-alive connection ${id} from ${record.remoteIP} inactive for ${plugins.prettyMs(inactivityTime)}. Will close in 10 minutes if no activity.`, {
            connectionId: id,
            remoteIP: record.remoteIP,
            inactiveFor: plugins.prettyMs(inactivityTime),
            closureWarning: '10 minutes',
            component: 'connection-manager'
          });

          // Set warning flag and add grace period
          record.inactivityWarningIssued = true;
          record.lastActivity = now - (effectiveTimeout - 600000);

          // Try to stimulate activity with a probe packet
          if (record.outgoing && !record.outgoing.destroyed) {
            try {
              record.outgoing.write(Buffer.alloc(0));

              if (this.settings.enableDetailedLogging) {
                logger.log('info', `Sent probe packet to test keep-alive connection ${id}`, { connectionId: id, component: 'connection-manager' });
              }
            } catch (err) {
              logger.log('error', `Error sending probe packet to connection ${id}: ${err}`, { connectionId: id, error: err, component: 'connection-manager' });
            }
          }
        } else {
          // For non-keep-alive or after warning, close the connection
          logger.log('warn', `Closing inactive connection ${id} from ${record.remoteIP} (inactive for ${plugins.prettyMs(inactivityTime)}, keep-alive: ${record.hasKeepAlive ? 'Yes' : 'No'})`, {
            connectionId: id,
            remoteIP: record.remoteIP,
            inactiveFor: plugins.prettyMs(inactivityTime),
            hasKeepAlive: record.hasKeepAlive,
            component: 'connection-manager'
          });
          this.cleanupConnection(record, 'inactivity');
        }
      } else if (inactivityTime <= effectiveTimeout && record.inactivityWarningIssued) {
        // If activity detected after warning, clear the warning
        if (this.settings.enableDetailedLogging) {
          logger.log('info', `Connection ${id} activity detected after inactivity warning`, {
            connectionId: id,
            component: 'connection-manager'
          });
        }
        record.inactivityWarningIssued = false;
      }
      
      // Parity check: if outgoing socket closed and incoming remains active
      if (
        record.outgoingClosedTime &&
        !record.incoming.destroyed &&
        !record.connectionClosed &&
        now - record.outgoingClosedTime > 120000
      ) {
        logger.log('warn', `Parity check: Connection ${id} from ${record.remoteIP} has incoming socket still active ${plugins.prettyMs(now - record.outgoingClosedTime)} after outgoing socket closed`, {
          connectionId: id,
          remoteIP: record.remoteIP,
          timeElapsed: plugins.prettyMs(now - record.outgoingClosedTime),
          component: 'connection-manager'
        });
        this.cleanupConnection(record, 'parity_check');
      }
    }
  }
  
  /**
   * Clear all connections (for shutdown)
   */
  public clearConnections(): void {
    // Create a copy of the keys to avoid modification during iteration
    const connectionIds = [...this.connectionRecords.keys()];
    
    // First pass: End all connections gracefully
    for (const id of connectionIds) {
      const record = this.connectionRecords.get(id);
      if (record) {
        try {
          // Clear any timers
          if (record.cleanupTimer) {
            clearTimeout(record.cleanupTimer);
            record.cleanupTimer = undefined;
          }

          // End sockets gracefully
          if (record.incoming && !record.incoming.destroyed) {
            record.incoming.end();
          }

          if (record.outgoing && !record.outgoing.destroyed) {
            record.outgoing.end();
          }
        } catch (err) {
          logger.log('error', `Error during graceful end of connection ${id}: ${err}`, { connectionId: id, error: err, component: 'connection-manager' });
        }
      }
    }

    // Short delay to allow graceful ends to process
    setTimeout(() => {
      // Second pass: Force destroy everything
      for (const id of connectionIds) {
        const record = this.connectionRecords.get(id);
        if (record) {
          try {
            // Remove all listeners to prevent memory leaks
            if (record.incoming) {
              record.incoming.removeAllListeners();
              if (!record.incoming.destroyed) {
                record.incoming.destroy();
              }
            }

            if (record.outgoing) {
              record.outgoing.removeAllListeners();
              if (!record.outgoing.destroyed) {
                record.outgoing.destroy();
              }
            }
          } catch (err) {
            logger.log('error', `Error during forced destruction of connection ${id}: ${err}`, { connectionId: id, error: err, component: 'connection-manager' });
          }
        }
      }
      
      // Clear all maps
      this.connectionRecords.clear();
      this.terminationStats = { incoming: {}, outgoing: {} };
    }, 100);
  }
}