From 459ee7130fb216111d3d965170d783d24268d618 Mon Sep 17 00:00:00 2001
From: Philipp Kunz <gitlab@philkunz.com>
Date: Mon, 7 Oct 2024 12:52:01 +0200
Subject: [PATCH] feat(NetworkProxy): Introduce WebSocket heartbeat to maintain
 active connections in NetworkProxy

---
 changelog.md                          |  7 +++++
 ts/00_commitinfo_data.ts              |  2 +-
 ts/smartproxy.classes.networkproxy.ts | 38 ++++++++++++++++++++++-----
 3 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/changelog.md b/changelog.md
index c58ae9d..1488644 100644
--- a/changelog.md
+++ b/changelog.md
@@ -1,5 +1,12 @@
 # Changelog
 
+## 2024-10-07 - 3.1.0 - feat(NetworkProxy)
+Introduce WebSocket heartbeat to maintain active connections in NetworkProxy
+
+- Added heartbeat mechanism to WebSocket connections to ensure they remain active.
+- Terminating WebSocket if no pong is received for 5 minutes.
+- Set up heartbeat interval to run every 1 minute for connection checks.
+
 ## 2024-10-07 - 3.0.61 - fix(networkproxy)
 Improve error handling for proxy requests
 
diff --git a/ts/00_commitinfo_data.ts b/ts/00_commitinfo_data.ts
index 66332c8..8a57fb8 100644
--- a/ts/00_commitinfo_data.ts
+++ b/ts/00_commitinfo_data.ts
@@ -3,6 +3,6 @@
  */
 export const commitinfo = {
   name: '@push.rocks/smartproxy',
-  version: '3.0.61',
+  version: '3.1.0',
   description: 'a proxy for handling high workloads of proxying'
 }
diff --git a/ts/smartproxy.classes.networkproxy.ts b/ts/smartproxy.classes.networkproxy.ts
index 1923855..3ea5313 100644
--- a/ts/smartproxy.classes.networkproxy.ts
+++ b/ts/smartproxy.classes.networkproxy.ts
@@ -5,6 +5,10 @@ export interface INetworkProxyOptions {
   port: number;
 }
 
+interface WebSocketWithHeartbeat extends plugins.wsDefault {
+  lastPong: number;
+}
+
 export class NetworkProxy {
   // INSTANCE
   public options: INetworkProxyOptions;
@@ -13,6 +17,7 @@ export class NetworkProxy {
   public router = new ProxyRouter();
   public socketMap = new plugins.lik.ObjectMap<plugins.net.Socket>();
   public defaultHeaders: { [key: string]: string } = {};
+  public heartbeatInterval: NodeJS.Timeout;
 
   public alreadyAddedReverseConfigs: {
     [hostName: string]: plugins.tsclass.network.IReverseProxyConfig;
@@ -198,7 +203,7 @@ JNj2Dr5H0XoLFFnvuvzcRbhlJ9J67JzR+7g=
               },
               keepAlive: true,
             },
-            true, // lets make this streaming
+            true, // lets make this streaming (keepAlive)
             (proxyRequest) => {
               originRequest.on('data', (data) => {
                 proxyRequest.write(data);
@@ -254,13 +259,35 @@ JNj2Dr5H0XoLFFnvuvzcRbhlJ9J67JzR+7g=
 
     // 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: plugins.wsDefault, reqArg: plugins.http.IncomingMessage) => {
+      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();
@@ -315,8 +342,6 @@ JNj2Dr5H0XoLFFnvuvzcRbhlJ9J67JzR+7g=
         };
         wsOutgoing.on('error', () => terminateWsIncoming());
         wsOutgoing.on('close', () => terminateWsIncoming());
-
-        
       }
     );
 
@@ -358,8 +383,6 @@ JNj2Dr5H0XoLFFnvuvzcRbhlJ9J67JzR+7g=
     this.proxyConfigs = proxyConfigsArg;
     this.router.setNewProxyConfigs(proxyConfigsArg);
     for (const hostCandidate of this.proxyConfigs) {
-      // console.log(hostCandidate);
-
       const existingHostNameConfig = this.alreadyAddedReverseConfigs[hostCandidate.hostName];
 
       if (!existingHostNameConfig) {
@@ -397,6 +420,7 @@ JNj2Dr5H0XoLFFnvuvzcRbhlJ9J67JzR+7g=
       socket.destroy();
     });
     await done.promise;
+    clearInterval(this.heartbeatInterval);
     console.log('NetworkProxy -> OK: Server has been stopped and all connections closed.');
   }
-}
\ No newline at end of file
+}