dcrouter/test/test.smartproxy-ratelimit.node.ts

import { expect, tap } from '@git.zone/tstest/tapbundle';
import { SmartProxy } from '@push.rocks/smartproxy';
import { Buffer } from 'node:buffer';
import * as http from 'node:http';
import * as net from 'node:net';

async function getFreePort(): Promise<number> {
  return await new Promise<number>((resolve, reject) => {
    const server = net.createServer();
    server.once('error', reject);
    server.listen(0, '127.0.0.1', () => {
      const address = server.address();
      const port = typeof address === 'object' && address ? address.port : 0;
      server.close(() => resolve(port));
    });
  });
}

async function startBackend(
  handler: http.RequestListener = (_request, response) => {
    response.writeHead(200, { 'content-type': 'text/plain' });
    response.end('ok');
  },
): Promise<{ server: http.Server; port: number }> {
  const server = http.createServer(handler);
  const port = await new Promise<number>((resolve, reject) => {
    server.once('error', reject);
    server.listen(0, '127.0.0.1', () => {
      const address = server.address();
      resolve(typeof address === 'object' && address ? address.port : 0);
    });
  });
  return { server, port };
}

async function closeServer(server: http.Server): Promise<void> {
  if (!server.listening) return;
  await new Promise<void>((resolve, reject) => server.close((error) => error ? reject(error) : resolve()));
}

async function requestHeaders(
  port: number,
  path: string,
  headers?: Record<string, string>,
): Promise<http.IncomingMessage> {
  return await new Promise<http.IncomingMessage>((resolve, reject) => {
    const request = http.get({ host: '127.0.0.1', port, path, headers, agent: false }, resolve);
    request.once('error', reject);
  });
}

async function readResponseBody(response: http.IncomingMessage): Promise<string> {
  const chunks: Buffer[] = [];
  for await (const chunk of response) {
    chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk));
  }
  return Buffer.concat(chunks).toString('utf8');
}

tap.test('SmartProxy route rateLimit returns 429 after threshold', async () => {
  const backend = await startBackend();
  const proxyPort = await getFreePort();
  const proxy = new SmartProxy({
    connectionRateLimitPerMinute: 1000,
    routes: [
      {
        name: 'rate-limit-smoke',
        match: {
          ports: proxyPort,
        },
        action: {
          type: 'forward',
          targets: [{ host: '127.0.0.1', port: backend.port }],
        },
        security: {
          rateLimit: {
            enabled: true,
            maxRequests: 1,
            window: 60,
            keyBy: 'ip',
            errorMessage: 'too many requests',
          },
        },
      },
    ],
  });

  try {
    await proxy.start();
    const firstResponse = await fetch(`http://127.0.0.1:${proxyPort}/`);
    const secondResponse = await fetch(`http://127.0.0.1:${proxyPort}/`);
    const firstBody = await firstResponse.text();
    const secondBody = await secondResponse.text();

    expect(firstResponse.status).toEqual(200);
    expect(firstBody).toEqual('ok');
    expect(secondResponse.status).toEqual(429);
    expect(secondBody).toContain('too many requests');
  } finally {
    await Promise.allSettled([
      proxy.stop(),
      closeServer(backend.server),
    ]);
  }
});

tap.test('SmartProxy rateLimit is terminal and does not fall through to a lower priority route', async () => {
  const limitedBackend = await startBackend((_request, response) => {
    response.writeHead(200, { 'content-type': 'text/plain' });
    response.end('limited');
  });
  const fallbackBackend = await startBackend((_request, response) => {
    response.writeHead(200, { 'content-type': 'text/plain' });
    response.end('fallback');
  });
  const proxyPort = await getFreePort();
  const proxy = new SmartProxy({
    connectionRateLimitPerMinute: 1000,
    routes: [
      {
        id: 'terminal-rate-limit',
        name: 'terminal-rate-limit',
        priority: 10,
        match: { ports: proxyPort, domains: 'limited.local' },
        action: {
          type: 'forward',
          targets: [{ host: '127.0.0.1', port: limitedBackend.port }],
        },
        security: {
          rateLimit: {
            enabled: true,
            maxRequests: 1,
            window: 60,
            keyBy: 'ip',
            errorMessage: 'limited route exceeded',
          },
        },
      },
      {
        id: 'lower-priority-fallback',
        name: 'lower-priority-fallback',
        priority: 0,
        match: { ports: proxyPort },
        action: {
          type: 'forward',
          targets: [{ host: '127.0.0.1', port: fallbackBackend.port }],
        },
      },
    ],
  });

  try {
    await proxy.start();
    const firstResponse = await requestHeaders(proxyPort, '/', { host: 'limited.local' });
    const secondResponse = await requestHeaders(proxyPort, '/', { host: 'limited.local' });
    const firstBody = await readResponseBody(firstResponse);
    const secondBody = await readResponseBody(secondResponse);

    expect(firstResponse.statusCode).toEqual(200);
    expect(firstBody).toEqual('limited');
    expect(secondResponse.statusCode).toEqual(429);
    expect(secondBody).toContain('limited route exceeded');
    expect(secondBody.includes('fallback')).toBeFalse();
  } finally {
    await Promise.allSettled([
      proxy.stop(),
      closeServer(limitedBackend.server),
      closeServer(fallbackBackend.server),
    ]);
  }
});

tap.test('SmartProxy route maxConnections returns 429 when concurrent limit is exceeded', async () => {
  let firstResponse: http.IncomingMessage | undefined;
  let secondResponse: http.IncomingMessage | undefined;
  let releaseResponse: (() => void) | undefined;
  const releasePromise = new Promise<void>((resolve) => {
    releaseResponse = resolve;
  });
  const backend = await startBackend((_request, response) => {
    response.writeHead(200, { 'content-type': 'text/plain' });
    response.flushHeaders();
    void releasePromise.then(() => response.end('released'));
  });
  const proxyPort = await getFreePort();
  const proxy = new SmartProxy({
    connectionRateLimitPerMinute: 1000,
    routes: [
      {
        id: 'max-connections-smoke',
        name: 'max-connections-smoke',
        match: { ports: proxyPort },
        action: {
          type: 'forward',
          targets: [{ host: '127.0.0.1', port: backend.port }],
        },
        security: {
          maxConnections: 1,
        },
      },
    ],
  });

  try {
    await proxy.start();
    firstResponse = await requestHeaders(proxyPort, '/hold');
    secondResponse = await requestHeaders(proxyPort, '/blocked');

    expect(firstResponse.statusCode).toEqual(200);
    expect(secondResponse.statusCode).toEqual(429);
    const secondBody = await readResponseBody(secondResponse);
    releaseResponse?.();
    expect(await readResponseBody(firstResponse)).toEqual('released');
    expect(secondBody.length > 0).toBeTrue();
  } finally {
    releaseResponse?.();
    firstResponse?.destroy();
    secondResponse?.destroy();
    await Promise.allSettled([
      proxy.stop(),
      closeServer(backend.server),
    ]);
  }
});

tap.test('SmartProxy maxConnections is terminal and does not fall through to a lower priority route', async () => {
  let firstResponse: http.IncomingMessage | undefined;
  let secondResponse: http.IncomingMessage | undefined;
  let releaseResponse: (() => void) | undefined;
  const releasePromise = new Promise<void>((resolve) => {
    releaseResponse = resolve;
  });
  const limitedBackend = await startBackend((_request, response) => {
    response.writeHead(200, { 'content-type': 'text/plain' });
    response.flushHeaders();
    void releasePromise.then(() => response.end('limited released'));
  });
  const fallbackBackend = await startBackend((_request, response) => {
    response.writeHead(200, { 'content-type': 'text/plain' });
    response.end('fallback');
  });
  const proxyPort = await getFreePort();
  const proxy = new SmartProxy({
    connectionRateLimitPerMinute: 1000,
    routes: [
      {
        id: 'terminal-max-connections',
        name: 'terminal-max-connections',
        priority: 10,
        match: { ports: proxyPort, domains: 'limited.local' },
        action: {
          type: 'forward',
          targets: [{ host: '127.0.0.1', port: limitedBackend.port }],
        },
        security: {
          maxConnections: 1,
        },
      },
      {
        id: 'max-connections-lower-priority-fallback',
        name: 'max-connections-lower-priority-fallback',
        priority: 0,
        match: { ports: proxyPort },
        action: {
          type: 'forward',
          targets: [{ host: '127.0.0.1', port: fallbackBackend.port }],
        },
      },
    ],
  });

  try {
    await proxy.start();
    firstResponse = await requestHeaders(proxyPort, '/hold', { host: 'limited.local' });
    secondResponse = await requestHeaders(proxyPort, '/blocked', { host: 'limited.local' });
    const secondBody = await readResponseBody(secondResponse);
    releaseResponse?.();
    const firstBody = await readResponseBody(firstResponse);

    expect(firstResponse.statusCode).toEqual(200);
    expect(firstBody).toEqual('limited released');
    expect(secondResponse.statusCode).toEqual(429);
    expect(secondBody.includes('fallback')).toBeFalse();
  } finally {
    releaseResponse?.();
    firstResponse?.destroy();
    secondResponse?.destroy();
    await Promise.allSettled([
      proxy.stop(),
      closeServer(limitedBackend.server),
      closeServer(fallbackBackend.server),
    ]);
  }
});

export default tap.start();