feat(proxy): Implement WrappedSocket class for PROXY protocol support and update connection handling

2025-06-05 17:57:24 +00:00
parent 2a75e7c490
commit 18d79ac7e1
7 changed files with 813 additions and 51 deletions
--- a/test/test.wrapped-socket.ts
+++ b/test/test.wrapped-socket.ts
@ -0,0 +1,366 @@
+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();