From 32bd5e2e28dab95f20f6d8616cb29e0735795dfa Mon Sep 17 00:00:00 2001
From: Blaine Schmeisser <blainesch@gmail.com>
Date: Tue, 12 May 2015 22:41:26 -0700
Subject: [PATCH] Add ability to use a function for the `lookup` option.

---
 README.md      | 16 +++++++++++++++-
 index.js       |  7 +++++++
 tests/index.js | 26 ++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index d223394..7c48baa 100644
--- a/README.md
+++ b/README.md
@@ -41,7 +41,7 @@ limiter(options)
 
  - `path`: `String` *optional* route path to the request
  - `method`: `String` *optional* http method. accepts `get`, `post`, `put`, `delete`, and of course Express' `all`
- - `lookup`: `String|Array.<String>` value lookup on the request object. Can be a single value or array. See [examples](#examples) for common usages
+ - `lookup`: `Function|String|Array.<String>` value lookup on the request object. Can be a single value, array or function. See [examples](#examples) for common usages
  - `total`: `Number` allowed number of requests before getting rate limited
  - `expire`: `Number` amount of time in `ms` before the rate-limited is reset
  - `whitelist`: `function(req)` optional param allowing the ability to whitelist. return `boolean`, `true` to whitelist, `false` to passthru to limiter.
@@ -114,6 +114,20 @@ limiter({
   }
 })
 
+// with a function for dynamic-ness
+limiter({
+  lookup: function(req, res, opts, next) {
+    if (validApiKey(req.query.api_key)) {
+      opts.lookup = 'query.api_key'
+      opts.total = 100
+    } else {
+      opts.lookup = 'connection.remoteAddress'
+      opts.total = 10
+    }
+    return next()
+  }
+})
+
 ```
 
 ### as direct middleware
diff --git a/index.js b/index.js
index 4d17f0a..a613cfa 100644
--- a/index.js
+++ b/index.js
@@ -48,6 +48,13 @@ module.exports = function (app, db) {
 
       })
     }
+    if (typeof(opts.lookup) === 'function') {
+      middleware = function (middleware, req, res, next) {
+        return opts.lookup(req, res, opts, function () {
+          return middleware(req, res, next)
+        })
+      }.bind(this, middleware)
+    }
     if (opts.method && opts.path) app[opts.method](opts.path, middleware)
     return middleware
   }
diff --git a/tests/index.js b/tests/index.js
index f6f1da9..a439dd4 100644
--- a/tests/index.js
+++ b/tests/index.js
@@ -127,6 +127,32 @@ describe('rate-limiter', function () {
         stub.restore()
       })
     })
+
+    it('should process lookup as a function', function (done) {
+      limiter({
+        path: '*',
+        method: 'all',
+        lookup: function (req, res, opts, next) {
+          opts.lookup = 'query.api_key';
+          opts.total = 20
+          return next()
+        },
+        total: 3,
+        expire: 1000 * 60 * 60
+      })
+
+      app.get('/route', function (req, res) {
+        res.send(200, 'hello')
+      })
+
+      request(app)
+      .get('/route?api_key=foobar')
+      .expect('X-RateLimit-Limit', 20)
+      .expect('X-RateLimit-Remaining', 19)
+      .expect(200, function (e) {
+        done(e)
+      })
+    })
   })
 
   context('direct middleware', function () {