smartproxy/test/test.connection-limits.node.ts

import { expect, tap } from '@git.zone/tstest/tapbundle';
import * as net from 'net';
import { SmartProxy } from '../ts/proxies/smart-proxy/index.js';
import { HttpProxy } from '../ts/proxies/http-proxy/index.js';

let testServer: net.Server;
let smartProxy: SmartProxy;
let httpProxy: HttpProxy;
const TEST_SERVER_PORT = 5100;
const PROXY_PORT = 5101;
const HTTP_PROXY_PORT = 5102;

// Track all created servers and connections for cleanup
const allServers: net.Server[] = [];
const allProxies: (SmartProxy | HttpProxy)[] = [];
const activeConnections: net.Socket[] = [];

// Helper: Creates a test TCP server
function createTestServer(port: number): Promise<net.Server> {
  return new Promise((resolve) => {
    const server = net.createServer((socket) => {
      socket.on('data', (data) => {
        socket.write(`Echo: ${data.toString()}`);
      });
      socket.on('error', () => {});
    });
    server.listen(port, 'localhost', () => {
      console.log(`[Test Server] Listening on localhost:${port}`);
      allServers.push(server);
      resolve(server);
    });
  });
}

// Helper: Creates multiple concurrent connections
async function createConcurrentConnections(
  port: number,
  count: number,
  fromIP?: string
): Promise<net.Socket[]> {
  const connections: net.Socket[] = [];
  const promises: Promise<net.Socket>[] = [];

  for (let i = 0; i < count; i++) {
    promises.push(
      new Promise((resolve, reject) => {
        const client = new net.Socket();
        const timeout = setTimeout(() => {
          client.destroy();
          reject(new Error(`Connection ${i} timeout`));
        }, 5000);

        client.connect(port, 'localhost', () => {
          clearTimeout(timeout);
          activeConnections.push(client);
          connections.push(client);
          resolve(client);
        });

        client.on('error', (err) => {
          clearTimeout(timeout);
          reject(err);
        });
      })
    );
  }

  await Promise.all(promises);
  return connections;
}

// Helper: Clean up connections
function cleanupConnections(connections: net.Socket[]): void {
  connections.forEach(conn => {
    if (!conn.destroyed) {
      conn.destroy();
    }
  });
}

tap.test('Setup test environment', async () => {
  testServer = await createTestServer(TEST_SERVER_PORT);
  
  // Create SmartProxy with low connection limits for testing
  smartProxy = new SmartProxy({
    routes: [{
      name: 'test-route',
      match: {
        ports: PROXY_PORT
      },
      action: {
        type: 'forward',
        target: {
          host: 'localhost',
          port: TEST_SERVER_PORT
        }
      },
      security: {
        maxConnections: 5 // Low limit for testing
      }
    }],
    maxConnectionsPerIP: 3, // Low per-IP limit
    connectionRateLimitPerMinute: 10, // Low rate limit
    defaults: {
      security: {
        maxConnections: 10 // Low global limit
      }
    }
  });
  
  await smartProxy.start();
  allProxies.push(smartProxy);
});

tap.test('Per-IP connection limits', async () => {
  // Test that we can create up to the per-IP limit
  const connections1 = await createConcurrentConnections(PROXY_PORT, 3);
  expect(connections1.length).toEqual(3);
  
  // Try to create one more connection - should fail
  try {
    await createConcurrentConnections(PROXY_PORT, 1);
    expect.fail('Should not allow more than 3 connections per IP');
  } catch (err) {
    expect(err.message).toInclude('ECONNRESET');
  }
  
  // Clean up first set of connections
  cleanupConnections(connections1);
  await new Promise(resolve => setTimeout(resolve, 100));
  
  // Should be able to create new connections after cleanup
  const connections2 = await createConcurrentConnections(PROXY_PORT, 2);
  expect(connections2.length).toEqual(2);
  
  cleanupConnections(connections2);
});

tap.test('Route-level connection limits', async () => {
  // Create multiple connections up to route limit
  const connections = await createConcurrentConnections(PROXY_PORT, 5);
  expect(connections.length).toEqual(5);
  
  // Try to exceed route limit
  try {
    await createConcurrentConnections(PROXY_PORT, 1);
    expect.fail('Should not allow more than 5 connections for this route');
  } catch (err) {
    expect(err.message).toInclude('ECONNRESET');
  }
  
  cleanupConnections(connections);
});

tap.test('Connection rate limiting', async () => {
  // Create connections rapidly
  const connections: net.Socket[] = [];
  
  // Create 10 connections rapidly (at rate limit)
  for (let i = 0; i < 10; i++) {
    try {
      const conn = await createConcurrentConnections(PROXY_PORT, 1);
      connections.push(...conn);
      // Small delay to avoid per-IP limit
      if (connections.length >= 3) {
        cleanupConnections(connections.splice(0, 3));
        await new Promise(resolve => setTimeout(resolve, 50));
      }
    } catch (err) {
      // Expected to fail at some point due to rate limit
      expect(i).toBeGreaterThan(0);
      break;
    }
  }
  
  cleanupConnections(connections);
});

tap.test('HttpProxy per-IP validation', async () => {
  // Create HttpProxy
  httpProxy = new HttpProxy({
    port: HTTP_PROXY_PORT,
    maxConnectionsPerIP: 2,
    connectionRateLimitPerMinute: 10,
    routes: []
  });
  
  await httpProxy.start();
  allProxies.push(httpProxy);
  
  // Update SmartProxy to use HttpProxy for TLS termination
  await smartProxy.stop();
  smartProxy = new SmartProxy({
    routes: [{
      name: 'https-route',
      match: {
        ports: PROXY_PORT + 10
      },
      action: {
        type: 'forward',
        target: {
          host: 'localhost',
          port: TEST_SERVER_PORT
        },
        tls: {
          mode: 'terminate'
        }
      }
    }],
    useHttpProxy: [PROXY_PORT + 10],
    httpProxyPort: HTTP_PROXY_PORT,
    maxConnectionsPerIP: 3
  });
  
  await smartProxy.start();
  
  // Test that HttpProxy enforces its own per-IP limits
  const connections = await createConcurrentConnections(PROXY_PORT + 10, 2);
  expect(connections.length).toEqual(2);
  
  // Should reject additional connections
  try {
    await createConcurrentConnections(PROXY_PORT + 10, 1);
    expect.fail('HttpProxy should enforce per-IP limits');
  } catch (err) {
    expect(err.message).toInclude('ECONNRESET');
  }
  
  cleanupConnections(connections);
});

tap.test('IP tracking cleanup', async (tools) => {
  // Create and close many connections from different IPs
  const connections: net.Socket[] = [];
  
  for (let i = 0; i < 5; i++) {
    const conn = await createConcurrentConnections(PROXY_PORT, 1);
    connections.push(...conn);
  }
  
  // Close all connections
  cleanupConnections(connections);
  
  // Wait for cleanup interval (set to 60s in production, but we'll check immediately)
  await tools.delayFor(100);
  
  // Verify that IP tracking has been cleaned up
  const securityManager = (smartProxy as any).securityManager;
  const ipCount = (securityManager.connectionsByIP as Map<string, any>).size;
  
  // Should have no IPs tracked after cleanup
  expect(ipCount).toEqual(0);
});

tap.test('Cleanup queue race condition handling', async () => {
  // Create many connections concurrently to trigger batched cleanup
  const promises: Promise<net.Socket[]>[] = [];
  
  for (let i = 0; i < 20; i++) {
    promises.push(createConcurrentConnections(PROXY_PORT, 1).catch(() => []));
  }
  
  const results = await Promise.all(promises);
  const allConnections = results.flat();
  
  // Close all connections rapidly
  allConnections.forEach(conn => conn.destroy());
  
  // Give cleanup queue time to process
  await new Promise(resolve => setTimeout(resolve, 500));
  
  // Verify all connections were cleaned up
  const connectionManager = (smartProxy as any).connectionManager;
  const remainingConnections = connectionManager.getConnectionCount();
  
  expect(remainingConnections).toEqual(0);
});

tap.test('Cleanup and shutdown', async () => {
  // Clean up any remaining connections
  cleanupConnections(activeConnections);
  activeConnections.length = 0;
  
  // Stop all proxies
  for (const proxy of allProxies) {
    await proxy.stop();
  }
  allProxies.length = 0;
  
  // Close all test servers
  for (const server of allServers) {
    await new Promise<void>((resolve) => {
      server.close(() => resolve());
    });
  }
  allServers.length = 0;
});

tap.start();