Added sync/async whitelist hooks options

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
 ==========
 

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);
+        }
+      })
 };

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();
+      });
 });