smartproxy/test/test.ts

import { expect, tap } from '@push.rocks/tapbundle';
import * as smartproxy from '../ts/index.js';
import { loadTestCertificates } from './helpers/certificates.js';
import * as https from 'https';
import * as http from 'http';
import { WebSocket, WebSocketServer } from 'ws';

let testProxy: smartproxy.NetworkProxy;
let testServer: http.Server;
let wsServer: WebSocketServer;
let testCertificates: { privateKey: string; publicKey: string };

// Helper function to make HTTPS requests
async function makeHttpsRequest(
  options: https.RequestOptions,
): Promise<{ statusCode: number; headers: http.IncomingHttpHeaders; body: string }> {
  console.log('[TEST] Making HTTPS request:', {
    hostname: options.hostname,
    port: options.port,
    path: options.path,
    method: options.method,
    headers: options.headers,
  });
  return new Promise((resolve, reject) => {
    const req = https.request(options, (res) => {
      console.log('[TEST] Received HTTPS response:', {
        statusCode: res.statusCode,
        headers: res.headers,
      });
      let data = '';
      res.on('data', (chunk) => (data += chunk));
      res.on('end', () => {
        console.log('[TEST] Response completed:', { data });
        resolve({
          statusCode: res.statusCode!,
          headers: res.headers,
          body: data,
        });
      });
    });
    req.on('error', (error) => {
      console.error('[TEST] Request error:', error);
      reject(error);
    });
    req.end();
  });
}

// Setup test environment
tap.test('setup test environment', async () => {
  // Load and validate certificates
  console.log('[TEST] Loading and validating certificates');
  testCertificates = loadTestCertificates();
  console.log('[TEST] Certificates loaded and validated');

  // Create a test HTTP server
  testServer = http.createServer((req, res) => {
    console.log('[TEST SERVER] Received HTTP request:', {
      url: req.url,
      method: req.method,
      headers: req.headers,
    });
    res.writeHead(200, { 'Content-Type': 'text/plain' });
    res.end('Hello from test server!');
  });

  // Handle WebSocket upgrade requests
  testServer.on('upgrade', (request, socket, head) => {
    console.log('[TEST SERVER] Received WebSocket upgrade request:', {
      url: request.url,
      method: request.method,
      headers: {
        host: request.headers.host,
        upgrade: request.headers.upgrade,
        connection: request.headers.connection,
        'sec-websocket-key': request.headers['sec-websocket-key'],
        'sec-websocket-version': request.headers['sec-websocket-version'],
        'sec-websocket-protocol': request.headers['sec-websocket-protocol'],
      },
    });

    if (request.headers.upgrade?.toLowerCase() !== 'websocket') {
      console.log('[TEST SERVER] Not a WebSocket upgrade request');
      socket.destroy();
      return;
    }

    console.log('[TEST SERVER] Handling WebSocket upgrade');
    wsServer.handleUpgrade(request, socket, head, (ws) => {
      console.log('[TEST SERVER] WebSocket connection upgraded');
      wsServer.emit('connection', ws, request);
    });
  });

  // Create a WebSocket server (for the test HTTP server)
  console.log('[TEST SERVER] Creating WebSocket server');
  wsServer = new WebSocketServer({
    noServer: true,
    perMessageDeflate: false,
    clientTracking: true,
    handleProtocols: () => 'echo-protocol',
  });

  wsServer.on('connection', (ws, request) => {
    console.log('[TEST SERVER] WebSocket connection established:', {
      url: request.url,
      headers: {
        host: request.headers.host,
        upgrade: request.headers.upgrade,
        connection: request.headers.connection,
        'sec-websocket-key': request.headers['sec-websocket-key'],
        'sec-websocket-version': request.headers['sec-websocket-version'],
        'sec-websocket-protocol': request.headers['sec-websocket-protocol'],
      },
    });

    // Set up connection timeout
    const connectionTimeout = setTimeout(() => {
      console.error('[TEST SERVER] WebSocket connection timed out');
      ws.terminate();
    }, 5000);

    // Clear timeout when connection is properly closed
    const clearConnectionTimeout = () => {
      clearTimeout(connectionTimeout);
    };

    ws.on('message', (message) => {
      const msg = message.toString();
      console.log('[TEST SERVER] Received message:', msg);
      try {
        const response = `Echo: ${msg}`;
        console.log('[TEST SERVER] Sending response:', response);
        ws.send(response);
        // Clear timeout on successful message exchange
        clearConnectionTimeout();
      } catch (error) {
        console.error('[TEST SERVER] Error sending message:', error);
      }
    });

    ws.on('error', (error) => {
      console.error('[TEST SERVER] WebSocket error:', error);
      clearConnectionTimeout();
    });

    ws.on('close', (code, reason) => {
      console.log('[TEST SERVER] WebSocket connection closed:', {
        code,
        reason: reason.toString(),
        wasClean: code === 1000 || code === 1001,
      });
      clearConnectionTimeout();
    });

    ws.on('ping', (data) => {
      try {
        console.log('[TEST SERVER] Received ping, sending pong');
        ws.pong(data);
      } catch (error) {
        console.error('[TEST SERVER] Error sending pong:', error);
      }
    });

    ws.on('pong', (data) => {
      console.log('[TEST SERVER] Received pong');
    });
  });

  wsServer.on('error', (error) => {
    console.error('Test server: WebSocket server error:', error);
  });

  wsServer.on('headers', (headers) => {
    console.log('Test server: WebSocket headers:', headers);
  });

  wsServer.on('close', () => {
    console.log('Test server: WebSocket server closed');
  });

  await new Promise<void>((resolve) => testServer.listen(3000, resolve));
  console.log('Test server listening on port 3000');
});

tap.test('should create proxy instance', async () => {
  testProxy = new smartproxy.NetworkProxy({
    port: 3001,
  });
  expect(testProxy).toEqual(testProxy); // Instance equality check
});

tap.test('should start the proxy server', async () => {
  // Ensure any previous server is closed
  if (testProxy && testProxy.httpsServer) {
    await new Promise<void>((resolve) =>
      testProxy.httpsServer.close(() => resolve())
    );
  }

  console.log('[TEST] Starting the proxy server');
  await testProxy.start();
  console.log('[TEST] Proxy server started');

  // Configure proxy with test certificates
  // Awaiting the update ensures that the SNI context is added before any requests come in.
  await testProxy.updateProxyConfigs([
    {
      destinationIp: '127.0.0.1',
      destinationPort: '3000',
      hostName: 'push.rocks',
      publicKey: testCertificates.publicKey,
      privateKey: testCertificates.privateKey,
    },
  ]);

  console.log('[TEST] Proxy configuration updated');
});

tap.test('should route HTTPS requests based on host header', async () => {
  // IMPORTANT: Connect to localhost (where the proxy is listening) but use the Host header "push.rocks"
  const response = await makeHttpsRequest({
    hostname: 'localhost', // changed from 'push.rocks' to 'localhost'
    port: 3001,
    path: '/',
    method: 'GET',
    headers: {
      host: 'push.rocks', // virtual host for routing
    },
    rejectUnauthorized: false,
  });

  expect(response.statusCode).toEqual(200);
  expect(response.body).toEqual('Hello from test server!');
});

tap.test('should handle unknown host headers', async () => {
  // Connect to localhost but use an unknown host header.
  const response = await makeHttpsRequest({
    hostname: 'localhost', // connecting to localhost
    port: 3001,
    path: '/',
    method: 'GET',
    headers: {
      host: 'unknown.host', // this should not match any proxy config
    },
    rejectUnauthorized: false,
  });

  // Expect a 404 response with the appropriate error message.
  expect(response.statusCode).toEqual(404);
  expect(response.body).toEqual('This route is not available on this server.');
});

tap.test('should support WebSocket connections', async () => {
  console.log('\n[TEST] ====== WebSocket Test Started ======');
  console.log('[TEST] Test server port:', 3000);
  console.log('[TEST] Proxy server port:', 3001);
  console.log('\n[TEST] Starting WebSocket test');

  // Reconfigure proxy with test certificates if necessary
  await testProxy.updateProxyConfigs([
    {
      destinationIp: '127.0.0.1',
      destinationPort: '3000',
      hostName: 'push.rocks',
      publicKey: testCertificates.publicKey,
      privateKey: testCertificates.privateKey,
    },
  ]);

  return new Promise<void>((resolve, reject) => {
    console.log('[TEST] Creating WebSocket client');

    // IMPORTANT: Connect to localhost but specify the SNI servername and Host header as "push.rocks"
    const wsUrl = 'wss://localhost:3001'; // changed from 'wss://push.rocks:3001'
    console.log('[TEST] Creating WebSocket connection to:', wsUrl);

    const ws = new WebSocket(wsUrl, {
      rejectUnauthorized: false, // Accept self-signed certificates
      handshakeTimeout: 5000,
      perMessageDeflate: false,
      headers: {
        Host: 'push.rocks', // required for SNI and routing on the proxy
        Connection: 'Upgrade',
        Upgrade: 'websocket',
        'Sec-WebSocket-Version': '13',
      },
      protocol: 'echo-protocol',
      agent: new https.Agent({
        rejectUnauthorized: false, // Also needed for the underlying HTTPS connection
      }),
    });

    console.log('[TEST] WebSocket client created');

    let resolved = false;
    const cleanup = () => {
      if (!resolved) {
        resolved = true;
        try {
          console.log('[TEST] Cleaning up WebSocket connection');
          ws.close();
          resolve();
        } catch (error) {
          console.error('[TEST] Error during cleanup:', error);
          reject(error);
        }
      }
    };

    const timeout = setTimeout(() => {
      console.error('[TEST] WebSocket test timed out');
      cleanup();
      reject(new Error('WebSocket test timed out after 5 seconds'));
    }, 5000);

    // Connection establishment events
    ws.on('upgrade', (response) => {
      console.log('[TEST] WebSocket upgrade response received:', {
        headers: response.headers,
        statusCode: response.statusCode,
      });
    });

    ws.on('open', () => {
      console.log('[TEST] WebSocket connection opened');
      try {
        console.log('[TEST] Sending test message');
        ws.send('Hello WebSocket');
      } catch (error) {
        console.error('[TEST] Error sending message:', error);
        cleanup();
        reject(error);
      }
    });

    ws.on('message', (message) => {
      console.log('[TEST] Received message:', message.toString());
      if (
        message.toString() === 'Hello WebSocket' ||
        message.toString() === 'Echo: Hello WebSocket'
      ) {
        console.log('[TEST] Message received correctly');
        clearTimeout(timeout);
        cleanup();
      }
    });

    ws.on('error', (error) => {
      console.error('[TEST] WebSocket error:', error);
      cleanup();
      reject(error);
    });

    ws.on('close', (code, reason) => {
      console.log('[TEST] WebSocket connection closed:', {
        code,
        reason: reason.toString(),
      });
      cleanup();
    });
  });
});

tap.test('should handle custom headers', async () => {
  await testProxy.addDefaultHeaders({
    'X-Proxy-Header': 'test-value',
  });

  const response = await makeHttpsRequest({
    hostname: 'localhost', // changed to 'localhost'
    port: 3001,
    path: '/',
    method: 'GET',
    headers: {
      host: 'push.rocks', // still routing to push.rocks
    },
    rejectUnauthorized: false,
  });

  expect(response.headers['x-proxy-header']).toEqual('test-value');
});

tap.test('cleanup', async () => {
  console.log('[TEST] Starting cleanup');

  // Clean up all servers
  console.log('[TEST] Terminating WebSocket clients');
  wsServer.clients.forEach((client) => {
    client.terminate();
  });

  console.log('[TEST] Closing WebSocket server');
  await new Promise<void>((resolve) =>
    wsServer.close(() => {
      console.log('[TEST] WebSocket server closed');
      resolve();
    })
  );

  console.log('[TEST] Closing test server');
  await new Promise<void>((resolve) =>
    testServer.close(() => {
      console.log('[TEST] Test server closed');
      resolve();
    })
  );

  console.log('[TEST] Stopping proxy');
  await testProxy.stop();
  console.log('[TEST] Cleanup complete');
});

process.on('exit', () => {
  console.log('[TEST] Shutting down test server');
  testServer.close(() => console.log('[TEST] Test server shut down'));
  wsServer.close(() => console.log('[TEST] WebSocket server shut down'));
  testProxy.stop().then(() => console.log('[TEST] Proxy server stopped'));
});

tap.start();