From 5a2d3512ca504470e981da432c0acbb5d36dc2c7 Mon Sep 17 00:00:00 2001
From: Ylian Saint-Hilaire <ysainthilaire@hotmail.com>
Date: Tue, 22 Jan 2019 19:34:55 -0800
Subject: [PATCH] Added agent trouble detection on the server.

---
 meshagent.js | 24 ++++++++++++++++++++++++
 package.json |  2 +-
 webserver.js |  5 +++++
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/meshagent.js b/meshagent.js
index f0433a77..3b8b7343 100644
--- a/meshagent.js
+++ b/meshagent.js
@@ -326,6 +326,13 @@ module.exports.CreateMeshAgent = function (parent, db, ws, req, args, domain) {
         if (obj.nodeid != null) {
             //console.log('Agent disconnect ' + obj.nodeid + ' (' + obj.remoteaddrport + ') id=' + obj.agentInfo.agentId);
             obj.parent.parent.debug(1, 'Agent disconnect ' + obj.nodeid + ' (' + obj.remoteaddrport + ') id=' + obj.agentInfo.agentId);
+
+            // Log the agent disconnection
+            if (obj.parent.wsagentsDisconnections[obj.nodeid] == null) {
+                obj.parent.wsagentsDisconnections[obj.nodeid] = 1;
+            } else {
+                obj.parent.wsagentsDisconnections[obj.nodeid] = ++obj.parent.wsagentsDisconnections[obj.nodeid];
+            }
         }
         obj.close(0);
     });
@@ -411,9 +418,25 @@ module.exports.CreateMeshAgent = function (parent, db, ws, req, args, domain) {
             delete obj.pendingCompleteAgentConnection;
             obj.authenticated = 2;
 
+            // Check how many times this agent disconnected in the last few minutes.
+            var disconnectCount = obj.parent.wsagentsDisconnections[obj.nodeid];
+            if (disconnectCount > 6) {
+                console.log('Agent in big trouble: NodeId=' + obj.nodeid + ', IP=' + obj.remoteaddrport + ', Agent=' + obj.agentInfo.agentId + '.');
+                // TODO: Log or do something to recover?
+                return;
+            }
+
             // Command 4, inform mesh agent that it's authenticated.
             obj.send(obj.common.ShortToStr(4));
 
+            if (disconnectCount > 4) {
+                // Too many disconnections, this agent has issues. Just clear the core.
+                obj.send(obj.common.ShortToStr(10) + obj.common.ShortToStr(0));
+                console.log('Agent in trouble: NodeId=' + obj.nodeid + ', IP=' + obj.remoteaddrport + ', Agent=' + obj.agentInfo.agentId + '.');
+                // TODO: Log or do something to recover?
+                return;
+            }
+
             // Check if we need to make an native update check
             obj.agentExeInfo = obj.parent.parent.meshAgentBinaries[obj.agentInfo.agentId];
             const corename = obj.parent.parent.meshAgentsArchitectureNumbers[obj.agentInfo.agentId].core;
@@ -468,6 +491,7 @@ module.exports.CreateMeshAgent = function (parent, db, ws, req, args, domain) {
                     }
                 });
             }
+            
         });
     }
 
diff --git a/package.json b/package.json
index d9346ac2..377ca709 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "meshcentral",
-  "version": "0.2.6-t",
+  "version": "0.2.6-v",
   "keywords": [
     "Remote Management",
     "Intel AMT",
diff --git a/webserver.js b/webserver.js
index 40103bf2..b0d615ad 100644
--- a/webserver.js
+++ b/webserver.js
@@ -140,6 +140,8 @@ module.exports.CreateWebServer = function (parent, db, args, certificates) {
 
     // Main lists
     obj.wsagents = {};
+    obj.wsagentsDisconnections = {};
+    obj.wsagentsDisconnectionsTimer = null;
     obj.wssessions = {};            // UserId --> Array Of Sessions
     obj.wssessions2 = {};           // "UserId + SessionRnd" --> Session  (Note that the SessionId is the UserId + / + SessionRnd)
     obj.wsPeerSessions = {};        // ServerId --> Array Of "UserId + SessionRnd"
@@ -2004,6 +2006,9 @@ module.exports.CreateWebServer = function (parent, db, args, certificates) {
 
             // Indicates to ExpressJS that the public folder should be used to serve static files.
             obj.app.use(url, obj.express.static(obj.path.join(__dirname, 'public')));
+
+            // Start regular disconnection list flush every 2 minutes.
+            obj.wsagentsDisconnectionsTimer = setInterval(function () { obj.wsagentsDisconnections = {}; }, 120000);
         }
 
         // Start server on a free port