update

ts/mail/delivery/smtpserver/connection-manager.ts

@@ -13,6 +13,7 @@ import { getSocketDetails, formatMultilineResponse } from './utils/helpers.js';
 /**
  * Manager for SMTP connections
  * Handles connection setup, event listeners, and lifecycle management
+ * Provides resource management, connection tracking, and monitoring
  */
 export class ConnectionManager implements IConnectionManager {
   /**
@@ -20,18 +21,50 @@ export class ConnectionManager implements IConnectionManager {
    */
   private activeConnections: Set<plugins.net.Socket | plugins.tls.TLSSocket> = new Set();
   
+  /**
+   * Connection tracking for resource management
+   */
+  private connectionStats = {
+    totalConnections: 0,
+    activeConnections: 0,
+    peakConnections: 0,
+    rejectedConnections: 0,
+    closedConnections: 0,
+    erroredConnections: 0,
+    timedOutConnections: 0
+  };
+
+  /**
+   * Per-IP connection tracking for rate limiting
+   */
+  private ipConnections: Map<string, {
+    count: number;
+    firstConnection: number;
+    lastConnection: number;
+  }> = new Map();
+  
+  /**
+   * Resource monitoring interval
+   */
+  private resourceCheckInterval: NodeJS.Timeout | null = null;
+  
   /**
    * Reference to the session manager
    */
   private sessionManager: ISessionManager;
   
   /**
-   * SMTP server options
+   * SMTP server options with enhanced resource controls
    */
   private options: {
     hostname: string;
     maxConnections: number;
     socketTimeout: number;
+    maxConnectionsPerIP: number;
+    connectionRateLimit: number;
+    connectionRateWindow: number;
+    bufferSizeLimit: number;
+    resourceCheckInterval: number;
   };
   
   /**
@@ -40,7 +73,7 @@ export class ConnectionManager implements IConnectionManager {
   private commandHandler: (socket: plugins.net.Socket | plugins.tls.TLSSocket, line: string) => void;
   
   /**
-   * Creates a new connection manager
+   * Creates a new connection manager with enhanced resource management
    * @param sessionManager - Session manager instance
    * @param commandHandler - Command handler function
    * @param options - Connection manager options
@@ -52,26 +85,250 @@ export class ConnectionManager implements IConnectionManager {
       hostname?: string;
       maxConnections?: number;
       socketTimeout?: number;
+      maxConnectionsPerIP?: number;
+      connectionRateLimit?: number;
+      connectionRateWindow?: number;
+      bufferSizeLimit?: number;
+      resourceCheckInterval?: number;
     } = {}
   ) {
     this.sessionManager = sessionManager;
     this.commandHandler = commandHandler;
     
+    // Default values for resource management
+    const DEFAULT_MAX_CONNECTIONS_PER_IP = 10;
+    const DEFAULT_CONNECTION_RATE_LIMIT = 30; // connections per window
+    const DEFAULT_CONNECTION_RATE_WINDOW = 60 * 1000; // 60 seconds window
+    const DEFAULT_BUFFER_SIZE_LIMIT = 10 * 1024 * 1024; // 10 MB
+    const DEFAULT_RESOURCE_CHECK_INTERVAL = 30 * 1000; // 30 seconds
+    
     this.options = {
       hostname: options.hostname || SMTP_DEFAULTS.HOSTNAME,
       maxConnections: options.maxConnections || SMTP_DEFAULTS.MAX_CONNECTIONS,
-      socketTimeout: options.socketTimeout || SMTP_DEFAULTS.SOCKET_TIMEOUT
+      socketTimeout: options.socketTimeout || SMTP_DEFAULTS.SOCKET_TIMEOUT,
+      maxConnectionsPerIP: options.maxConnectionsPerIP || DEFAULT_MAX_CONNECTIONS_PER_IP,
+      connectionRateLimit: options.connectionRateLimit || DEFAULT_CONNECTION_RATE_LIMIT,
+      connectionRateWindow: options.connectionRateWindow || DEFAULT_CONNECTION_RATE_WINDOW,
+      bufferSizeLimit: options.bufferSizeLimit || DEFAULT_BUFFER_SIZE_LIMIT,
+      resourceCheckInterval: options.resourceCheckInterval || DEFAULT_RESOURCE_CHECK_INTERVAL
     };
+    
+    // Start resource monitoring
+    this.startResourceMonitoring();
   }
   
   /**
-   * Handle a new connection
+   * Start resource monitoring interval to check resource usage
+   */
+  private startResourceMonitoring(): void {
+    // Clear any existing interval
+    if (this.resourceCheckInterval) {
+      clearInterval(this.resourceCheckInterval);
+    }
+    
+    // Set up new interval
+    this.resourceCheckInterval = setInterval(() => {
+      this.monitorResourceUsage();
+    }, this.options.resourceCheckInterval);
+  }
+  
+  /**
+   * Monitor resource usage and log statistics
+   */
+  private monitorResourceUsage(): void {
+    // Calculate memory usage
+    const memoryUsage = process.memoryUsage();
+    const memoryUsageMB = {
+      rss: Math.round(memoryUsage.rss / 1024 / 1024),
+      heapTotal: Math.round(memoryUsage.heapTotal / 1024 / 1024),
+      heapUsed: Math.round(memoryUsage.heapUsed / 1024 / 1024),
+      external: Math.round(memoryUsage.external / 1024 / 1024)
+    };
+    
+    // Calculate connection rate metrics
+    const activeIPs = Array.from(this.ipConnections.entries())
+      .filter(([_, data]) => data.count > 0).length;
+      
+    const highVolumeIPs = Array.from(this.ipConnections.entries())
+      .filter(([_, data]) => data.count > this.options.connectionRateLimit / 2).length;
+    
+    // Log resource usage with more detailed metrics
+    SmtpLogger.info('Resource usage stats', {
+      connections: {
+        active: this.activeConnections.size,
+        total: this.connectionStats.totalConnections,
+        peak: this.connectionStats.peakConnections,
+        rejected: this.connectionStats.rejectedConnections,
+        closed: this.connectionStats.closedConnections,
+        errored: this.connectionStats.erroredConnections,
+        timedOut: this.connectionStats.timedOutConnections
+      },
+      memory: memoryUsageMB,
+      ipTracking: {
+        uniqueIPs: this.ipConnections.size,
+        activeIPs: activeIPs,
+        highVolumeIPs: highVolumeIPs
+      },
+      resourceLimits: {
+        maxConnections: this.options.maxConnections,
+        maxConnectionsPerIP: this.options.maxConnectionsPerIP,
+        connectionRateLimit: this.options.connectionRateLimit,
+        bufferSizeLimit: Math.round(this.options.bufferSizeLimit / 1024 / 1024) + 'MB'
+      }
+    });
+    
+    // Check for potential DoS conditions
+    if (highVolumeIPs > 3) {
+      SmtpLogger.warn(`Potential DoS detected: ${highVolumeIPs} IPs with high connection rates`);
+    }
+    
+    // Assess memory usage trends
+    if (memoryUsageMB.heapUsed > 500) { // Over 500MB heap used
+      SmtpLogger.warn(`High memory usage detected: ${memoryUsageMB.heapUsed}MB heap used`);
+    }
+    
+    // Clean up expired IP rate limits and validate resource tracking
+    this.cleanupIpRateLimits();
+  }
+  
+  /**
+   * Clean up expired IP rate limits and perform additional resource monitoring
+   */
+  private cleanupIpRateLimits(): void {
+    const now = Date.now();
+    const windowThreshold = now - this.options.connectionRateWindow;
+    let activeIps = 0;
+    let removedEntries = 0;
+    
+    // Iterate through IP connections and manage entries
+    for (const [ip, data] of this.ipConnections.entries()) {
+      // If the last connection was before the window threshold + one extra window, remove the entry
+      if (data.lastConnection < windowThreshold - this.options.connectionRateWindow) {
+        // Remove stale entries to prevent memory growth
+        this.ipConnections.delete(ip);
+        removedEntries++;
+      } 
+      // If last connection was before the window threshold, reset the count
+      else if (data.lastConnection < windowThreshold) {
+        if (data.count > 0) {
+          // Reset but keep the IP in the map with a zero count
+          this.ipConnections.set(ip, { 
+            count: 0, 
+            firstConnection: now, 
+            lastConnection: now 
+          });
+        }
+      } else {
+        // This IP is still active within the current window
+        activeIps++;
+      }
+    }
+    
+    // Log cleanup activity if significant changes occurred
+    if (removedEntries > 0) {
+      SmtpLogger.debug(`IP rate limit cleanup: removed ${removedEntries} stale entries, ${this.ipConnections.size} remaining, ${activeIps} active in current window`);
+    }
+    
+    // Check for memory leaks in connection tracking
+    if (this.activeConnections.size > 0 && this.connectionStats.activeConnections !== this.activeConnections.size) {
+      SmtpLogger.warn(`Connection tracking inconsistency detected: stats.active=${this.connectionStats.activeConnections}, actual=${this.activeConnections.size}`);
+      // Fix the inconsistency
+      this.connectionStats.activeConnections = this.activeConnections.size;
+    }
+    
+    // Validate and clean leaked resources if needed
+    this.validateResourceTracking();
+  }
+  
+  /**
+   * Validate and repair resource tracking to prevent leaks
+   */
+  private validateResourceTracking(): void {
+    // Prepare a detailed report if inconsistencies are found
+    const inconsistenciesFound = [];
+    
+    // 1. Check active connections count matches activeConnections set size
+    if (this.connectionStats.activeConnections !== this.activeConnections.size) {
+      inconsistenciesFound.push({
+        issue: 'Active connection count mismatch',
+        stats: this.connectionStats.activeConnections,
+        actual: this.activeConnections.size,
+        action: 'Auto-corrected'
+      });
+      this.connectionStats.activeConnections = this.activeConnections.size;
+    }
+    
+    // 2. Check for destroyed sockets in active connections
+    let destroyedSocketsCount = 0;
+    for (const socket of this.activeConnections) {
+      if (socket.destroyed) {
+        destroyedSocketsCount++;
+        // This should not happen - remove destroyed sockets from tracking
+        this.activeConnections.delete(socket);
+      }
+    }
+    
+    if (destroyedSocketsCount > 0) {
+      inconsistenciesFound.push({
+        issue: 'Destroyed sockets in active list',
+        count: destroyedSocketsCount,
+        action: 'Removed from tracking'
+      });
+      // Update active connections count after cleanup
+      this.connectionStats.activeConnections = this.activeConnections.size;
+    }
+    
+    // 3. Check for sessions without corresponding active connections
+    const sessionCount = this.sessionManager.getSessionCount();
+    if (sessionCount > this.activeConnections.size) {
+      inconsistenciesFound.push({
+        issue: 'Orphaned sessions',
+        sessions: sessionCount,
+        connections: this.activeConnections.size,
+        action: 'Session cleanup recommended'
+      });
+    }
+    
+    // If any inconsistencies found, log a detailed report
+    if (inconsistenciesFound.length > 0) {
+      SmtpLogger.warn('Resource tracking inconsistencies detected and repaired', { inconsistencies: inconsistenciesFound });
+    }
+  }
+  
+  /**
+   * Handle a new connection with resource management
    * @param socket - Client socket
    */
   public handleNewConnection(socket: plugins.net.Socket): void {
-    // Check if maximum connections reached
+    // Update connection stats
+    this.connectionStats.totalConnections++;
+    this.connectionStats.activeConnections = this.activeConnections.size + 1;
+    
+    if (this.connectionStats.activeConnections > this.connectionStats.peakConnections) {
+      this.connectionStats.peakConnections = this.connectionStats.activeConnections;
+    }
+    
+    // Get client IP
+    const remoteAddress = socket.remoteAddress || '0.0.0.0';
+    
+    // Check rate limits by IP
+    if (this.isIPRateLimited(remoteAddress)) {
+      this.rejectConnection(socket, 'Rate limit exceeded');
+      this.connectionStats.rejectedConnections++;
+      return;
+    }
+    
+    // Check per-IP connection limit
+    if (this.hasReachedIPConnectionLimit(remoteAddress)) {
+      this.rejectConnection(socket, 'Too many connections from this IP');
+      this.connectionStats.rejectedConnections++;
+      return;
+    }
+    
+    // Check if maximum global connections reached
     if (this.hasReachedMaxConnections()) {
       this.rejectConnection(socket, 'Too many connections');
+      this.connectionStats.rejectedConnections++;
       return;
     }
     
@@ -82,6 +339,27 @@ export class ConnectionManager implements IConnectionManager {
     socket.setKeepAlive(true);
     socket.setTimeout(this.options.socketTimeout);
     
+    // Explicitly set socket buffer sizes to prevent memory issues
+    socket.setNoDelay(true); // Disable Nagle's algorithm for better responsiveness
+    
+    // Set limits on socket buffer size if supported by Node.js version
+    try {
+      // Here we set reasonable buffer limits to prevent memory exhaustion attacks
+      const highWaterMark = 64 * 1024; // 64 KB
+      if (typeof socket.// setReadableHighWaterMark === 'function') {
+        socket.// setReadableHighWaterMark(highWaterMark);
+      }
+      if (typeof socket.// setWritableHighWaterMark === 'function') {
+        socket.// setWritableHighWaterMark(highWaterMark);
+      }
+    } catch (error) {
+      // Ignore errors from older Node.js versions that don't support these methods
+      SmtpLogger.debug(`Could not set socket buffer limits: ${error instanceof Error ? error.message : String(error)}`);
+    }
+    
+    // Track this IP connection
+    this.trackIPConnection(remoteAddress);
+    
     // Set up event handlers
     this.setupSocketEventHandlers(socket);
     
@@ -97,13 +375,117 @@ export class ConnectionManager implements IConnectionManager {
   }
   
   /**
-   * Handle a new secure TLS connection
+   * Check if an IP has exceeded the rate limit
+   * @param ip - Client IP address
+   * @returns True if rate limited
+   */
+  private isIPRateLimited(ip: string): boolean {
+    const now = Date.now();
+    const ipData = this.ipConnections.get(ip);
+    
+    if (!ipData) {
+      return false; // No previous connections
+    }
+    
+    // Check if we're within the rate window
+    const isWithinWindow = now - ipData.firstConnection <= this.options.connectionRateWindow;
+    
+    // If within window and count exceeds limit, rate limit is applied
+    if (isWithinWindow && ipData.count >= this.options.connectionRateLimit) {
+      SmtpLogger.warn(`Rate limit exceeded for IP ${ip}: ${ipData.count} connections in ${Math.round((now - ipData.firstConnection) / 1000)}s`);
+      return true;
+    }
+    
+    return false;
+  }
+  
+  /**
+   * Track a new connection from an IP
+   * @param ip - Client IP address
+   */
+  private trackIPConnection(ip: string): void {
+    const now = Date.now();
+    const ipData = this.ipConnections.get(ip);
+    
+    if (!ipData) {
+      // First connection from this IP
+      this.ipConnections.set(ip, {
+        count: 1,
+        firstConnection: now,
+        lastConnection: now
+      });
+    } else {
+      // Check if we need to reset the window
+      if (now - ipData.lastConnection > this.options.connectionRateWindow) {
+        // Reset the window
+        this.ipConnections.set(ip, {
+          count: 1,
+          firstConnection: now,
+          lastConnection: now
+        });
+      } else {
+        // Increment within the current window
+        this.ipConnections.set(ip, {
+          count: ipData.count + 1,
+          firstConnection: ipData.firstConnection,
+          lastConnection: now
+        });
+      }
+    }
+  }
+  
+  /**
+   * Check if an IP has reached its connection limit
+   * @param ip - Client IP address
+   * @returns True if limit reached
+   */
+  private hasReachedIPConnectionLimit(ip: string): boolean {
+    let ipConnectionCount = 0;
+    
+    // Count active connections from this IP
+    for (const socket of this.activeConnections) {
+      if (socket.remoteAddress === ip) {
+        ipConnectionCount++;
+      }
+    }
+    
+    return ipConnectionCount >= this.options.maxConnectionsPerIP;
+  }
+  
+  /**
+   * Handle a new secure TLS connection with resource management
    * @param socket - Client TLS socket
    */
   public handleNewSecureConnection(socket: plugins.tls.TLSSocket): void {
-    // Check if maximum connections reached
+    // Update connection stats
+    this.connectionStats.totalConnections++;
+    this.connectionStats.activeConnections = this.activeConnections.size + 1;
+    
+    if (this.connectionStats.activeConnections > this.connectionStats.peakConnections) {
+      this.connectionStats.peakConnections = this.connectionStats.activeConnections;
+    }
+    
+    // Get client IP
+    const remoteAddress = socket.remoteAddress || '0.0.0.0';
+    
+    // Check rate limits by IP
+    if (this.isIPRateLimited(remoteAddress)) {
+      this.rejectConnection(socket, 'Rate limit exceeded');
+      this.connectionStats.rejectedConnections++;
+      return;
+    }
+    
+    // Check per-IP connection limit
+    if (this.hasReachedIPConnectionLimit(remoteAddress)) {
+      this.rejectConnection(socket, 'Too many connections from this IP');
+      this.connectionStats.rejectedConnections++;
+      return;
+    }
+    
+    // Check if maximum global connections reached
     if (this.hasReachedMaxConnections()) {
       this.rejectConnection(socket, 'Too many connections');
+      this.connectionStats.rejectedConnections++;
       return;
     }
     
@@ -114,6 +496,27 @@ export class ConnectionManager implements IConnectionManager {
     socket.setKeepAlive(true);
     socket.setTimeout(this.options.socketTimeout);
     
+    // Explicitly set socket buffer sizes to prevent memory issues
+    socket.setNoDelay(true); // Disable Nagle's algorithm for better responsiveness
+    
+    // Set limits on socket buffer size if supported by Node.js version
+    try {
+      // Here we set reasonable buffer limits to prevent memory exhaustion attacks
+      const highWaterMark = 64 * 1024; // 64 KB
+      if (typeof socket.// setReadableHighWaterMark === 'function') {
+        socket.// setReadableHighWaterMark(highWaterMark);
+      }
+      if (typeof socket.// setWritableHighWaterMark === 'function') {
+        socket.// setWritableHighWaterMark(highWaterMark);
+      }
+    } catch (error) {
+      // Ignore errors from older Node.js versions that don't support these methods
+      SmtpLogger.debug(`Could not set socket buffer limits: ${error instanceof Error ? error.message : String(error)}`);
+    }
+    
+    // Track this IP connection
+    this.trackIPConnection(remoteAddress);
+    
     // Set up event handlers
     this.setupSocketEventHandlers(socket);
     
@@ -128,7 +531,7 @@ export class ConnectionManager implements IConnectionManager {
   }
   
   /**
-   * Set up event handlers for a socket
+   * Set up event handlers for a socket with enhanced resource management
    * @param socket - Client socket
    */
   public setupSocketEventHandlers(socket: plugins.net.Socket | plugins.tls.TLSSocket): void {
@@ -144,30 +547,96 @@ export class ConnectionManager implements IConnectionManager {
     if (existingErrorHandler) socket.removeListener('error', existingErrorHandler);
     if (existingTimeoutHandler) socket.removeListener('timeout', existingTimeoutHandler);
     
-    // Data event - process incoming data from the client
+    // Data event - process incoming data from the client with resource limits
     let buffer = '';
+    let totalBytesReceived = 0;
+    
     socket.on('data', (data) => {
-      // Get current session and update activity timestamp
-      const session = this.sessionManager.getSession(socket);
-      if (session) {
-        this.sessionManager.updateSessionActivity(session);
-      }
-      
-      // Buffer incoming data
-      buffer += data.toString();
-      
-      // Process complete lines
-      let lineEndPos;
-      while ((lineEndPos = buffer.indexOf(SMTP_DEFAULTS.CRLF)) !== -1) {
-        // Extract a complete line
-        const line = buffer.substring(0, lineEndPos);
-        buffer = buffer.substring(lineEndPos + 2); // +2 to skip CRLF
-        
-        // Process non-empty lines
-        if (line.length > 0) {
-          // In DATA state, the command handler will process the data differently
-          this.commandHandler(socket, line);
+      try {
+        // Get current session and update activity timestamp
+        const session = this.sessionManager.getSession(socket);
+        if (session) {
+          this.sessionManager.updateSessionActivity(session);
         }
+        
+        // Check buffer size limits to prevent memory attacks
+        totalBytesReceived += data.length;
+        
+        if (buffer.length > this.options.bufferSizeLimit) {
+          // Buffer is too large, reject the connection
+          SmtpLogger.warn(`Buffer size limit exceeded: ${buffer.length} bytes for ${socket.remoteAddress}`);
+          this.sendResponse(socket, `${SmtpResponseCode.EXCEEDED_STORAGE} Message too large, disconnecting`);
+          socket.destroy();
+          return;
+        }
+        
+        // Impose a total transfer limit to prevent DoS
+        if (totalBytesReceived > this.options.bufferSizeLimit * 2) {
+          SmtpLogger.warn(`Total transfer limit exceeded: ${totalBytesReceived} bytes for ${socket.remoteAddress}`);
+          this.sendResponse(socket, `${SmtpResponseCode.EXCEEDED_STORAGE} Transfer limit exceeded, disconnecting`);
+          socket.destroy();
+          return;
+        }
+        
+        // Convert buffer to string safely with explicit encoding
+        const dataString = data.toString('utf8');
+        
+        // Buffer incoming data
+        buffer += dataString;
+        
+        // Process complete lines
+        let lineEndPos;
+        while ((lineEndPos = buffer.indexOf(SMTP_DEFAULTS.CRLF)) !== -1) {
+          // Extract a complete line
+          const line = buffer.substring(0, lineEndPos);
+          buffer = buffer.substring(lineEndPos + 2); // +2 to skip CRLF
+          
+          // Check line length to prevent extremely long lines
+          if (line.length > 4096) { // 4KB line limit is reasonable for SMTP
+            SmtpLogger.warn(`Line length limit exceeded: ${line.length} bytes for ${socket.remoteAddress}`);
+            this.sendResponse(socket, `${SmtpResponseCode.SYNTAX_ERROR} Line too long, disconnecting`);
+            socket.destroy();
+            return;
+          }
+          
+          // Process non-empty lines
+          if (line.length > 0) {
+            try {
+              // In DATA state, the command handler will process the data differently
+              this.commandHandler(socket, line);
+            } catch (cmdError) {
+              // Handle any errors in command processing
+              SmtpLogger.error(`Command handler error: ${cmdError instanceof Error ? cmdError.message : String(cmdError)}`);
+              
+              // If there's a severe error, close the connection
+              if (cmdError instanceof Error && 
+                  (cmdError.message.includes('fatal') || cmdError.message.includes('critical'))) {
+                socket.destroy();
+                return;
+              }
+            }
+          }
+        }
+        
+        // If buffer is getting too large without CRLF, it might be a DoS attempt
+        if (buffer.length > 10240) { // 10KB is a reasonable limit for a line without CRLF
+          SmtpLogger.warn(`Incomplete line too large: ${buffer.length} bytes for ${socket.remoteAddress}`);
+          this.sendResponse(socket, `${SmtpResponseCode.SYNTAX_ERROR} Incomplete line too large, disconnecting`);
+          socket.destroy();
+        }
+      } catch (error) {
+        // Handle any unexpected errors during data processing
+        SmtpLogger.error(`Data handler error: ${error instanceof Error ? error.message : String(error)}`);
+        socket.destroy();
+      }
+    });
+    
+    // Add drain event handler to manage flow control
+    socket.on('drain', () => {
+      // Socket buffer has been emptied, resume data flow if needed
+      if (socket.isPaused()) {
+        socket.resume();
+        SmtpLogger.debug(`Resumed socket for ${socket.remoteAddress} after drain`);
       }
     });
     
@@ -236,17 +705,45 @@ export class ConnectionManager implements IConnectionManager {
    * @param hadError - Whether the socket was closed due to error
    */
   private handleSocketClose(socket: plugins.net.Socket | plugins.tls.TLSSocket, hadError: boolean): void {
-    // Remove from active connections
-    this.activeConnections.delete(socket);
-    
-    // Get the session before removing it
-    const session = this.sessionManager.getSession(socket);
-    
-    // Remove from session manager
-    this.sessionManager.removeSession(socket);
-    
-    // Log connection close
-    SmtpLogger.logConnection(socket, 'close', session);
+    try {
+      // Update connection statistics
+      this.connectionStats.closedConnections++;
+      this.connectionStats.activeConnections = this.activeConnections.size - 1;
+      
+      // Get socket details for logging
+      const socketDetails = getSocketDetails(socket);
+      const socketId = `${socketDetails.remoteAddress}:${socketDetails.remotePort}`;
+      
+      // Log with appropriate level based on whether there was an error
+      if (hadError) {
+        SmtpLogger.warn(`Socket closed with error: ${socketId}`);
+      } else {
+        SmtpLogger.debug(`Socket closed normally: ${socketId}`);
+      }
+      
+      // Get the session before removing it
+      const session = this.sessionManager.getSession(socket);
+      
+      // Remove from active connections
+      this.activeConnections.delete(socket);
+      
+      // Remove from session manager
+      this.sessionManager.removeSession(socket);
+      
+      // Cancel any timeout ID stored in the session
+      if (session?.dataTimeoutId) {
+        clearTimeout(session.dataTimeoutId);
+      }
+      
+      // Log connection close with session details if available
+      SmtpLogger.logConnection(socket, 'close', session);
+    } catch (error) {
+      // Handle any unexpected errors during cleanup
+      SmtpLogger.error(`Error in handleSocketClose: ${error instanceof Error ? error.message : String(error)}`);
+      
+      // Ensure socket is removed from active connections even if an error occurs
+      this.activeConnections.delete(socket);
+    }
   }
   
   /**
@@ -255,15 +752,54 @@ export class ConnectionManager implements IConnectionManager {
    * @param error - Error object
    */
   private handleSocketError(socket: plugins.net.Socket | plugins.tls.TLSSocket, error: Error): void {
-    // Get the session
-    const session = this.sessionManager.getSession(socket);
-    
-    // Log the error
-    SmtpLogger.logConnection(socket, 'error', session, error);
-    
-    // Close the socket if not already closed
-    if (!socket.destroyed) {
-      socket.destroy();
+    try {
+      // Update connection statistics
+      this.connectionStats.erroredConnections++;
+      
+      // Get socket details for context
+      const socketDetails = getSocketDetails(socket);
+      const socketId = `${socketDetails.remoteAddress}:${socketDetails.remotePort}`;
+      
+      // Get the session
+      const session = this.sessionManager.getSession(socket);
+      
+      // Detailed error logging with context information
+      SmtpLogger.error(`Socket error for ${socketId}: ${error.message}`, {
+        errorCode: (error as any).code,
+        errorStack: error.stack,
+        sessionId: session?.id,
+        sessionState: session?.state,
+        remoteAddress: socketDetails.remoteAddress,
+        remotePort: socketDetails.remotePort
+      });
+      
+      // Log the error for connection tracking
+      SmtpLogger.logConnection(socket, 'error', session, error);
+      
+      // Cancel any timeout ID stored in the session
+      if (session?.dataTimeoutId) {
+        clearTimeout(session.dataTimeoutId);
+      }
+      
+      // Close the socket if not already closed
+      if (!socket.destroyed) {
+        socket.destroy();
+      }
+      
+      // Remove from active connections (cleanup after error)
+      this.activeConnections.delete(socket);
+      
+      // Remove from session manager
+      this.sessionManager.removeSession(socket);
+    } catch (handlerError) {
+      // Meta-error handling (errors in the error handler)
+      SmtpLogger.error(`Error in handleSocketError: ${handlerError instanceof Error ? handlerError.message : String(handlerError)}`);
+      
+      // Ensure socket is destroyed and removed from active connections
+      if (!socket.destroyed) {
+        socket.destroy();
+      }
+      this.activeConnections.delete(socket);
     }
   }
   
@@ -272,31 +808,77 @@ export class ConnectionManager implements IConnectionManager {
    * @param socket - Client socket
    */
   private handleSocketTimeout(socket: plugins.net.Socket | plugins.tls.TLSSocket): void {
-    // Get the session
-    const session = this.sessionManager.getSession(socket);
-    
-    if (session) {
-      // Log the timeout
-      SmtpLogger.warn(`Socket timeout from ${session.remoteAddress}`, {
-        sessionId: session.id,
-        remoteAddress: session.remoteAddress,
-        state: session.state,
-        timeout: this.options.socketTimeout
-      });
-      
-      // Send timeout notification
-      this.sendResponse(socket, `${SmtpResponseCode.SERVICE_NOT_AVAILABLE} Connection timeout - closing connection`);
-    } else {
-      // Log timeout without session context
-      const socketDetails = getSocketDetails(socket);
-      SmtpLogger.warn(`Socket timeout without session from ${socketDetails.remoteAddress}:${socketDetails.remotePort}`);
-    }
-    
-    // Close the socket
     try {
-      socket.end();
-    } catch (error) {
-      SmtpLogger.error(`Error ending timed out socket: ${error instanceof Error ? error.message : String(error)}`);
+      // Update connection statistics
+      this.connectionStats.timedOutConnections++;
+      
+      // Get socket details for context
+      const socketDetails = getSocketDetails(socket);
+      const socketId = `${socketDetails.remoteAddress}:${socketDetails.remotePort}`;
+      
+      // Get the session
+      const session = this.sessionManager.getSession(socket);
+      
+      // Get timing information for better debugging
+      const now = Date.now();
+      const idleTime = session?.lastActivity ? now - session.lastActivity : 'unknown';
+      
+      if (session) {
+        // Log the timeout with extended details
+        SmtpLogger.warn(`Socket timeout from ${session.remoteAddress}`, {
+          sessionId: session.id,
+          remoteAddress: session.remoteAddress,
+          state: session.state,
+          timeout: this.options.socketTimeout,
+          idleTime: idleTime,
+          emailState: session.envelope?.mailFrom ? 'has-sender' : 'no-sender',
+          recipientCount: session.envelope?.rcptTo?.length || 0
+        });
+        
+        // Cancel any timeout ID stored in the session
+        if (session.dataTimeoutId) {
+          clearTimeout(session.dataTimeoutId);
+        }
+        
+        // Send timeout notification to client
+        this.sendResponse(socket, `${SmtpResponseCode.SERVICE_NOT_AVAILABLE} Connection timeout - closing connection`);
+      } else {
+        // Log timeout without session context
+        SmtpLogger.warn(`Socket timeout without session from ${socketId}`);
+      }
+      
+      // Close the socket gracefully
+      try {
+        socket.end();
+        
+        // Set a forced close timeout in case socket.end() doesn't close the connection
+        setTimeout(() => {
+          if (!socket.destroyed) {
+            SmtpLogger.warn(`Forcing destroy of timed out socket: ${socketId}`);
+            socket.destroy();
+          }
+        }, 5000); // 5 second grace period for socket to end properly
+      } catch (error) {
+        SmtpLogger.error(`Error ending timed out socket: ${error instanceof Error ? error.message : String(error)}`);
+        
+        // Ensure socket is destroyed even if end() fails
+        if (!socket.destroyed) {
+          socket.destroy();
+        }
+      }
+      
+      // Clean up resources
+      this.activeConnections.delete(socket);
+      this.sessionManager.removeSession(socket);
+    } catch (handlerError) {
+      // Handle any unexpected errors during timeout handling
+      SmtpLogger.error(`Error in handleSocketTimeout: ${handlerError instanceof Error ? handlerError.message : String(handlerError)}`);
+      
+      // Ensure socket is destroyed and removed from tracking
+      if (!socket.destroyed) {
+        socket.destroy();
+      }
+      this.activeConnections.delete(socket);
     }
   }