smartproxy/test/test.wrapped-socket.ts

import { expect, tap } from '@git.zone/tstest/tapbundle';
import * as plugins from '../ts/plugins.js';
import { WrappedSocket } from '../ts/core/models/wrapped-socket.js';
import * as net from 'net';

tap.test('WrappedSocket - should wrap a regular socket', async () => {
  // Create a simple test server
  const server = net.createServer();
  await new Promise<void>((resolve) => {
    server.listen(0, 'localhost', () => resolve());
  });
  
  const serverPort = (server.address() as net.AddressInfo).port;
  
  // Create a client connection
  const clientSocket = net.connect(serverPort, 'localhost');
  
  // Wrap the socket
  const wrappedSocket = new WrappedSocket(clientSocket);
  
  // Test initial state - should use underlying socket values
  expect(wrappedSocket.remoteAddress).toEqual(clientSocket.remoteAddress);
  expect(wrappedSocket.remotePort).toEqual(clientSocket.remotePort);
  expect(wrappedSocket.localAddress).toEqual(clientSocket.localAddress);
  expect(wrappedSocket.localPort).toEqual(clientSocket.localPort);
  expect(wrappedSocket.isFromTrustedProxy).toBeFalse();
  
  // Clean up
  clientSocket.destroy();
  server.close();
});

tap.test('WrappedSocket - should provide real client info when set', async () => {
  // Create a simple test server
  const server = net.createServer();
  await new Promise<void>((resolve) => {
    server.listen(0, 'localhost', () => resolve());
  });
  
  const serverPort = (server.address() as net.AddressInfo).port;
  
  // Create a client connection
  const clientSocket = net.connect(serverPort, 'localhost');
  
  // Wrap the socket with initial proxy info
  const wrappedSocket = new WrappedSocket(clientSocket, '192.168.1.100', 54321);
  
  // Test that real client info is returned
  expect(wrappedSocket.remoteAddress).toEqual('192.168.1.100');
  expect(wrappedSocket.remotePort).toEqual(54321);
  expect(wrappedSocket.isFromTrustedProxy).toBeTrue();
  
  // Local info should still come from underlying socket
  expect(wrappedSocket.localAddress).toEqual(clientSocket.localAddress);
  expect(wrappedSocket.localPort).toEqual(clientSocket.localPort);
  
  // Clean up
  clientSocket.destroy();
  server.close();
});

tap.test('WrappedSocket - should update proxy info via setProxyInfo', async () => {
  // Create a simple test server
  const server = net.createServer();
  await new Promise<void>((resolve) => {
    server.listen(0, 'localhost', () => resolve());
  });
  
  const serverPort = (server.address() as net.AddressInfo).port;
  
  // Create a client connection
  const clientSocket = net.connect(serverPort, 'localhost');
  
  // Wrap the socket without initial proxy info
  const wrappedSocket = new WrappedSocket(clientSocket);
  
  // Initially should use underlying socket
  expect(wrappedSocket.isFromTrustedProxy).toBeFalse();
  expect(wrappedSocket.remoteAddress).toEqual(clientSocket.remoteAddress);
  
  // Update proxy info
  wrappedSocket.setProxyInfo('10.0.0.5', 12345);
  
  // Now should return proxy info
  expect(wrappedSocket.remoteAddress).toEqual('10.0.0.5');
  expect(wrappedSocket.remotePort).toEqual(12345);
  expect(wrappedSocket.isFromTrustedProxy).toBeTrue();
  
  // Clean up
  clientSocket.destroy();
  server.close();
});

tap.test('WrappedSocket - should correctly determine IP family', async () => {
  // Create a simple test server
  const server = net.createServer();
  await new Promise<void>((resolve) => {
    server.listen(0, 'localhost', () => resolve());
  });
  
  const serverPort = (server.address() as net.AddressInfo).port;
  
  // Create a client connection
  const clientSocket = net.connect(serverPort, 'localhost');
  
  // Test IPv4
  const wrappedSocketIPv4 = new WrappedSocket(clientSocket, '192.168.1.1', 80);
  expect(wrappedSocketIPv4.remoteFamily).toEqual('IPv4');
  
  // Test IPv6
  const wrappedSocketIPv6 = new WrappedSocket(clientSocket, '2001:0db8:85a3:0000:0000:8a2e:0370:7334', 443);
  expect(wrappedSocketIPv6.remoteFamily).toEqual('IPv6');
  
  // Test fallback to underlying socket
  const wrappedSocketNoProxy = new WrappedSocket(clientSocket);
  expect(wrappedSocketNoProxy.remoteFamily).toEqual(clientSocket.remoteFamily);
  
  // Clean up
  clientSocket.destroy();
  server.close();
});

tap.test('WrappedSocket - should forward events correctly', async () => {
  // Create a simple echo server
  let serverConnection: net.Socket;
  const server = net.createServer((socket) => {
    serverConnection = socket;
    socket.on('data', (data) => {
      socket.write(data); // Echo back
    });
  });
  
  await new Promise<void>((resolve) => {
    server.listen(0, 'localhost', () => resolve());
  });
  
  const serverPort = (server.address() as net.AddressInfo).port;
  
  // Create a client connection
  const clientSocket = net.connect(serverPort, 'localhost');
  
  // Wrap the socket
  const wrappedSocket = new WrappedSocket(clientSocket);
  
  // Set up event tracking
  let connectReceived = false;
  let dataReceived = false;
  let endReceived = false;
  let closeReceived = false;
  
  wrappedSocket.on('connect', () => {
    connectReceived = true;
  });
  
  wrappedSocket.on('data', (chunk) => {
    dataReceived = true;
    expect(chunk.toString()).toEqual('test data');
  });
  
  wrappedSocket.on('end', () => {
    endReceived = true;
  });
  
  wrappedSocket.on('close', () => {
    closeReceived = true;
  });
  
  // Wait for connection
  await new Promise<void>((resolve) => {
    if (clientSocket.readyState === 'open') {
      resolve();
    } else {
      clientSocket.once('connect', () => resolve());
    }
  });
  
  // Send data
  wrappedSocket.write('test data');
  
  // Wait for echo
  await new Promise(resolve => setTimeout(resolve, 100));
  
  // Close the connection
  serverConnection.end();
  
  // Wait for events
  await new Promise(resolve => setTimeout(resolve, 100));
  
  // Verify all events were received
  expect(dataReceived).toBeTrue();
  expect(endReceived).toBeTrue();
  expect(closeReceived).toBeTrue();
  
  // Clean up
  server.close();
});

tap.test('WrappedSocket - should pass through socket methods', async () => {
  // Create a simple test server
  const server = net.createServer();
  await new Promise<void>((resolve) => {
    server.listen(0, 'localhost', () => resolve());
  });
  
  const serverPort = (server.address() as net.AddressInfo).port;
  
  // Create a client connection
  const clientSocket = net.connect(serverPort, 'localhost');
  await new Promise<void>((resolve) => {
    clientSocket.once('connect', () => resolve());
  });
  
  // Wrap the socket
  const wrappedSocket = new WrappedSocket(clientSocket);
  
  // Test various pass-through methods
  expect(wrappedSocket.readable).toEqual(clientSocket.readable);
  expect(wrappedSocket.writable).toEqual(clientSocket.writable);
  expect(wrappedSocket.destroyed).toEqual(clientSocket.destroyed);
  expect(wrappedSocket.bytesRead).toEqual(clientSocket.bytesRead);
  expect(wrappedSocket.bytesWritten).toEqual(clientSocket.bytesWritten);
  
  // Test method calls
  wrappedSocket.pause();
  expect(clientSocket.isPaused()).toBeTrue();
  
  wrappedSocket.resume();
  expect(clientSocket.isPaused()).toBeFalse();
  
  // Test setTimeout
  let timeoutCalled = false;
  wrappedSocket.setTimeout(100, () => {
    timeoutCalled = true;
  });
  await new Promise(resolve => setTimeout(resolve, 150));
  expect(timeoutCalled).toBeTrue();
  
  // Clean up
  wrappedSocket.destroy();
  server.close();
});

tap.test('WrappedSocket - should handle write and pipe operations', async () => {
  // Create a simple echo server
  const server = net.createServer((socket) => {
    socket.pipe(socket); // Echo everything back
  });
  
  await new Promise<void>((resolve) => {
    server.listen(0, 'localhost', () => resolve());
  });
  
  const serverPort = (server.address() as net.AddressInfo).port;
  
  // Create a client connection
  const clientSocket = net.connect(serverPort, 'localhost');
  await new Promise<void>((resolve) => {
    clientSocket.once('connect', () => resolve());
  });
  
  // Wrap the socket
  const wrappedSocket = new WrappedSocket(clientSocket);
  
  // Test write with callback
  const writeResult = wrappedSocket.write('test', 'utf8', () => {
    // Write completed
  });
  expect(typeof writeResult).toEqual('boolean');
  
  // Test pipe
  const { PassThrough } = await import('stream');
  const passThrough = new PassThrough();
  const piped = wrappedSocket.pipe(passThrough);
  expect(piped).toEqual(passThrough);
  
  // Clean up
  wrappedSocket.destroy();
  server.close();
});

tap.test('WrappedSocket - should handle encoding and address methods', async () => {
  // Create a simple test server
  const server = net.createServer();
  await new Promise<void>((resolve) => {
    server.listen(0, 'localhost', () => resolve());
  });
  
  const serverPort = (server.address() as net.AddressInfo).port;
  
  // Create a client connection
  const clientSocket = net.connect(serverPort, 'localhost');
  await new Promise<void>((resolve) => {
    clientSocket.once('connect', () => resolve());
  });
  
  // Wrap the socket
  const wrappedSocket = new WrappedSocket(clientSocket);
  
  // Test setEncoding
  wrappedSocket.setEncoding('utf8');
  
  // Test address method
  const addr = wrappedSocket.address();
  expect(addr).toEqual(clientSocket.address());
  
  // Test cork/uncork (if available)
  wrappedSocket.cork();
  wrappedSocket.uncork();
  
  // Clean up
  wrappedSocket.destroy();
  server.close();
});

tap.test('WrappedSocket - should work with ConnectionManager', async () => {
  // This test verifies that WrappedSocket can be used seamlessly with ConnectionManager
  const { ConnectionManager } = await import('../ts/proxies/smart-proxy/connection-manager.js');
  const { SecurityManager } = await import('../ts/proxies/smart-proxy/security-manager.js');
  const { TimeoutManager } = await import('../ts/proxies/smart-proxy/timeout-manager.js');
  
  // Create minimal settings
  const settings = {
    routes: [],
    defaults: {
      security: {
        maxConnections: 100
      }
    }
  };
  
  const securityManager = new SecurityManager(settings);
  const timeoutManager = new TimeoutManager(settings);
  const connectionManager = new ConnectionManager(settings, securityManager, timeoutManager);
  
  // Create a simple test server
  const server = net.createServer();
  await new Promise<void>((resolve) => {
    server.listen(0, 'localhost', () => resolve());
  });
  
  const serverPort = (server.address() as net.AddressInfo).port;
  
  // Create a client connection
  const clientSocket = net.connect(serverPort, 'localhost');
  
  // Wait for connection to establish
  await new Promise<void>((resolve) => {
    clientSocket.once('connect', () => resolve());
  });
  
  // Wrap with proxy info
  const wrappedSocket = new WrappedSocket(clientSocket, '203.0.113.45', 65432);
  
  // Create connection using wrapped socket
  const record = connectionManager.createConnection(wrappedSocket);
  
  expect(record).toBeTruthy();
  expect(record!.remoteIP).toEqual('203.0.113.45'); // Should use the real client IP
  expect(record!.localPort).toEqual(clientSocket.localPort);
  
  // Clean up
  connectionManager.cleanupConnection(record!, 'test-complete');
  server.close();
});

export default tap.start();