fix(networkproxy): Refactor and improve WebSocket handling and request processing

changelog.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## 2025-02-04 - 3.1.3 - fix(networkproxy)
+Refactor and improve WebSocket handling and request processing
+
+- Improved error handling in WebSocket connection and request processing.
+- Refactored the WebSocket handling in NetworkProxy to use a unified error logging mechanism.
+
 ## 2025-02-04 - 3.1.2 - fix(core)
 Refactor certificate handling across the project
 

test/test.ts

@@ -14,19 +14,34 @@ let testCertificates: { privateKey: string; publicKey: string };
 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', () =>
+      res.on('end', () => {
+        console.log('[TEST] Response completed:', { data });
         resolve({
           statusCode: res.statusCode!,
           headers: res.headers,
           body: data,
-        }),
-      );
         });
-    req.on('error', reject);
+      });
+    });
+    req.on('error', (error) => {
+      console.error('[TEST] Request error:', error);
+      reject(error);
+    });
     req.end();
   });
 }
@@ -37,12 +52,13 @@ tap.test('setup test environment', async () => {
   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
+      headers: req.headers,
     });
     res.writeHead(200, { 'Content-Type': 'text/plain' });
     res.end('Hello from test server!');
@@ -59,8 +75,8 @@ tap.test('setup test environment', async () => {
         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']
+        'sec-websocket-protocol': request.headers['sec-websocket-protocol'],
-      }
+      },
     });
 
     if (request.headers.upgrade?.toLowerCase() !== 'websocket') {
@@ -76,13 +92,13 @@ tap.test('setup test environment', async () => {
     });
   });
 
-  // Create a WebSocket server
+  // 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'
+    handleProtocols: () => 'echo-protocol',
   });
 
   wsServer.on('connection', (ws, request) => {
@@ -94,8 +110,8 @@ tap.test('setup test environment', async () => {
         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']
+        'sec-websocket-protocol': request.headers['sec-websocket-protocol'],
-      }
+      },
     });
 
     // Set up connection timeout
@@ -132,7 +148,7 @@ tap.test('setup test environment', async () => {
       console.log('[TEST SERVER] WebSocket connection closed:', {
         code,
         reason: reason.toString(),
-        wasClean: code === 1000 || code === 1001
+        wasClean: code === 1000 || code === 1001,
       });
       clearConnectionTimeout();
     });
@@ -175,30 +191,42 @@ tap.test('should create proxy instance', async () => {
 });
 
 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
-  testProxy.updateProxyConfigs([
+  // 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: 'push.rocks',
+    hostname: 'localhost', // changed from 'push.rocks' to 'localhost'
     port: 3001,
     path: '/',
     method: 'GET',
+    headers: {
+      host: 'push.rocks', // virtual host for routing
+    },
     rejectUnauthorized: false,
   });
 
@@ -207,15 +235,21 @@ tap.test('should route HTTPS requests based on host header', async () => {
 });
 
 tap.test('should handle unknown host headers', async () => {
+  // Connect to localhost but use an unknown host header.
   const response = await makeHttpsRequest({
-    hostname: 'unknown.host',
+    hostname: 'localhost', // connecting to localhost
     port: 3001,
     path: '/',
     method: 'GET',
+    headers: {
+      host: 'unknown.host', // this should not match any proxy config
+    },
     rejectUnauthorized: false,
-  }).catch((e) => e);
+  });
 
-  expect(response instanceof Error).toEqual(true);
+  // 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 () => {
@@ -224,23 +258,22 @@ tap.test('should support WebSocket connections', async () => {
   console.log('[TEST] Proxy server port:', 3001);
   console.log('\n[TEST] Starting WebSocket test');
 
-  // First configure the proxy with test certificates
+  // Reconfigure proxy with test certificates if necessary
-  console.log('[TEST] Configuring proxy with test certificates');
+  await testProxy.updateProxyConfigs([
-  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');
 
-    // Create WebSocket client with SSL/TLS options
+    // IMPORTANT: Connect to localhost but specify the SNI servername and Host header as "push.rocks"
-    const wsUrl = 'wss://push.rocks:3001';
+    const wsUrl = 'wss://localhost:3001'; // changed from 'wss://push.rocks:3001'
     console.log('[TEST] Creating WebSocket connection to:', wsUrl);
 
     const ws = new WebSocket(wsUrl, {
@@ -248,15 +281,15 @@ tap.test('should support WebSocket connections', async () => {
       handshakeTimeout: 5000,
       perMessageDeflate: false,
       headers: {
-        'Host': 'push.rocks',
+        Host: 'push.rocks', // required for SNI and routing on the proxy
-        'Connection': 'Upgrade',
+        Connection: 'Upgrade',
-        'Upgrade': 'websocket',
+        Upgrade: 'websocket',
-        'Sec-WebSocket-Version': '13'
+        'Sec-WebSocket-Version': '13',
       },
       protocol: 'echo-protocol',
       agent: new https.Agent({
-        rejectUnauthorized: false // Also needed for the underlying HTTPS connection
+        rejectUnauthorized: false, // Also needed for the underlying HTTPS connection
-      })
+      }),
     });
 
     console.log('[TEST] WebSocket client created');
@@ -286,7 +319,7 @@ tap.test('should support WebSocket connections', async () => {
     ws.on('upgrade', (response) => {
       console.log('[TEST] WebSocket upgrade response received:', {
         headers: response.headers,
-        statusCode: response.statusCode
+        statusCode: response.statusCode,
       });
     });
 
@@ -304,7 +337,10 @@ tap.test('should support WebSocket connections', async () => {
 
     ws.on('message', (message) => {
       console.log('[TEST] Received message:', message.toString());
-      if (message.toString() === 'Hello WebSocket') {
+      if (
+        message.toString() === 'Hello WebSocket' ||
+        message.toString() === 'Echo: Hello WebSocket'
+      ) {
         console.log('[TEST] Message received correctly');
         clearTimeout(timeout);
         cleanup();
@@ -320,7 +356,7 @@ tap.test('should support WebSocket connections', async () => {
     ws.on('close', (code, reason) => {
       console.log('[TEST] WebSocket connection closed:', {
         code,
-        reason: reason.toString()
+        reason: reason.toString(),
       });
       cleanup();
     });
@@ -328,15 +364,18 @@ tap.test('should support WebSocket connections', async () => {
 });
 
 tap.test('should handle custom headers', async () => {
-  testProxy.addDefaultHeaders({
+  await testProxy.addDefaultHeaders({
     'X-Proxy-Header': 'test-value',
   });
 
   const response = await makeHttpsRequest({
-    hostname: 'push.rocks',
+    hostname: 'localhost', // changed to 'localhost'
     port: 3001,
     path: '/',
     method: 'GET',
+    headers: {
+      host: 'push.rocks', // still routing to push.rocks
+    },
     rejectUnauthorized: false,
   });
 
@@ -353,16 +392,20 @@ tap.test('cleanup', async () => {
   });
 
   console.log('[TEST] Closing WebSocket server');
-  await new Promise<void>((resolve) => wsServer.close(() => {
+  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(() => {
+  await new Promise<void>((resolve) =>
+    testServer.close(() => {
       console.log('[TEST] Test server closed');
       resolve();
-  }));
+    })
+  );
 
   console.log('[TEST] Stopping proxy');
   await testProxy.stop();

ts/00_commitinfo_data.ts

@@ -3,6 +3,6 @@
  */
 export const commitinfo = {
   name: '@push.rocks/smartproxy',
-  version: '3.1.2',
+  version: '3.1.3',
   description: 'a proxy for handling high workloads of proxying'
 }

ts/smartproxy.classes.networkproxy.ts

@@ -8,12 +8,11 @@ export interface INetworkProxyOptions {
   port: number;
 }
 
-interface WebSocketWithHeartbeat extends plugins.wsDefault {
+interface IWebSocketWithHeartbeat extends plugins.wsDefault {
   lastPong: number;
 }
 
 export class NetworkProxy {
-  // INSTANCE
   public options: INetworkProxyOptions;
   public proxyConfigs: plugins.tsclass.network.IReverseProxyConfig[] = [];
   public httpsServer: plugins.https.Server;
@@ -43,20 +42,149 @@ export class NetworkProxy {
     }
   }
 
-  /**
-   * starts the proxyInstance
-   */
   public async start() {
+    // Instead of marking the callback async (which Node won't await),
+    // we call our async handler and catch errors.
     this.httpsServer = plugins.https.createServer(
       {
         key: this.defaultCertificates.key,
         cert: this.defaultCertificates.cert
       },
-      async (originRequest, originResponse) => {
+      (originRequest, originResponse) => {
+        this.handleRequest(originRequest, originResponse).catch((error) => {
+          console.error('Unhandled error in request handler:', error);
+          try {
+            originResponse.end();
+          } catch (err) {
+            // ignore errors during cleanup
+          }
+        });
+      },
+    );
+
+    // Enable websockets
+    const wsServer = new plugins.ws.WebSocketServer({ server: this.httpsServer });
+
+    // Set up the heartbeat interval
+    this.heartbeatInterval = setInterval(() => {
+      wsServer.clients.forEach((ws: plugins.wsDefault) => {
+        const wsIncoming = ws as IWebSocketWithHeartbeat;
+        if (!wsIncoming.lastPong) {
+          wsIncoming.lastPong = Date.now();
+        }
+        if (Date.now() - wsIncoming.lastPong > 5 * 60 * 1000) {
+          console.log('Terminating websocket due to missing pong for 5 minutes.');
+          wsIncoming.terminate();
+        } else {
+          wsIncoming.ping();
+        }
+      });
+    }, 60000); // runs every 1 minute
+
+    wsServer.on(
+      'connection',
+      (wsIncoming: IWebSocketWithHeartbeat, reqArg: plugins.http.IncomingMessage) => {
+        console.log(
+          `wss proxy: got connection for wsc for https://${reqArg.headers.host}${reqArg.url}`,
+        );
+
+        wsIncoming.lastPong = Date.now();
+        wsIncoming.on('pong', () => {
+          wsIncoming.lastPong = Date.now();
+        });
+
+        let wsOutgoing: plugins.wsDefault;
+        const outGoingDeferred = plugins.smartpromise.defer();
+
+        // --- Improvement 2: Only call routeReq once ---
+        const wsDestinationConfig = this.router.routeReq(reqArg);
+        if (!wsDestinationConfig) {
+          wsIncoming.terminate();
+          return;
+        }
+        try {
+          wsOutgoing = new plugins.wsDefault(
+            `ws://${wsDestinationConfig.destinationIp}:${wsDestinationConfig.destinationPort}${reqArg.url}`,
+          );
+          console.log('wss proxy: initiated outgoing proxy');
+          wsOutgoing.on('open', async () => {
+            outGoingDeferred.resolve();
+          });
+        } catch (err) {
+          console.error('Error initiating outgoing WebSocket:', err);
+          wsIncoming.terminate();
+          return;
+        }
+
+        wsIncoming.on('message', async (message, isBinary) => {
+          try {
+            await outGoingDeferred.promise;
+            wsOutgoing.send(message, { binary: isBinary });
+          } catch (error) {
+            console.error('Error sending message to wsOutgoing:', error);
+          }
+        });
+
+        wsOutgoing.on('message', async (message, isBinary) => {
+          try {
+            wsIncoming.send(message, { binary: isBinary });
+          } catch (error) {
+            console.error('Error sending message to wsIncoming:', error);
+          }
+        });
+
+        const terminateWsOutgoing = () => {
+          if (wsOutgoing) {
+            wsOutgoing.terminate();
+            console.log('Terminated outgoing ws.');
+          }
+        };
+        wsIncoming.on('error', terminateWsOutgoing);
+        wsIncoming.on('close', terminateWsOutgoing);
+
+        const terminateWsIncoming = () => {
+          if (wsIncoming) {
+            wsIncoming.terminate();
+            console.log('Terminated incoming ws.');
+          }
+        };
+        wsOutgoing.on('error', terminateWsIncoming);
+        wsOutgoing.on('close', terminateWsIncoming);
+      },
+    );
+
+    this.httpsServer.keepAliveTimeout = 600 * 1000;
+    this.httpsServer.headersTimeout = 600 * 1000;
+
+    this.httpsServer.on('connection', (connection: plugins.net.Socket) => {
+      this.socketMap.add(connection);
+      console.log(`Added connection. Now ${this.socketMap.getArray().length} sockets connected.`);
+      const cleanupConnection = () => {
+        if (this.socketMap.checkForObject(connection)) {
+          this.socketMap.remove(connection);
+          console.log(`Removed connection. ${this.socketMap.getArray().length} sockets remaining.`);
+          connection.destroy();
+        }
+      };
+      connection.on('close', cleanupConnection);
+      connection.on('error', cleanupConnection);
+      connection.on('end', cleanupConnection);
+      connection.on('timeout', cleanupConnection);
+    });
+
+    this.httpsServer.listen(this.options.port);
+    console.log(
+      `NetworkProxy -> OK: now listening for new connections on port ${this.options.port}`,
+    );
+  }
+
   /**
-         * endRequest function
+   * Internal async handler for processing HTTP/HTTPS requests.
-         * can be used to prematurely end a request
    */
+  private async handleRequest(
+    originRequest: plugins.http.IncomingMessage,
+    originResponse: plugins.http.ServerResponse,
+  ): Promise<void> {
     const endOriginReqRes = (
       statusArg: number = 404,
       messageArg: string = 'This route is not available on this server.',
@@ -87,26 +215,30 @@ export class NetworkProxy {
     if (destinationConfig.authentication) {
       const authInfo = destinationConfig.authentication;
       switch (authInfo.type) {
-            case 'Basic':
+        case 'Basic': {
           const authHeader = originRequest.headers.authorization;
-              if (authHeader) {
+          if (!authHeader) {
+            return endOriginReqRes(401, 'Authentication required', {
+              'WWW-Authenticate': 'Basic realm="Access to the staging site", charset="UTF-8"',
+            });
+          }
           if (!authHeader.includes('Basic ')) {
             return endOriginReqRes(401, 'Authentication required', {
               'WWW-Authenticate': 'Basic realm="Access to the staging site", charset="UTF-8"',
             });
           }
-                const authStringBase64 = originRequest.headers.authorization.replace('Basic ', '');
+          const authStringBase64 = authHeader.replace('Basic ', '');
           const authString: string = plugins.smartstring.base64.decode(authStringBase64);
           const userPassArray = authString.split(':');
           const user = userPassArray[0];
           const pass = userPassArray[1];
           if (user === authInfo.user && pass === authInfo.pass) {
-                  console.log('request successfully authenticated');
+            console.log('Request successfully authenticated');
           } else {
             return endOriginReqRes(403, 'Forbidden: Wrong credentials');
           }
-              }
           break;
+        }
         default:
           return endOriginReqRes(
             403,
@@ -134,7 +266,7 @@ export class NetworkProxy {
           },
           keepAlive: true,
         },
-            true, // lets make this streaming (keepAlive)
+        true, // streaming (keepAlive)
         (proxyRequest) => {
           originRequest.on('data', (data) => {
             proxyRequest.write(data);
@@ -185,131 +317,11 @@ export class NetworkProxy {
       console.error('Error while processing request:', error);
       endOriginReqRes(502, 'Bad Gateway: Error processing the request');
     }
-      },
-    );
-
-    // Enable websockets
-    const wsServer = new plugins.ws.WebSocketServer({ server: this.httpsServer });
-
-    // Set up the heartbeat interval
-    this.heartbeatInterval = setInterval(() => {
-      wsServer.clients.forEach((ws: plugins.wsDefault) => {
-        const wsIncoming = ws as WebSocketWithHeartbeat;
-        if (!wsIncoming.lastPong) {
-          wsIncoming.lastPong = Date.now();
-        }
-        if (Date.now() - wsIncoming.lastPong > 5 * 60 * 1000) {
-          console.log('Terminating websocket due to missing pong for 5 minutes.');
-          wsIncoming.terminate();
-        } else {
-          wsIncoming.ping();
-        }
-      });
-    }, 60000); // runs every 1 minute
-
-    wsServer.on(
-      'connection',
-      async (wsIncoming: WebSocketWithHeartbeat, reqArg: plugins.http.IncomingMessage) => {
-        console.log(
-          `wss proxy: got connection for wsc for https://${reqArg.headers.host}${reqArg.url}`,
-        );
-
-        wsIncoming.lastPong = Date.now();
-        wsIncoming.on('pong', () => {
-          wsIncoming.lastPong = Date.now();
-        });
-
-        let wsOutgoing: plugins.wsDefault;
-
-        const outGoingDeferred = plugins.smartpromise.defer();
-
-        try {
-          wsOutgoing = new plugins.wsDefault(
-            `ws://${this.router.routeReq(reqArg).destinationIp}:${
-              this.router.routeReq(reqArg).destinationPort
-            }${reqArg.url}`,
-          );
-          console.log('wss proxy: initiated outgoing proxy');
-          wsOutgoing.on('open', async () => {
-            outGoingDeferred.resolve();
-          });
-        } catch (err) {
-          console.log(err);
-          wsIncoming.terminate();
-          return;
   }
 
-        wsIncoming.on('message', async (message, isBinary) => {
+  public async updateProxyConfigs(
-          try {
+    proxyConfigsArg: plugins.tsclass.network.IReverseProxyConfig[],
-            await outGoingDeferred.promise;
+  ) {
-            wsOutgoing.send(message, { binary: isBinary });
-          } catch (error) {
-            console.error('Error sending message to wsOutgoing:', error);
-          }
-        });
-
-        wsOutgoing.on('message', async (message, isBinary) => {
-          try {
-            wsIncoming.send(message, { binary: isBinary });
-          } catch (error) {
-            console.error('Error sending message to wsIncoming:', error);
-          }
-        });
-
-        const terminateWsOutgoing = () => {
-          if (wsOutgoing) {
-            wsOutgoing.terminate();
-            console.log('terminated outgoing ws.');
-          }
-        };
-        wsIncoming.on('error', () => terminateWsOutgoing());
-        wsIncoming.on('close', () => terminateWsOutgoing());
-
-        const terminateWsIncoming = () => {
-          if (wsIncoming) {
-            wsIncoming.terminate();
-            console.log('terminated incoming ws.');
-          }
-        };
-        wsOutgoing.on('error', () => terminateWsIncoming());
-        wsOutgoing.on('close', () => terminateWsIncoming());
-      },
-    );
-
-    this.httpsServer.keepAliveTimeout = 600 * 1000;
-    this.httpsServer.headersTimeout = 600 * 1000;
-
-    this.httpsServer.on('connection', (connection: plugins.net.Socket) => {
-      this.socketMap.add(connection);
-      console.log(`added connection. now ${this.socketMap.getArray().length} sockets connected.`);
-      const cleanupConnection = () => {
-        if (this.socketMap.checkForObject(connection)) {
-          this.socketMap.remove(connection);
-          console.log(`removed connection. ${this.socketMap.getArray().length} sockets remaining.`);
-          connection.destroy();
-        }
-      };
-      connection.on('close', () => {
-        cleanupConnection();
-      });
-      connection.on('error', () => {
-        cleanupConnection();
-      });
-      connection.on('end', () => {
-        cleanupConnection();
-      });
-      connection.on('timeout', () => {
-        cleanupConnection();
-      });
-    });
-
-    this.httpsServer.listen(this.options.port);
-    console.log(
-      `NetworkProxy -> OK: now listening for new connections on port ${this.options.port}`,
-    );
-  }
-
-  public async updateProxyConfigs(proxyConfigsArg: plugins.tsclass.network.IReverseProxyConfig[]) {
     console.log(`got new proxy configs`);
     this.proxyConfigs = proxyConfigsArg;
     this.router.setNewProxyConfigs(proxyConfigsArg);
@@ -347,9 +359,9 @@ export class NetworkProxy {
     this.httpsServer.close(() => {
       done.resolve();
     });
-    await this.socketMap.forEach(async (socket) => {
+    for (const socket of this.socketMap.getArray()) {
       socket.destroy();
-    });
+    }
     await done.promise;
     clearInterval(this.heartbeatInterval);
     console.log('NetworkProxy -> OK: Server has been stopped and all connections closed.');