dcrouter/test/suite/smtpserver_connection/test.cm-04.connection-limits.ts

import { tap, expect } from '@git.zone/tstest/tapbundle';
import * as net from 'net';
import { startTestServer, stopTestServer } from '../../helpers/server.loader.js'
import type { ITestServer } from '../../helpers/server.loader.js';
// Test configuration
const TEST_PORT = 2525;
const TEST_TIMEOUT = 5000;

let testServer: ITestServer;

// Setup
tap.test('setup - start SMTP server', async () => {
  testServer = await startTestServer({ port: TEST_PORT });
  await new Promise(resolve => setTimeout(resolve, 1000));
  
});

// Test: Basic connection limit enforcement
tap.test('Connection Limits - should handle multiple connections gracefully', async (tools) => {
  const done = tools.defer();
  
  const maxConnections = 20; // Test with reasonable number
  const testConnections = maxConnections + 5; // Try 5 more than limit
  const connections: net.Socket[] = [];
  const connectionPromises: Promise<{ index: number; success: boolean; error?: string }>[] = [];
  
  // Helper to create a connection with index
  const createConnectionWithIndex = (index: number): Promise<{ index: number; success: boolean; error?: string }> => {
    return new Promise((resolve) => {
      let timeoutHandle: NodeJS.Timeout;
      
      try {
        const socket = net.createConnection({
          host: 'localhost',
          port: TEST_PORT,
          timeout: TEST_TIMEOUT
        });
        
        socket.on('connect', () => {
          clearTimeout(timeoutHandle);
          connections[index] = socket;
          
          // Wait for server greeting
          socket.on('data', (data) => {
            if (data.toString().includes('220')) {
              resolve({ index, success: true });
            }
          });
        });
        
        socket.on('error', (err) => {
          clearTimeout(timeoutHandle);
          resolve({ index, success: false, error: err.message });
        });
        
        timeoutHandle = setTimeout(() => {
          socket.destroy();
          resolve({ index, success: false, error: 'Connection timeout' });
        }, TEST_TIMEOUT);
      } catch (err: any) {
        resolve({ index, success: false, error: err.message });
      }
    });
  };
  
  // Create connections
  for (let i = 0; i < testConnections; i++) {
    connectionPromises.push(createConnectionWithIndex(i));
  }
  
  const results = await Promise.all(connectionPromises);
  
  // Count successful connections
  const successfulConnections = results.filter(r => r.success).length;
  const failedConnections = results.filter(r => !r.success).length;
  
  // Clean up connections
  for (const socket of connections) {
    if (socket && !socket.destroyed) {
      socket.write('QUIT\r\n');
      setTimeout(() => socket.destroy(), 100);
    }
  }
  
  // Verify results
  expect(successfulConnections).toBeGreaterThan(0);
  
  // If some connections were rejected, that's good (limit enforced)
  // If all connections succeeded, that's also acceptable (high/no limit)
  if (failedConnections > 0) {
    console.log(`Server enforced connection limit: ${successfulConnections} accepted, ${failedConnections} rejected`);
  } else {
    console.log(`Server accepted all ${successfulConnections} connections`);
  }
  
  done.resolve();
  
  await done.promise;
});

// Test: Connection limit recovery
tap.test('Connection Limits - should accept new connections after closing old ones', async (tools) => {
  const done = tools.defer();
  
  const batchSize = 10;
  const firstBatch: net.Socket[] = [];
  const secondBatch: net.Socket[] = [];
  
  // Create first batch of connections
  const firstBatchPromises = [];
  for (let i = 0; i < batchSize; i++) {
    firstBatchPromises.push(
      new Promise<boolean>((resolve) => {
        const socket = net.createConnection({
          host: 'localhost',
          port: TEST_PORT,
          timeout: TEST_TIMEOUT
        });
        
        socket.on('connect', () => {
          firstBatch.push(socket);
          socket.on('data', (data) => {
            if (data.toString().includes('220')) {
              resolve(true);
            }
          });
        });
        
        socket.on('error', () => resolve(false));
      })
    );
  }
  
  const firstResults = await Promise.all(firstBatchPromises);
  const firstSuccessCount = firstResults.filter(r => r).length;
  
  // Close first batch
  for (const socket of firstBatch) {
    if (socket && !socket.destroyed) {
      socket.write('QUIT\r\n');
    }
  }
  
  // Wait for connections to close
  await new Promise(resolve => setTimeout(resolve, 1000));
  
  // Destroy sockets
  for (const socket of firstBatch) {
    if (socket && !socket.destroyed) {
      socket.destroy();
    }
  }
  
  // Create second batch
  const secondBatchPromises = [];
  for (let i = 0; i < batchSize; i++) {
    secondBatchPromises.push(
      new Promise<boolean>((resolve) => {
        const socket = net.createConnection({
          host: 'localhost',
          port: TEST_PORT,
          timeout: TEST_TIMEOUT
        });
        
        socket.on('connect', () => {
          secondBatch.push(socket);
          socket.on('data', (data) => {
            if (data.toString().includes('220')) {
              resolve(true);
            }
          });
        });
        
        socket.on('error', () => resolve(false));
      })
    );
  }
  
  const secondResults = await Promise.all(secondBatchPromises);
  const secondSuccessCount = secondResults.filter(r => r).length;
  
  // Clean up second batch
  for (const socket of secondBatch) {
    if (socket && !socket.destroyed) {
      socket.write('QUIT\r\n');
      setTimeout(() => socket.destroy(), 100);
    }
  }
  
  // Both batches should have successful connections
  expect(firstSuccessCount).toBeGreaterThan(0);
  expect(secondSuccessCount).toBeGreaterThan(0);
  
  done.resolve();
  
  await done.promise;
});

// Test: Rapid connection attempts
tap.test('Connection Limits - should handle rapid connection attempts', async (tools) => {
  const done = tools.defer();
  
  const rapidConnections = 50;
  const connections: net.Socket[] = [];
  let successCount = 0;
  let errorCount = 0;
  
  // Create connections as fast as possible
  for (let i = 0; i < rapidConnections; i++) {
    const socket = net.createConnection({
      host: 'localhost',
      port: TEST_PORT
    });
    
    socket.on('connect', () => {
      connections.push(socket);
      successCount++;
    });
    
    socket.on('error', () => {
      errorCount++;
    });
  }
  
  // Wait for all connection attempts to settle
  await new Promise(resolve => setTimeout(resolve, 3000));
  
  // Clean up
  for (const socket of connections) {
    if (socket && !socket.destroyed) {
      socket.destroy();
    }
  }
  
  // Should handle at least some connections
  expect(successCount).toBeGreaterThan(0);
  console.log(`Rapid connections: ${successCount} succeeded, ${errorCount} failed`);
  
  done.resolve();
  
  await done.promise;
});

// Test: Connection limit with different client IPs (simulated)
tap.test('Connection Limits - should track connections per IP or globally', async (tools) => {
  const done = tools.defer();
  
  // Note: In real test, this would use different source IPs
  // For now, we test from same IP but document the behavior
  const connectionsPerIP = 5;
  const connections: net.Socket[] = [];
  const results: boolean[] = [];
  
  for (let i = 0; i < connectionsPerIP; i++) {
    const result = await new Promise<boolean>((resolve) => {
      const socket = net.createConnection({
        host: 'localhost',
        port: TEST_PORT,
        timeout: TEST_TIMEOUT
      });
      
      socket.on('connect', () => {
        connections.push(socket);
        socket.on('data', (data) => {
          if (data.toString().includes('220')) {
            resolve(true);
          }
        });
      });
      
      socket.on('error', () => resolve(false));
    });
    
    results.push(result);
  }
  
  const successCount = results.filter(r => r).length;
  
  // Clean up
  for (const socket of connections) {
    if (socket && !socket.destroyed) {
      socket.write('QUIT\r\n');
      setTimeout(() => socket.destroy(), 100);
    }
  }
  
  // Should accept connections from same IP
  expect(successCount).toBeGreaterThan(0);
  console.log(`Per-IP connections: ${successCount} of ${connectionsPerIP} succeeded`);
  
  done.resolve();
  
  await done.promise;
});

// Test: Connection limit error messages
tap.test('Connection Limits - should provide meaningful error when limit reached', async (tools) => {
  const done = tools.defer();
  
  const manyConnections = 100;
  const connections: net.Socket[] = [];
  const errors: string[] = [];
  let rejected = false;
  
  // Create many connections to try to hit limit
  const promises = [];
  for (let i = 0; i < manyConnections; i++) {
    promises.push(
      new Promise<void>((resolve) => {
        const socket = net.createConnection({
          host: 'localhost',
          port: TEST_PORT,
          timeout: 1000
        });
        
        socket.on('connect', () => {
          connections.push(socket);
          
          socket.on('data', (data) => {
            const response = data.toString();
            // Check if server sends connection limit message
            if (response.includes('421') || response.includes('too many connections')) {
              rejected = true;
              errors.push(response);
            }
            resolve();
          });
        });
        
        socket.on('error', (err) => {
          if (err.message.includes('ECONNREFUSED') || err.message.includes('ECONNRESET')) {
            rejected = true;
            errors.push(err.message);
          }
          resolve();
        });
        
        socket.on('timeout', () => {
          resolve();
        });
      })
    );
  }
  
  await Promise.all(promises);
  
  // Clean up
  for (const socket of connections) {
    if (socket && !socket.destroyed) {
      socket.destroy();
    }
  }
  
  // Log results
  console.log(`Connection limit test: ${connections.length} connected, ${errors.length} rejected`);
  if (rejected) {
    console.log(`Sample rejection: ${errors[0]}`);
  }
  
  // Should have handled connections (either accepted or properly rejected)
  expect(connections.length + errors.length).toBeGreaterThan(0);
  
  done.resolve();
  
  await done.promise;
});

// Teardown
tap.test('teardown - stop SMTP server', async () => {
  await stopTestServer(testServer);
});

// Start the test
export default tap.start();