smartproxy/test/test.port-mapping.ts

import { expect, tap } from '@git.zone/tstest/tapbundle';
import * as net from 'net';
import { SmartProxy } from '../ts/proxies/smart-proxy/index.js';
import type { IRouteConfig, IRouteContext } from '../ts/proxies/smart-proxy/models/route-types.js';
import { findFreePorts, assertPortsFree } from './helpers/port-allocator.js';

let testServers: Array<{ server: net.Server; port: number }> = [];
let smartProxy: SmartProxy;

let TEST_PORTS: number[];
let PROXY_PORTS: number[];
const TEST_DATA = 'Hello through dynamic port mapper!';

function cleanup() {
  console.log('Starting cleanup...');
  const promises = [];

  for (const { server, port } of testServers) {
    promises.push(new Promise<void>(resolve => {
      console.log(`Closing test server on port ${port}`);
      server.close(() => {
        console.log(`Test server on port ${port} closed`);
        resolve();
      });
    }));
  }

  if (smartProxy) {
    console.log('Stopping SmartProxy...');
    promises.push(smartProxy.stop().then(() => {
      console.log('SmartProxy stopped');
    }));
  }

  return Promise.all(promises);
}

function createTestServer(port: number): Promise<net.Server> {
  return new Promise((resolve) => {
    const server = net.createServer((socket) => {
      socket.on('data', (data) => {
        socket.write(`Server ${port} says: ${data.toString()}`);
      });
      socket.on('error', (error) => {
        console.error(`[Test Server] Socket error on port ${port}:`, error);
      });
    });

    server.listen(port, () => {
      console.log(`[Test Server] Listening on port ${port}`);
      testServers.push({ server, port });
      resolve(server);
    });
  });
}

function createTestClient(port: number, data: string): Promise<string> {
  return new Promise((resolve, reject) => {
    const client = new net.Socket();
    let response = '';

    const timeout = setTimeout(() => {
      client.destroy();
      reject(new Error(`Client connection timeout to port ${port}`));
    }, 5000);

    client.connect(port, 'localhost', () => {
      console.log(`[Test Client] Connected to server on port ${port}`);
      client.write(data);
    });

    client.on('data', (chunk) => {
      response += chunk.toString();
      client.end();
    });

    client.on('end', () => {
      clearTimeout(timeout);
      resolve(response);
    });

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

tap.test('setup port mapping test environment', async () => {
  const allPorts = await findFreePorts(9);
  TEST_PORTS = allPorts.slice(0, 3);
  PROXY_PORTS = allPorts.slice(3, 9);

  await Promise.all([
    createTestServer(TEST_PORTS[0]),
    createTestServer(TEST_PORTS[1]),
    createTestServer(TEST_PORTS[2]),
  ]);

  const portOffset = TEST_PORTS[1] - PROXY_PORTS[1];

  smartProxy = new SmartProxy({
    routes: [
      {
        match: { ports: PROXY_PORTS[0] },
        action: {
          type: 'forward',
          targets: [{
            host: 'localhost',
            port: (context: IRouteContext) => TEST_PORTS[0]
          }]
        },
        name: 'Identity Port Mapping'
      },

      {
        match: { ports: PROXY_PORTS[1] },
        action: {
          type: 'forward',
          targets: [{
            host: 'localhost',
            port: (context: IRouteContext) => context.port + portOffset
          }]
        },
        name: `Offset Port Mapping (${portOffset})`
      },

      {
        match: { ports: [PROXY_PORTS[2], PROXY_PORTS[3]] },
        action: {
          type: 'forward',
          targets: [{
            host: (context: IRouteContext) => {
              return context.port === PROXY_PORTS[2] ? 'localhost' : '127.0.0.1';
            },
            port: (context: IRouteContext) => {
              if (context.port === PROXY_PORTS[2]) {
                return TEST_PORTS[0];
              } else {
                return TEST_PORTS[2];
              }
            }
          }]
        },
        name: 'Dynamic Host and Port Mapping'
      },

      {
        match: { ports: PROXY_PORTS[4] },
        action: {
          type: 'forward',
          targets: [{
            host: (context: IRouteContext) => {
              if (context.domain === 'test1.example.com') return 'localhost';
              if (context.domain === 'test2.example.com') return '127.0.0.1';
              return 'localhost';
            },
            port: (context: IRouteContext) => {
              if (context.domain === 'test1.example.com') {
                return TEST_PORTS[0];
              } else {
                return TEST_PORTS[1];
              }
            }
          }]
        },
        name: 'Smart Domain Load Balancer'
      }
    ]
  });

  await smartProxy.start();
});

tap.test('should map port using identity function', async () => {
  const response = await createTestClient(PROXY_PORTS[0], TEST_DATA);
  expect(response).toEqual(`Server ${TEST_PORTS[0]} says: ${TEST_DATA}`);
});

tap.test('should map port using offset function', async () => {
  const response = await createTestClient(PROXY_PORTS[1], TEST_DATA);
  expect(response).toEqual(`Server ${TEST_PORTS[1]} says: ${TEST_DATA}`);
});

tap.test('should map port using dynamic logic', async () => {
  const response = await createTestClient(PROXY_PORTS[2], TEST_DATA);
  expect(response).toEqual(`Server ${TEST_PORTS[0]} says: ${TEST_DATA}`);
});

tap.test('should handle errors in port mapping functions', async () => {
  const errorRoute: IRouteConfig = {
    match: {
      ports: PROXY_PORTS[5]
    },
    action: {
      type: 'forward',
      targets: [{
        host: 'localhost',
        port: () => {
          throw new Error('Test error in port mapping function');
        }
      }]
    },
    name: 'Error Route'
  };

  await smartProxy.updateRoutes([...smartProxy.settings.routes, errorRoute]);

  try {
    await createTestClient(PROXY_PORTS[5], TEST_DATA);
    expect(false).toBeTrue();
  } catch (error) {
    expect(true).toBeTrue();
  }
});

tap.test('cleanup port mapping test environment', async () => {
  const cleanupPromise = cleanup();
  const timeoutPromise = new Promise((_, reject) =>
    setTimeout(() => reject(new Error('Cleanup timeout after 5 seconds')), 5000)
  );

  try {
    await Promise.race([cleanupPromise, timeoutPromise]);
  } catch (error) {
    console.error('Cleanup error:', error);
    testServers = [];
    smartProxy = null as any;
  }

  await assertPortsFree([...TEST_PORTS, ...PROXY_PORTS]);
});

export default tap.start();