smartvpn/test/test.loadtest.node.ts

import { tap, expect } from '@git.zone/tstest/tapbundle';
import * as net from 'net';
import * as stream from 'stream';
import { VpnClient, VpnServer } from '../ts/index.js';
import type { IVpnKeypair, IVpnServerConfig } from '../ts/index.js';

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

async function findFreePort(): Promise<number> {
  const server = net.createServer();
  await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', resolve));
  const port = (server.address() as net.AddressInfo).port;
  await new Promise<void>((resolve) => server.close(() => resolve()));
  return port;
}

function delay(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

async function waitFor(
  fn: () => Promise<boolean>,
  timeoutMs: number = 10000,
  pollMs: number = 500,
): Promise<void> {
  const deadline = Date.now() + timeoutMs;
  while (Date.now() < deadline) {
    if (await fn()) return;
    await delay(pollMs);
  }
  throw new Error(`waitFor timed out after ${timeoutMs}ms`);
}

// ---------------------------------------------------------------------------
// ThrottleProxy (adapted from remoteingress)
// ---------------------------------------------------------------------------

class ThrottleTransform extends stream.Transform {
  private bytesPerSec: number;
  private bucket: number;
  private lastRefill: number;
  private destroyed_: boolean = false;

  constructor(bytesPerSecond: number) {
    super();
    this.bytesPerSec = bytesPerSecond;
    this.bucket = bytesPerSecond;
    this.lastRefill = Date.now();
  }

  _transform(chunk: Buffer, _encoding: BufferEncoding, callback: stream.TransformCallback) {
    if (this.destroyed_) return;

    const now = Date.now();
    const elapsed = (now - this.lastRefill) / 1000;
    this.bucket = Math.min(this.bytesPerSec, this.bucket + elapsed * this.bytesPerSec);
    this.lastRefill = now;

    if (chunk.length <= this.bucket) {
      this.bucket -= chunk.length;
      callback(null, chunk);
    } else {
      const deficit = chunk.length - this.bucket;
      this.bucket = 0;
      const delayMs = Math.min((deficit / this.bytesPerSec) * 1000, 1000);
      setTimeout(() => {
        if (this.destroyed_) { callback(); return; }
        this.lastRefill = Date.now();
        this.bucket = 0;
        callback(null, chunk);
      }, delayMs);
    }
  }

  _destroy(err: Error | null, callback: (error: Error | null) => void) {
    this.destroyed_ = true;
    callback(err);
  }
}

interface ThrottleProxy {
  server: net.Server;
  close: () => Promise<void>;
}

async function startThrottleProxy(
  listenPort: number,
  targetHost: string,
  targetPort: number,
  bytesPerSecond: number,
): Promise<ThrottleProxy> {
  const connections = new Set<net.Socket>();
  const server = net.createServer((clientSock) => {
    connections.add(clientSock);
    const upstream = net.createConnection({ host: targetHost, port: targetPort });
    connections.add(upstream);

    const throttleUp = new ThrottleTransform(bytesPerSecond);
    const throttleDown = new ThrottleTransform(bytesPerSecond);

    clientSock.pipe(throttleUp).pipe(upstream);
    upstream.pipe(throttleDown).pipe(clientSock);

    let cleaned = false;
    const cleanup = () => {
      if (cleaned) return;
      cleaned = true;
      throttleUp.destroy();
      throttleDown.destroy();
      clientSock.destroy();
      upstream.destroy();
      connections.delete(clientSock);
      connections.delete(upstream);
    };
    clientSock.on('error', () => cleanup());
    upstream.on('error', () => cleanup());
    throttleUp.on('error', () => cleanup());
    throttleDown.on('error', () => cleanup());
    clientSock.on('close', () => cleanup());
    upstream.on('close', () => cleanup());
  });

  await new Promise<void>((resolve) => server.listen(listenPort, '127.0.0.1', resolve));
  return {
    server,
    close: async () => {
      for (const c of connections) c.destroy();
      connections.clear();
      await new Promise<void>((resolve) => server.close(() => resolve()));
    },
  };
}

// ---------------------------------------------------------------------------
// Test state
// ---------------------------------------------------------------------------

let server: VpnServer;
let serverPort: number;
let proxyPort: number;
let keypair: IVpnKeypair;
let throttle: ThrottleProxy;
const allClients: VpnClient[] = [];

async function createConnectedClient(port: number): Promise<VpnClient> {
  const c = new VpnClient({ transport: { transport: 'stdio' } });
  await c.start();
  await c.connect({
    serverUrl: `ws://127.0.0.1:${port}`,
    serverPublicKey: keypair.publicKey,
    keepaliveIntervalSecs: 3,
  });
  allClients.push(c);
  return c;
}

async function stopClient(c: VpnClient): Promise<void> {
  if (c.running) {
    try { await c.disconnect(); } catch { /* already disconnected */ }
    c.stop();
  }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

tap.test('setup: start throttled VPN tunnel (1 MB/s)', async () => {
  serverPort = await findFreePort();
  proxyPort = await findFreePort();

  // Start VPN server
  server = new VpnServer({ transport: { transport: 'stdio' } });
  const started = await server['bridge'].start();
  expect(started).toBeTrue();

  keypair = await server.generateKeypair();
  const serverConfig: IVpnServerConfig = {
    listenAddr: `127.0.0.1:${serverPort}`,
    privateKey: keypair.privateKey,
    publicKey: keypair.publicKey,
    subnet: '10.8.0.0/24',
  };
  await server['bridge'].sendCommand('start', { config: serverConfig });

  const status = await server.getStatus();
  expect(status.state).toEqual('connected');

  // Start throttle proxy: 1 MB/s
  throttle = await startThrottleProxy(proxyPort, '127.0.0.1', serverPort, 1024 * 1024);
});

tap.test('throttled connection: handshake succeeds through throttle', async () => {
  const client = await createConnectedClient(proxyPort);

  const status = await client.getStatus();
  expect(status.state).toEqual('connected');
  expect(status.assignedIp).toStartWith('10.8.0.');

  await waitFor(async () => {
    const clients = await server.listClients();
    return clients.length === 1;
  });
});

tap.test('sustained keepalive under throttle', async () => {
  // Wait for at least 2 keepalive cycles (3s interval)
  await delay(8000);

  const client = allClients[0];
  const stats = await client.getStatistics();
  expect(stats.keepalivesSent).toBeGreaterThanOrEqual(1);
  expect(stats.keepalivesReceived).toBeGreaterThanOrEqual(1);

  // Throttle adds latency — RTT should be measurable
  const quality = await client.getConnectionQuality();
  expect(quality.srttMs).toBeGreaterThanOrEqual(0);
  expect(quality.jitterMs).toBeTypeofNumber();
});

tap.test('3 concurrent throttled clients', async () => {
  for (let i = 0; i < 3; i++) {
    await createConnectedClient(proxyPort);
  }

  // All 4 clients should be visible
  await waitFor(async () => {
    const clients = await server.listClients();
    return clients.length === 4;
  });

  const clients = await server.listClients();
  expect(clients.length).toEqual(4);

  // Verify all IPs are unique
  const ips = new Set(clients.map((c) => c.assignedIp));
  expect(ips.size).toEqual(4);
});

tap.test('rate limiting combined with network throttle', async () => {
  const clients = await server.listClients();
  const targetId = clients[0].clientId;

  // Set rate limit on first client
  await server.setClientRateLimit(targetId, 500, 500);
  const telemetry = await server.getClientTelemetry(targetId);
  expect(telemetry.rateLimitBytesPerSec).toEqual(500);
  expect(telemetry.burstBytes).toEqual(500);

  // Verify another client has no rate limit
  const otherTelemetry = await server.getClientTelemetry(clients[1].clientId);
  expect(otherTelemetry.rateLimitBytesPerSec).toBeNullOrUndefined();

  // Clean up the rate limit
  await server.removeClientRateLimit(targetId);
});

tap.test('burst waves: 3 waves of 3 clients', async () => {
  const initialCount = (await server.listClients()).length;

  for (let wave = 0; wave < 3; wave++) {
    const waveClients: VpnClient[] = [];

    // Connect 3 clients
    for (let i = 0; i < 3; i++) {
      const c = await createConnectedClient(proxyPort);
      waveClients.push(c);
    }

    // Verify all connected
    await waitFor(async () => {
      const all = await server.listClients();
      return all.length === initialCount + 3;
    });

    // Disconnect all wave clients
    for (const c of waveClients) {
      await stopClient(c);
    }

    // Wait for server to detect disconnections
    await waitFor(async () => {
      const all = await server.listClients();
      return all.length === initialCount;
    }, 15000);

    await delay(500);
  }

  // Verify total connections accumulated
  const stats = await server.getStatistics();
  expect(stats.totalConnections).toBeGreaterThanOrEqual(9 + initialCount);

  // Original clients still connected
  const remaining = await server.listClients();
  expect(remaining.length).toEqual(initialCount);
});

tap.test('aggressive throttle: 10 KB/s', async () => {
  // Close current throttle proxy and start an aggressive one
  await throttle.close();
  const aggressivePort = await findFreePort();
  throttle = await startThrottleProxy(aggressivePort, '127.0.0.1', serverPort, 10 * 1024);

  // Connect a client through the aggressive throttle
  const client = await createConnectedClient(aggressivePort);
  const status = await client.getStatus();
  expect(status.state).toEqual('connected');

  // Wait for keepalive exchange (might take longer due to throttle)
  await delay(10000);

  const stats = await client.getStatistics();
  expect(stats.keepalivesSent).toBeGreaterThanOrEqual(1);
  expect(stats.keepalivesReceived).toBeGreaterThanOrEqual(1);
});

tap.test('post-load health: direct connection still works', async () => {
  // Server should still be healthy after all load tests
  const serverStatus = await server.getStatus();
  expect(serverStatus.state).toEqual('connected');

  // Connect one more client directly (no throttle)
  const directClient = await createConnectedClient(serverPort);
  const status = await directClient.getStatus();
  expect(status.state).toEqual('connected');

  await delay(5000);

  const stats = await directClient.getStatistics();
  expect(stats.keepalivesSent).toBeGreaterThanOrEqual(1);
});

tap.test('teardown: stop all', async () => {
  // Stop all clients
  for (const c of allClients) {
    await stopClient(c);
  }

  await delay(500);

  // Close throttle proxy
  if (throttle) {
    await throttle.close();
  }

  // Stop server
  await server.stopServer();
  server.stop();
  await delay(500);

  expect(server.running).toBeFalse();
});

export default tap.start();