From eb2e67fecc5ce71843f8448b5db20cfb0ee04b50 Mon Sep 17 00:00:00 2001 From: Philipp Kunz Date: Sun, 1 Jun 2025 07:51:20 +0000 Subject: [PATCH] feat(socket-utils): implement socket cleanup utilities and enhance socket handling in forwarding handlers --- test/test.httpproxy.ts | 9 +- ts/core/utils/index.ts | 1 + ts/core/utils/socket-utils.ts | 96 +++++++++++++++ ts/forwarding/handlers/http-handler.ts | 10 +- .../handlers/https-passthrough-handler.ts | 79 +++--------- .../https-terminate-to-http-handler.ts | 116 ++++++++++-------- .../https-terminate-to-https-handler.ts | 91 ++++++-------- ts/proxies/smart-proxy/http-proxy-bridge.ts | 16 ++- .../smart-proxy/route-connection-handler.ts | 36 ++++++ 9 files changed, 273 insertions(+), 181 deletions(-) create mode 100644 ts/core/utils/socket-utils.ts diff --git a/test/test.httpproxy.ts b/test/test.httpproxy.ts index 5765963..e066a69 100644 --- a/test/test.httpproxy.ts +++ b/test/test.httpproxy.ts @@ -591,13 +591,6 @@ tap.test('cleanup', async () => { // Exit handler removed to prevent interference with test cleanup -// Add a post-hook to force exit after tap completion -tap.test('teardown', async () => { - // Force exit after all tests complete - setTimeout(() => { - console.log('[TEST] Force exit after tap completion'); - process.exit(0); - }, 1000); -}); +// Teardown test removed - let tap handle proper cleanup export default tap.start(); \ No newline at end of file diff --git a/ts/core/utils/index.ts b/ts/core/utils/index.ts index a79f145..1992994 100644 --- a/ts/core/utils/index.ts +++ b/ts/core/utils/index.ts @@ -16,3 +16,4 @@ export * from './fs-utils.js'; export * from './lifecycle-component.js'; export * from './binary-heap.js'; export * from './enhanced-connection-pool.js'; +export * from './socket-utils.js'; diff --git a/ts/core/utils/socket-utils.ts b/ts/core/utils/socket-utils.ts new file mode 100644 index 0000000..fc41b8a --- /dev/null +++ b/ts/core/utils/socket-utils.ts @@ -0,0 +1,96 @@ +import * as plugins from '../../plugins.js'; + +/** + * Safely cleanup a socket by removing all listeners and destroying it + * @param socket The socket to cleanup + * @param socketName Optional name for logging + */ +export function cleanupSocket(socket: plugins.net.Socket | plugins.tls.TLSSocket | null, socketName?: string): void { + if (!socket) return; + + try { + // Remove all event listeners + socket.removeAllListeners(); + + // Unpipe any streams + socket.unpipe(); + + // Destroy if not already destroyed + if (!socket.destroyed) { + socket.destroy(); + } + } catch (err) { + console.error(`Error cleaning up socket${socketName ? ` (${socketName})` : ''}: ${err}`); + } +} + +/** + * Create a cleanup handler for paired sockets (client and server) + * @param clientSocket The client socket + * @param serverSocket The server socket (optional) + * @param onCleanup Optional callback when cleanup is done + * @returns A cleanup function that can be called multiple times safely + */ +export function createSocketCleanupHandler( + clientSocket: plugins.net.Socket | plugins.tls.TLSSocket, + serverSocket?: plugins.net.Socket | plugins.tls.TLSSocket | null, + onCleanup?: (reason: string) => void +): (reason: string) => void { + let cleanedUp = false; + + return (reason: string) => { + if (cleanedUp) return; + cleanedUp = true; + + // Cleanup both sockets + cleanupSocket(clientSocket, 'client'); + if (serverSocket) { + cleanupSocket(serverSocket, 'server'); + } + + // Call cleanup callback if provided + if (onCleanup) { + onCleanup(reason); + } + }; +} + +/** + * Setup socket error and close handlers with proper cleanup + * @param socket The socket to setup handlers for + * @param handleClose The cleanup function to call + * @param errorPrefix Optional prefix for error messages + */ +export function setupSocketHandlers( + socket: plugins.net.Socket | plugins.tls.TLSSocket, + handleClose: (reason: string) => void, + errorPrefix?: string +): void { + socket.on('error', (error) => { + const prefix = errorPrefix || 'Socket'; + handleClose(`${prefix}_error: ${error.message}`); + }); + + socket.on('close', () => { + const prefix = errorPrefix || 'socket'; + handleClose(`${prefix}_closed`); + }); + + socket.on('timeout', () => { + const prefix = errorPrefix || 'socket'; + handleClose(`${prefix}_timeout`); + }); +} + +/** + * Pipe two sockets together with proper cleanup on either end + * @param socket1 First socket + * @param socket2 Second socket + */ +export function pipeSockets( + socket1: plugins.net.Socket | plugins.tls.TLSSocket, + socket2: plugins.net.Socket | plugins.tls.TLSSocket +): void { + socket1.pipe(socket2); + socket2.pipe(socket1); +} \ No newline at end of file diff --git a/ts/forwarding/handlers/http-handler.ts b/ts/forwarding/handlers/http-handler.ts index 258ad53..c2cddf3 100644 --- a/ts/forwarding/handlers/http-handler.ts +++ b/ts/forwarding/handlers/http-handler.ts @@ -2,6 +2,7 @@ import * as plugins from '../../plugins.js'; import { ForwardingHandler } from './base-handler.js'; import type { IForwardConfig } from '../config/forwarding-types.js'; import { ForwardingHandlerEvents } from '../config/forwarding-types.js'; +import { setupSocketHandlers } from '../../core/utils/socket-utils.js'; /** * Handler for HTTP-only forwarding @@ -40,12 +41,15 @@ export class HttpForwardingHandler extends ForwardingHandler { const remoteAddress = socket.remoteAddress || 'unknown'; const localPort = socket.localPort || 80; - socket.on('close', (hadError) => { + // Set up socket handlers with proper cleanup + const handleClose = (reason: string) => { this.emit(ForwardingHandlerEvents.DISCONNECTED, { remoteAddress, - hadError + reason }); - }); + }; + + setupSocketHandlers(socket, handleClose, 'http'); socket.on('error', (error) => { this.emit(ForwardingHandlerEvents.ERROR, { diff --git a/ts/forwarding/handlers/https-passthrough-handler.ts b/ts/forwarding/handlers/https-passthrough-handler.ts index c4ba40e..a804630 100644 --- a/ts/forwarding/handlers/https-passthrough-handler.ts +++ b/ts/forwarding/handlers/https-passthrough-handler.ts @@ -2,6 +2,7 @@ import * as plugins from '../../plugins.js'; import { ForwardingHandler } from './base-handler.js'; import type { IForwardConfig } from '../config/forwarding-types.js'; import { ForwardingHandlerEvents } from '../config/forwarding-types.js'; +import { createSocketCleanupHandler, setupSocketHandlers, pipeSockets } from '../../core/utils/socket-utils.js'; /** * Handler for HTTPS passthrough (SNI forwarding without termination) @@ -50,36 +51,24 @@ export class HttpsPassthroughHandler extends ForwardingHandler { // Create a connection to the target server const serverSocket = plugins.net.connect(target.port, target.host); - // Handle errors on the server socket - serverSocket.on('error', (error) => { - this.emit(ForwardingHandlerEvents.ERROR, { - remoteAddress, - error: `Target connection error: ${error.message}` - }); - - // Close the client socket if it's still open - if (!clientSocket.destroyed) { - clientSocket.destroy(); - } - }); - - // Handle errors on the client socket - clientSocket.on('error', (error) => { - this.emit(ForwardingHandlerEvents.ERROR, { - remoteAddress, - error: `Client connection error: ${error.message}` - }); - - // Close the server socket if it's still open - if (!serverSocket.destroyed) { - serverSocket.destroy(); - } - }); - // Track data transfer for logging let bytesSent = 0; let bytesReceived = 0; + // Create cleanup handler with our utility + const handleClose = createSocketCleanupHandler(clientSocket, serverSocket, (reason) => { + this.emit(ForwardingHandlerEvents.DISCONNECTED, { + remoteAddress, + bytesSent, + bytesReceived, + reason + }); + }); + + // Setup error and close handlers for both sockets + setupSocketHandlers(serverSocket, handleClose, 'server'); + setupSocketHandlers(clientSocket, handleClose, 'client'); + // Forward data from client to server clientSocket.on('data', (data) => { bytesSent += data.length; @@ -128,48 +117,10 @@ export class HttpsPassthroughHandler extends ForwardingHandler { }); }); - // Handle connection close - const handleClose = () => { - if (!clientSocket.destroyed) { - clientSocket.destroy(); - } - - if (!serverSocket.destroyed) { - serverSocket.destroy(); - } - - this.emit(ForwardingHandlerEvents.DISCONNECTED, { - remoteAddress, - bytesSent, - bytesReceived - }); - }; - - // Set up close handlers - clientSocket.on('close', handleClose); - serverSocket.on('close', handleClose); - // Set timeouts const timeout = this.getTimeout(); clientSocket.setTimeout(timeout); serverSocket.setTimeout(timeout); - - // Handle timeouts - clientSocket.on('timeout', () => { - this.emit(ForwardingHandlerEvents.ERROR, { - remoteAddress, - error: 'Client connection timeout' - }); - handleClose(); - }); - - serverSocket.on('timeout', () => { - this.emit(ForwardingHandlerEvents.ERROR, { - remoteAddress, - error: 'Server connection timeout' - }); - handleClose(); - }); } /** diff --git a/ts/forwarding/handlers/https-terminate-to-http-handler.ts b/ts/forwarding/handlers/https-terminate-to-http-handler.ts index 4f68e61..f8a93d8 100644 --- a/ts/forwarding/handlers/https-terminate-to-http-handler.ts +++ b/ts/forwarding/handlers/https-terminate-to-http-handler.ts @@ -2,6 +2,7 @@ import * as plugins from '../../plugins.js'; import { ForwardingHandler } from './base-handler.js'; import type { IForwardConfig } from '../config/forwarding-types.js'; import { ForwardingHandlerEvents } from '../config/forwarding-types.js'; +import { createSocketCleanupHandler, setupSocketHandlers } from '../../core/utils/socket-utils.js'; /** * Handler for HTTPS termination with HTTP backend @@ -95,62 +96,24 @@ export class HttpsTerminateToHttpHandler extends ForwardingHandler { tls: true }); - // Handle TLS errors - tlsSocket.on('error', (error) => { - this.emit(ForwardingHandlerEvents.ERROR, { - remoteAddress, - error: `TLS error: ${error.message}` - }); - - if (!tlsSocket.destroyed) { - tlsSocket.destroy(); - } - }); - - // The TLS socket will now emit HTTP traffic that can be processed - // In a real implementation, we would create an HTTP parser and handle - // the requests here, but for simplicity, we'll just log the data - + // Variables to track connections + let backendSocket: plugins.net.Socket | null = null; let dataBuffer = Buffer.alloc(0); + let connectionEstablished = false; - tlsSocket.on('data', (data) => { - // Append to buffer - dataBuffer = Buffer.concat([dataBuffer, data]); - - // Very basic HTTP parsing - in a real implementation, use http-parser - if (dataBuffer.includes(Buffer.from('\r\n\r\n'))) { - const target = this.getTargetFromConfig(); - - // Simple example: forward the data to an HTTP server - const socket = plugins.net.connect(target.port, target.host, () => { - socket.write(dataBuffer); - dataBuffer = Buffer.alloc(0); - - // Set up bidirectional data flow - tlsSocket.pipe(socket); - socket.pipe(tlsSocket); - }); - - socket.on('error', (error) => { - this.emit(ForwardingHandlerEvents.ERROR, { - remoteAddress, - error: `Target connection error: ${error.message}` - }); - - if (!tlsSocket.destroyed) { - tlsSocket.destroy(); - } - }); - } - }); - - // Handle close - tlsSocket.on('close', () => { + // Create cleanup handler for all sockets + const handleClose = createSocketCleanupHandler(tlsSocket, backendSocket, (reason) => { this.emit(ForwardingHandlerEvents.DISCONNECTED, { - remoteAddress + remoteAddress, + reason }); + dataBuffer = Buffer.alloc(0); + connectionEstablished = false; }); + // Set up error handling with our cleanup utility + setupSocketHandlers(tlsSocket, handleClose, 'tls'); + // Set timeout const timeout = this.getTimeout(); tlsSocket.setTimeout(timeout); @@ -160,9 +123,58 @@ export class HttpsTerminateToHttpHandler extends ForwardingHandler { remoteAddress, error: 'TLS connection timeout' }); + handleClose('timeout'); + }); + + // Handle TLS data + tlsSocket.on('data', (data) => { + // If backend connection already established, just forward the data + if (connectionEstablished && backendSocket && !backendSocket.destroyed) { + backendSocket.write(data); + return; + } - if (!tlsSocket.destroyed) { - tlsSocket.destroy(); + // Append to buffer + dataBuffer = Buffer.concat([dataBuffer, data]); + + // Very basic HTTP parsing - in a real implementation, use http-parser + if (dataBuffer.includes(Buffer.from('\r\n\r\n')) && !connectionEstablished) { + const target = this.getTargetFromConfig(); + + // Create backend connection + backendSocket = plugins.net.connect(target.port, target.host, () => { + connectionEstablished = true; + + // Send buffered data + if (dataBuffer.length > 0) { + backendSocket!.write(dataBuffer); + dataBuffer = Buffer.alloc(0); + } + + // Set up bidirectional data flow + tlsSocket.pipe(backendSocket!); + backendSocket!.pipe(tlsSocket); + }); + + // Update the cleanup handler with the backend socket + const newHandleClose = createSocketCleanupHandler(tlsSocket, backendSocket, (reason) => { + this.emit(ForwardingHandlerEvents.DISCONNECTED, { + remoteAddress, + reason + }); + dataBuffer = Buffer.alloc(0); + connectionEstablished = false; + }); + + // Set up handlers for backend socket + setupSocketHandlers(backendSocket, newHandleClose, 'backend'); + + backendSocket.on('error', (error) => { + this.emit(ForwardingHandlerEvents.ERROR, { + remoteAddress, + error: `Target connection error: ${error.message}` + }); + }); } }); } diff --git a/ts/forwarding/handlers/https-terminate-to-https-handler.ts b/ts/forwarding/handlers/https-terminate-to-https-handler.ts index 1505bfa..d2970fd 100644 --- a/ts/forwarding/handlers/https-terminate-to-https-handler.ts +++ b/ts/forwarding/handlers/https-terminate-to-https-handler.ts @@ -2,6 +2,7 @@ import * as plugins from '../../plugins.js'; import { ForwardingHandler } from './base-handler.js'; import type { IForwardConfig } from '../config/forwarding-types.js'; import { ForwardingHandlerEvents } from '../config/forwarding-types.js'; +import { createSocketCleanupHandler, setupSocketHandlers } from '../../core/utils/socket-utils.js'; /** * Handler for HTTPS termination with HTTPS backend @@ -93,28 +94,38 @@ export class HttpsTerminateToHttpsHandler extends ForwardingHandler { tls: true }); - // Handle TLS errors - tlsSocket.on('error', (error) => { - this.emit(ForwardingHandlerEvents.ERROR, { + // Variable to track backend socket + let backendSocket: plugins.tls.TLSSocket | null = null; + + // Create cleanup handler for both sockets + const handleClose = createSocketCleanupHandler(tlsSocket, backendSocket, (reason) => { + this.emit(ForwardingHandlerEvents.DISCONNECTED, { remoteAddress, - error: `TLS error: ${error.message}` + reason }); - - if (!tlsSocket.destroyed) { - tlsSocket.destroy(); - } }); - // The TLS socket will now emit HTTP traffic that can be processed - // In a real implementation, we would create an HTTP parser and handle - // the requests here, but for simplicity, we'll just forward the data + // Set up error handling with our cleanup utility + setupSocketHandlers(tlsSocket, handleClose, 'tls'); + + // Set timeout + const timeout = this.getTimeout(); + tlsSocket.setTimeout(timeout); + + tlsSocket.on('timeout', () => { + this.emit(ForwardingHandlerEvents.ERROR, { + remoteAddress, + error: 'TLS connection timeout' + }); + handleClose('timeout'); + }); // Get the target from configuration const target = this.getTargetFromConfig(); // Set up the connection to the HTTPS backend const connectToBackend = () => { - const backendSocket = plugins.tls.connect({ + backendSocket = plugins.tls.connect({ host: target.host, port: target.port, // In a real implementation, we would configure TLS options @@ -127,30 +138,29 @@ export class HttpsTerminateToHttpsHandler extends ForwardingHandler { }); // Set up bidirectional data flow - tlsSocket.pipe(backendSocket); - backendSocket.pipe(tlsSocket); + tlsSocket.pipe(backendSocket!); + backendSocket!.pipe(tlsSocket); }); + // Update the cleanup handler with the backend socket + const newHandleClose = createSocketCleanupHandler(tlsSocket, backendSocket, (reason) => { + this.emit(ForwardingHandlerEvents.DISCONNECTED, { + remoteAddress, + reason + }); + }); + + // Set up handlers for backend socket + setupSocketHandlers(backendSocket, newHandleClose, 'backend'); + backendSocket.on('error', (error) => { this.emit(ForwardingHandlerEvents.ERROR, { remoteAddress, error: `Backend connection error: ${error.message}` }); - - if (!tlsSocket.destroyed) { - tlsSocket.destroy(); - } }); - // Handle close - backendSocket.on('close', () => { - if (!tlsSocket.destroyed) { - tlsSocket.destroy(); - } - }); - - // Set timeout - const timeout = this.getTimeout(); + // Set timeout for backend socket backendSocket.setTimeout(timeout); backendSocket.on('timeout', () => { @@ -158,10 +168,7 @@ export class HttpsTerminateToHttpsHandler extends ForwardingHandler { remoteAddress, error: 'Backend connection timeout' }); - - if (!backendSocket.destroyed) { - backendSocket.destroy(); - } + newHandleClose('backend_timeout'); }); }; @@ -169,28 +176,6 @@ export class HttpsTerminateToHttpsHandler extends ForwardingHandler { tlsSocket.on('secure', () => { connectToBackend(); }); - - // Handle close - tlsSocket.on('close', () => { - this.emit(ForwardingHandlerEvents.DISCONNECTED, { - remoteAddress - }); - }); - - // Set timeout - const timeout = this.getTimeout(); - tlsSocket.setTimeout(timeout); - - tlsSocket.on('timeout', () => { - this.emit(ForwardingHandlerEvents.ERROR, { - remoteAddress, - error: 'TLS connection timeout' - }); - - if (!tlsSocket.destroyed) { - tlsSocket.destroy(); - } - }); } /** diff --git a/ts/proxies/smart-proxy/http-proxy-bridge.ts b/ts/proxies/smart-proxy/http-proxy-bridge.ts index 40292fa..4306bb7 100644 --- a/ts/proxies/smart-proxy/http-proxy-bridge.ts +++ b/ts/proxies/smart-proxy/http-proxy-bridge.ts @@ -128,10 +128,24 @@ export class HttpProxyBridge { proxySocket.pipe(socket); // Handle cleanup + let cleanedUp = false; const cleanup = (reason: string) => { + if (cleanedUp) return; + cleanedUp = true; + + // Remove all event listeners to prevent memory leaks + socket.removeAllListeners('end'); + socket.removeAllListeners('error'); + proxySocket.removeAllListeners('end'); + proxySocket.removeAllListeners('error'); + socket.unpipe(proxySocket); proxySocket.unpipe(socket); - proxySocket.destroy(); + + if (!proxySocket.destroyed) { + proxySocket.destroy(); + } + cleanupCallback(reason); }; diff --git a/ts/proxies/smart-proxy/route-connection-handler.ts b/ts/proxies/smart-proxy/route-connection-handler.ts index 5a3962b..ed348f7 100644 --- a/ts/proxies/smart-proxy/route-connection-handler.ts +++ b/ts/proxies/smart-proxy/route-connection-handler.ts @@ -822,6 +822,38 @@ export class RouteConnectionHandler { return; } + // Track event listeners added by the handler so we can clean them up + const originalOn = socket.on.bind(socket); + const originalOnce = socket.once.bind(socket); + const trackedListeners: Array<{event: string; listener: Function}> = []; + + // Override socket.on to track listeners + socket.on = function(event: string, listener: Function) { + trackedListeners.push({event, listener}); + return originalOn(event, listener); + } as any; + + // Override socket.once to track listeners + socket.once = function(event: string, listener: Function) { + trackedListeners.push({event, listener}); + return originalOnce(event, listener); + } as any; + + // Set up automatic cleanup when socket closes + const cleanupHandler = () => { + // Remove all tracked listeners + for (const {event, listener} of trackedListeners) { + socket.removeListener(event, listener); + } + // Restore original methods + socket.on = originalOn; + socket.once = originalOnce; + }; + + // Listen for socket close to trigger cleanup + originalOnce('close', cleanupHandler); + originalOnce('error', cleanupHandler); + // Create route context for the handler const routeContext = this.createRouteContext({ connectionId: record.id, @@ -855,6 +887,8 @@ export class RouteConnectionHandler { error: error.message, component: 'route-handler' }); + // Remove all event listeners before destroying to prevent memory leaks + socket.removeAllListeners(); if (!socket.destroyed) { socket.destroy(); } @@ -875,6 +909,8 @@ export class RouteConnectionHandler { error: error.message, component: 'route-handler' }); + // Remove all event listeners before destroying to prevent memory leaks + socket.removeAllListeners(); if (!socket.destroyed) { socket.destroy(); }