smartproxy/test/test.udp-forwarding.ts

import { expect, tap } from '@git.zone/tstest/tapbundle';
import * as dgram from 'dgram';
import { SmartProxy } from '../ts/index.js';
import { findFreePorts, assertPortsFree } from './helpers/port-allocator.js';

let smartProxy: SmartProxy;
let backendServer: dgram.Socket;
let PROXY_PORT: number;
let BACKEND_PORT: number;

// Helper: send a single UDP datagram and wait for a response
function sendDatagram(port: number, msg: string, timeoutMs = 5000): Promise<string> {
  return new Promise((resolve, reject) => {
    const client = dgram.createSocket('udp4');
    const timeout = setTimeout(() => {
      client.close();
      reject(new Error(`UDP response timeout after ${timeoutMs}ms`));
    }, timeoutMs);
    client.send(Buffer.from(msg), port, '127.0.0.1');
    client.on('message', (data) => {
      clearTimeout(timeout);
      client.close();
      resolve(data.toString());
    });
    client.on('error', (err) => {
      clearTimeout(timeout);
      client.close();
      reject(err);
    });
  });
}

// Helper: create a UDP echo server
function createUdpEchoServer(port: number): Promise<dgram.Socket> {
  return new Promise((resolve) => {
    const server = dgram.createSocket('udp4');
    server.on('message', (msg, rinfo) => {
      server.send(Buffer.from(`Echo: ${msg.toString()}`), rinfo.port, rinfo.address);
    });
    server.bind(port, '127.0.0.1', () => resolve(server));
  });
}

tap.test('setup: start UDP echo server and SmartProxy', async () => {
  [PROXY_PORT, BACKEND_PORT] = await findFreePorts(2);

  // Start backend UDP echo server
  backendServer = await createUdpEchoServer(BACKEND_PORT);

  // Start SmartProxy with a UDP forwarding route
  smartProxy = new SmartProxy({
    routes: [
      {
        name: 'udp-forward-test',
        match: {
          ports: PROXY_PORT,
          transport: 'udp' as const,
        },
        action: {
          type: 'forward',
          targets: [{ host: '127.0.0.1', port: BACKEND_PORT }],
          udp: {
            sessionTimeout: 5000,
          },
        },
      },
    ],
    defaults: {
      security: {
        ipAllowList: ['127.0.0.1', '::1', '::ffff:127.0.0.1'],
      },
    },
  });

  await smartProxy.start();
});

tap.test('UDP forwarding: basic datagram round-trip', async () => {
  const response = await sendDatagram(PROXY_PORT, 'Hello UDP');
  expect(response).toEqual('Echo: Hello UDP');
});

tap.test('UDP forwarding: multiple datagrams same session', async () => {
  // Use a single client socket for session reuse
  const client = dgram.createSocket('udp4');
  const responses: string[] = [];

  const done = new Promise<void>((resolve, reject) => {
    const timeout = setTimeout(() => {
      client.close();
      reject(new Error('Timeout waiting for 3 responses'));
    }, 5000);

    client.on('message', (data) => {
      responses.push(data.toString());
      if (responses.length === 3) {
        clearTimeout(timeout);
        client.close();
        resolve();
      }
    });
    client.on('error', (err) => {
      clearTimeout(timeout);
      client.close();
      reject(err);
    });
  });

  client.send(Buffer.from('msg1'), PROXY_PORT, '127.0.0.1');
  client.send(Buffer.from('msg2'), PROXY_PORT, '127.0.0.1');
  client.send(Buffer.from('msg3'), PROXY_PORT, '127.0.0.1');

  await done;

  expect(responses).toContain('Echo: msg1');
  expect(responses).toContain('Echo: msg2');
  expect(responses).toContain('Echo: msg3');
});

tap.test('UDP forwarding: multiple clients', async () => {
  const [resp1, resp2] = await Promise.all([
    sendDatagram(PROXY_PORT, 'client1'),
    sendDatagram(PROXY_PORT, 'client2'),
  ]);

  expect(resp1).toEqual('Echo: client1');
  expect(resp2).toEqual('Echo: client2');
});

tap.test('UDP forwarding: large datagram (1400 bytes)', async () => {
  const payload = 'X'.repeat(1400);
  const response = await sendDatagram(PROXY_PORT, payload);
  expect(response).toEqual(`Echo: ${payload}`);
});

tap.test('cleanup: stop SmartProxy and backend', async () => {
  await smartProxy.stop();
  await new Promise<void>((resolve) => backendServer.close(() => resolve()));
  await assertPortsFree([PROXY_PORT, BACKEND_PORT]);
});

export default tap.start();