Merge pull request #58 from pierotofy/flood

Added flood monitor

Author: pierotofy

config-default.json

@@ -12,6 +12,7 @@
     "debug": false,
     "log-level": "info",
     "upload-max-speed": 0,
+    "flood-limit": 0,
     "ssl-key": "",
     "ssl-cert": "",
     "asr-provider": "",

config.js

@@ -45,7 +45,7 @@ let argDefs = {
     default: defaultConfig,
 
     int: ['port', 'admin-cli-port', 'admin-web-port', 
-          'secure-port', 'upload-max-speed'] // for cast only, not used by minimist
+          'secure-port', 'upload-max-speed', 'flood-limit'] // for cast only, not used by minimist
 };
 let argv = require('minimist')(process.argv.slice(2), argDefs);
 
@@ -66,6 +66,7 @@ Options:
     --downloads-from-s3 <URL>	Manually set the S3 URL prefix where to redirect /task/<uuid>/download requests. (default: do not use S3, forward download requests to nodes, unless the autoscaler is setup, in which case the autoscaler's S3 configuration is used) 
     --no-splitmerge	By default the program will set itself as being a cluster node for all split/merge tasks. Setting this option disables it. (default: false)
     --public-address <http(s)://host:port>	Should be set to a public URL that nodes can use to reach ClusterODM. (default: match the "host" header from client's HTTP request)
+    --flood-limit <number>	Limit the number of simultaneous task uploads that a user can initiate concurrently (default: no limit)
     --token <token> Sets a token that needs to be passed for every request. This can be used to limit access to the node only to token holders. (default: none)
     --debug 	Disable caches and other settings to facilitate debug (default: false)
     --ssl-key <file>	Path to .pem SSL key file

index.js

@@ -22,6 +22,7 @@ const logger = require('./libs/logger');
 const package_info = require('./package_info');
 const nodes = require('./libs/nodes');
 const proxy = require('./libs/proxy');
+const floodMonitor = require('./libs/floodMonitor');
 const routetable = require('./libs/routetable');
 
 (async function(){
@@ -33,6 +34,7 @@ const routetable = require('./libs/routetable');
     const cloudProvider = (require('./libs/cloudProvider')).initialize(config.cloud_provider);
     await (require('./libs/asrProvider')).initialize(config.asr);
     await nodes.initialize();
+    floodMonitor.initialize();
 
     const proxies = await proxy.initialize(cloudProvider);
 

libs/floodMonitor.js

@@ -0,0 +1,74 @@
+/**
+ *  ClusterODM - A reverse proxy, load balancer and task tracker for NodeODM
+ *  Copyright (C) 2018-present MasseranoLabs LLC
+ *
+ *  This program is free software: you can redistribute it and/or modify
+ *  it under the terms of the GNU Affero General Public License as
+ *  published by the Free Software Foundation, either version 3 of the
+ *  License, or (at your option) any later version.
+ *
+ *  This program is distributed in the hope that it will be useful,
+ *  but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *  GNU Affero General Public License for more details.
+ *
+ *  You should have received a copy of the GNU Affero General Public License
+ *  along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+const config = require('../config');
+
+let userTasks = null;
+
+module.exports = {
+    initialize: function(){
+        userTasks = {};
+
+        const forgive = () => {
+            Object.keys(userTasks).forEach(userToken => {
+                if (userTasks[userToken].count > 0){
+                    userTasks[userToken].count = Math.floor(userTasks[userToken].count * 0.66);
+                }
+                
+                if (userTasks[userToken].count <= 0){
+                    delete(userTasks[userToken]);
+                }
+            });
+        }
+
+        setInterval(forgive, 1000 * 60 * this.FORGIVE_TIME);
+    },
+
+    FORGIVE_TIME: 15, // minutes
+
+    recordTaskInit: function(userToken){
+        this.modifyRecord(userToken, record => {
+            record.count = record.count ? (record.count + 1) : 1;
+        });
+    },
+
+    recordTaskCommit: function(userToken){
+        this.modifyRecord(userToken, record => {
+            record.count = Math.max(record.count - 1, 0);
+        });
+    },
+
+    isFlooding: function(userToken){
+        if (config.flood_limit <= 0) return false; // Disabled
+        if (!userToken) userToken = "default";
+
+        const record = userTasks[userToken];
+        if (!record) return false; // No record
+
+        return record.count > config.flood_limit;
+    },
+
+    modifyRecord: function(userToken, callback){
+        if (!userToken) userToken = "default";
+        
+        const record = userTasks[userToken] || {};
+
+        callback(record);
+
+        userTasks[userToken] = record;
+    }
+};

libs/proxy.js

@@ -35,6 +35,7 @@ const taskNew = require('./taskNew');
 const async = require('async');
 const odmOptions = require('./odmOptions');
 const asrProvider = require('./asrProvider');
+const floodMonitor = require('./floodMonitor');
 
 module.exports = {
 	initialize: async function(cloudProvider){
@@ -259,6 +260,13 @@ module.exports = {
                             return;
                         }
 
+                        floodMonitor.recordTaskInit(query.token);
+                        
+                        if (floodMonitor.isFlooding(query.token)){
+                            die(`Uuh, slow down! It seems like you are sending a lot of tasks. Check that your connection is not dropping, or wait ${floodMonitor.FORGIVE_TIME} minutes and try again.`);
+                            return;
+                        }
+
                         // Save
                         fs.writeFile(path.join(tmpPath, "body.json"),
                                     JSON.stringify(params), {encoding: 'utf8'}, err => {
@@ -307,6 +315,8 @@ module.exports = {
                             return;
                         }
 
+                        floodMonitor.recordTaskCommit(query.token);
+
                         async.series([
                             cb => {
                                 fs.readFile(bodyFile, 'utf8', (err, data) => {