smartvpn/test/test.flowcontrol.node.ts

import { tap, expect } from '@git.zone/tstest/tapbundle';
import * as net from 'net';
import { VpnClient, VpnServer } from '../ts/index.js';
import type { IVpnClientOptions, IVpnServerOptions, 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`);
}

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

let server: VpnServer;
let serverPort: number;
let keypair: IVpnKeypair;
let client: VpnClient;
const extraClients: VpnClient[] = [];

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

tap.test('setup: start VPN server', async () => {
  serverPort = await findFreePort();

  const options: IVpnServerOptions = {
    transport: { transport: 'stdio' },
  };
  server = new VpnServer(options);

  // Phase 1: start the daemon bridge
  const started = await server['bridge'].start();
  expect(started).toBeTrue();
  expect(server.running).toBeTrue();

  // Phase 2: generate a keypair
  keypair = await server.generateKeypair();
  expect(keypair.publicKey).toBeTypeofString();
  expect(keypair.privateKey).toBeTypeofString();

  // Phase 3: start the VPN listener
  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 });

  // Verify server is now running
  const status = await server.getStatus();
  expect(status.state).toEqual('connected');
});

tap.test('single client connects and gets IP', async () => {
  const options: IVpnClientOptions = {
    transport: { transport: 'stdio' },
  };
  client = new VpnClient(options);
  const started = await client.start();
  expect(started).toBeTrue();

  const result = await client.connect({
    serverUrl: `ws://127.0.0.1:${serverPort}`,
    serverPublicKey: keypair.publicKey,
    keepaliveIntervalSecs: 3,
  });

  expect(result.assignedIp).toBeTypeofString();
  expect(result.assignedIp).toStartWith('10.8.0.');

  // Verify client status
  const clientStatus = await client.getStatus();
  expect(clientStatus.state).toEqual('connected');

  // Verify server sees the client
  await waitFor(async () => {
    const clients = await server.listClients();
    return clients.length === 1;
  });
  const clients = await server.listClients();
  expect(clients.length).toEqual(1);
  expect(clients[0].assignedIp).toEqual(result.assignedIp);
});

tap.test('keepalive exchange', async () => {
  // Wait for at least 2 keepalive cycles (interval=3s, so 8s should be enough)
  await delay(8000);

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

  const serverStats = await server.getStatistics();
  expect(serverStats.keepalivesReceived).toBeGreaterThanOrEqual(1);
  expect(serverStats.keepalivesSent).toBeGreaterThanOrEqual(1);

  // Verify per-client keepalive tracking
  const clients = await server.listClients();
  expect(clients[0].keepalivesReceived).toBeGreaterThanOrEqual(1);
});

tap.test('connection quality telemetry', async () => {
  const quality = await client.getConnectionQuality();

  expect(quality.srttMs).toBeGreaterThanOrEqual(0);
  expect(quality.jitterMs).toBeTypeofNumber();
  expect(quality.minRttMs).toBeGreaterThanOrEqual(0);
  expect(quality.maxRttMs).toBeGreaterThanOrEqual(0);
  expect(quality.lossRatio).toBeTypeofNumber();
  expect(['healthy', 'degraded', 'critical']).toContain(quality.linkHealth);
});

tap.test('rate limiting: set and verify', async () => {
  const clients = await server.listClients();
  const clientId = clients[0].clientId;

  // Set a tight rate limit
  await server.setClientRateLimit(clientId, 100, 100);

  // Verify via telemetry
  const telemetry = await server.getClientTelemetry(clientId);
  expect(telemetry.rateLimitBytesPerSec).toEqual(100);
  expect(telemetry.burstBytes).toEqual(100);
  expect(telemetry.clientId).toEqual(clientId);
});

tap.test('rate limiting: removal', async () => {
  const clients = await server.listClients();
  const clientId = clients[0].clientId;

  await server.removeClientRateLimit(clientId);

  // Verify telemetry no longer shows rate limit
  const telemetry = await server.getClientTelemetry(clientId);
  expect(telemetry.rateLimitBytesPerSec).toBeNullOrUndefined();
  expect(telemetry.burstBytes).toBeNullOrUndefined();

  // Connection still healthy
  const status = await client.getStatus();
  expect(status.state).toEqual('connected');
});

tap.test('5 concurrent clients', async () => {
  const assignedIps = new Set<string>();

  // Get the first client's IP
  const existingClients = await server.listClients();
  assignedIps.add(existingClients[0].assignedIp);

  for (let i = 0; i < 5; i++) {
    const c = new VpnClient({ transport: { transport: 'stdio' } });
    await c.start();
    const result = await c.connect({
      serverUrl: `ws://127.0.0.1:${serverPort}`,
      serverPublicKey: keypair.publicKey,
      keepaliveIntervalSecs: 3,
    });
    expect(result.assignedIp).toStartWith('10.8.0.');
    assignedIps.add(result.assignedIp);
    extraClients.push(c);
  }

  // All IPs should be unique (6 total: original + 5 new)
  expect(assignedIps.size).toEqual(6);

  // Server should see 6 clients
  await waitFor(async () => {
    const clients = await server.listClients();
    return clients.length === 6;
  });
  const allClients = await server.listClients();
  expect(allClients.length).toEqual(6);
});

tap.test('client disconnect tracking', async () => {
  // Disconnect 3 of the 5 extra clients
  for (let i = 0; i < 3; i++) {
    const c = extraClients[i];
    await c.disconnect();
    c.stop();
  }

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

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

  const stats = await server.getStatistics();
  expect(stats.totalConnections).toBeGreaterThanOrEqual(6);
});

tap.test('server-side client disconnection', async () => {
  const clients = await server.listClients();
  // Pick one of the remaining extra clients (not the original)
  const targetClient = clients.find((c) => {
    // Find a client that belongs to extraClients[3] or extraClients[4]
    return c.clientId !== clients[0].clientId;
  });
  expect(targetClient).toBeTruthy();

  await server.disconnectClient(targetClient!.clientId);

  // Wait for server to update
  await waitFor(async () => {
    const remaining = await server.listClients();
    return remaining.length === 2;
  });

  const remaining = await server.listClients();
  expect(remaining.length).toEqual(2);
});

tap.test('teardown: stop all', async () => {
  // Stop the original client
  await client.disconnect();
  client.stop();

  // Stop remaining extra clients
  for (const c of extraClients) {
    if (c.running) {
      try {
        await c.disconnect();
      } catch {
        // May already be disconnected
      }
      c.stop();
    }
  }

  await delay(500);

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

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

export default tap.start();