diff --git a/README.md b/README.md index 5b8305e..0f08d66 100644 --- a/README.md +++ b/README.md @@ -294,6 +294,37 @@ If this callback is specified, it will be called with the `req` object on a deni const ddos = new Ddos({ limit: 2, onDenial }); ``` +### whiteListHookSync + +A synchronous function for whitelisting. In case you need to whitelist by a criteria different than IP. +For example, if Admin users can make infinite requests. It must return true or false. +```js + const whiteListHookSync = function(req) { + return user.isAdmin(req.headers.authorization); + }; + + const ddos = new Ddos({ limit: 2, whiteListHookSync }); +``` + +### whiteListHook + +An asynchronous function for whitelisting. In case you need to whitelist by a criteria different than IP. +For example, if Admin users can make infinite requests. It must return true or false. +```js + const whiteListHook = function(req) { + return new Promise((resolve,reject)=>{ + if(req.headers.authorization === 'pass'){ + return resolve(); // Whitelist this request + }else { + return reject(); // Continue process + } + }); + }; + + const ddos = new Ddos({ limit: 2, whiteListHook }); +``` + + Contribute ========== diff --git a/lib/index.js b/lib/index.js index 2252783..6b20fdf 100644 --- a/lib/index.js +++ b/lib/index.js @@ -58,55 +58,77 @@ const _handle = function(options, table, req) { if (options.whitelist.indexOf(host) != -1) { return resolve(); } - if (options.includeUserAgent) - host = host.concat("#" + req.headers["user-agent"]); - if (!table[host]) - table[host] = { count: 1, expiry: 1 }; - else { - table[host].count++; - if (table[host].count > options.maxcount) - table[host].count = options.maxcount; - if ((table[host].count > options.burst) && (table[host].expiry <= options.maxexpiry)) { - table[host].expiry = Math.min( - options.maxexpiry, - table[host].expiry * 2 - ); - } else { - table[host].expiry = 1; + if(!options.whiteListHook){ + if(options.whiteListHookSync){ + if(options.whiteListHookSync(req)){ + console.log("Whitelisted!"); + return resolve(); + } } + return completeHandle(options,table,req,resolve,reject,host); } - if (options.testmode) { - console.log("ddos: handle: end:", table); + else{ + options + .whiteListHook(req) + .then(()=>{ + return resolve(); + }) + .catch(()=>{ + return completeHandle(options,table,req,resolve,reject,host); + }) } + }) +}; - if (table[host].count > options.limit) { - (!options.testmode) && (console.log("ddos: denied: entry:", host, table[host])); - if (options.testmode) { - return reject({action:'respond', code:429, message:JSON.stringify(table[host])}); - } else { - return reject({action:'respond', code:options.responseStatus, message:options.errormessage}) - } +const completeHandle = function(options, table, req,resolve,reject, host){ + if (options.includeUserAgent) + host = host.concat("#" + req.headers["user-agent"]); + if (!table[host]) + table[host] = { count: 1, expiry: 1 }; + else { + table[host].count++; + if (table[host].count > options.maxcount) + table[host].count = options.maxcount; + if ((table[host].count > options.burst) && (table[host].expiry <= options.maxexpiry)) { + table[host].expiry = Math.min( + options.maxexpiry, + table[host].expiry * 2 + ); } else { - return resolve(); + table[host].expiry = 1; } - }) -}; + } + if (options.testmode) { + console.log("ddos: handle: end:", table); + } + + if (table[host].count > options.limit) { + (!options.testmode) && (console.log("ddos: denied: entry:", host, table[host])); + if (options.testmode) { + return reject({action:'respond', code:429, message:JSON.stringify(table[host])}); + } else { + return reject({action:'respond', code:options.responseStatus, message:options.errormessage}) + } + } else { + return resolve(); + } +} exports._handle = _handle; exports.handle = function (req, res, next) { return _handle(this.params, this.table, req) - .then(() => next()) - .catch((e) => { - if (e.action === "nothing") { - return next(); - } - if (e.action === "respond") { - if (this.params.onDenial) { - this.params.onDenial(req) - } + .then(() => next()) + .catch((e) => { + if (e.action === "nothing") { + return next(); + } + if (e.action === "respond") { + if (this.params.onDenial) { + this.params.onDenial(req) + } - res.writeHead(e.code, {'Content-Type':'application/json'}); - return res.end(e.message); - } - }) + res.writeHead(e.code, {'Content-Type':'application/json'}); + return res.end(e.message); + } + }) }; diff --git a/test/test-options.js b/test/test-options.js index 6a609ca..f580571 100644 --- a/test/test-options.js +++ b/test/test-options.js @@ -134,10 +134,77 @@ tape("options - onDenial ", function(t) { }; doCall() - .then(() => doCall()) - .then(() => doCall()) - .then(res => { - t.equals(count, 2); - ddos.end(); + .then(() => doCall()) + .then(() => doCall()) + .then(res => { + t.equals(count, 2); + ddos.end(); + }); +}); + + +tape("options - whiteListHookSync ", function(t) { + t.plan(1); + let count = 0; + const whiteListHookSync = function(req) { + return req.headers.authorization === 'pass'; + }; + const ddos = new Ddos({ limit: 1, whiteListHookSync }); + const app = express(); + app.use(ddos.express); + app.get("/user", (req, res) => { + console.log("reply"); + res.status(200).json({ name: "john" }); + count++; + }); + + const doPassCall = function() { + return request(app) + .set('Authorization','pass') + .get("/user"); + }; + + doPassCall() + .then(() => doPassCall()) + .then(() => doPassCall()) + .then(() => { + t.equals(count, 3); + ddos.end(); + }); +}); + +tape("options - whiteListHook ", function(t) { + t.plan(1); + let count = 0; + const whiteListHook = function(req) { + return new Promise((resolve,reject)=>{ + if(req.headers.authorization === 'pass'){ + return resolve(); + }else { + return reject(); + } }); + }; + const ddos = new Ddos({ limit: 1, whiteListHook }); + const app = express(); + app.use(ddos.express); + app.get("/user", (req, res) => { + console.log("reply"); + res.status(200).json({ name: "john" }); + count++; + }); + + const doPassCall = function() { + return request(app) + .set('Authorization','pass') + .get("/user"); + }; + + doPassCall(200) + .then(() => doPassCall(200)) + .then(() => doPassCall(200)) + .then(() => { + t.equals(count, 3); + ddos.end(); + }); });