Eps1lon/add throttle headers

lib/plugins/throttle.js

@@ -204,6 +204,8 @@ TokenTable.prototype.get = function get(key) {
  *                   - {Boolean} ip (optional).
  *                   - {Boolean} username (optional).
  *                   - {Boolean} xff (optional).
+ *                   - {Boolean} setHeaders: Set response headers for rate,
+ *                               limit (burst) and remaining. Default is false.
  *                   - {Object} overrides (optional).
  *                   - {Object} tokensTable: a storage engine this plugin will
  *                              use to store throttling keys -> bucket mappings.
@@ -220,6 +222,7 @@ function throttle(options) {
     assert.object(options, 'options');
     assert.number(options.burst, 'options.burst');
     assert.number(options.rate, 'options.rate');
+    assert.optionalBool(options.setHeaders, 'options.setHeaders');
 
     if (!xor(options.ip, options.xff, options.username)) {
         throw new Error('(ip ^ username ^ xff)');
@@ -278,7 +281,16 @@ function throttle(options) {
         req.log.trace('Throttle(%s): num_tokens= %d',
             attr, bucket.tokens);
 
-        if (!bucket.consume(1)) {
+        var tooManyRequests = !bucket.consume(1);
+
+        // set throttle headers after consume which changes the remaining tokens
+        if (options.setHeaders) {
+            res.header('X-RateLimit-Remaining', Math.floor(bucket.tokens));
+            res.header('X-RateLimit-Limit', burst);
+            res.header('X-RateLimit-Rate', rate);
+        }
+
+        if (tooManyRequests) {
             req.log.info({
                 address: req.connection.remoteAddress || '?',
                 method: req.method,

test/plugins/throttle.test.js

@@ -15,25 +15,45 @@ var SERVER;
 
 var errorMessage = 'Error message should include rate 0.5 r/s. Received: ';
 
-///--- Tests
+function setupClientServer(ip, throttleOptions, done) {
+    var server = restify.createServer({
+        dtrace: helper.dtrace,
+        log: helper.getLog('server')
+    });
 
-describe('throttle plugin', function () {
+    server.use(function ghettoAuthenticate(req, res, next) {
+        if (req.params.name) {
+            req.username = req.params.name;
+        }
 
-    before(function setup(done) {
-        SERVER = restify.createServer({
+        next();
+    });
+
+    server.use(restify.plugins.throttle(throttleOptions));
+
+    server.get('/test/:name', function (req, res, next) {
+        res.send();
+        next();
+    });
+
+    server.listen(PORT, ip, function () {
+        PORT = server.address().port;
+        var client = restifyClients.createJsonClient({
+            url: 'http://' + ip + ':' + PORT,
             dtrace: helper.dtrace,
-            log: helper.getLog('server')
+            retry: false
         });
 
-        SERVER.use(function ghettoAuthenticate(req, res, next) {
-            if (req.params.name) {
-                req.username = req.params.name;
-            }
+        done(client, server);
+    });
+}
 
-            next();
-        });
+///--- Tests
+
+describe('throttle plugin', function () {
 
-        SERVER.use(restify.plugins.throttle({
+    before(function setup(done) {
+        setupClientServer('127.0.0.1', {
             burst: 1,
             rate: 0.5,
             username: true,
@@ -47,21 +67,9 @@ describe('throttle plugin', function () {
                     rate: 1
                 }
             }
-        }));
-
-        SERVER.get('/test/:name', function (req, res, next) {
-            res.send();
-            next();
-        });
-
-        SERVER.listen(PORT, '127.0.0.1', function () {
-            PORT = SERVER.address().port;
-            CLIENT = restifyClients.createJsonClient({
-                url: 'http://127.0.0.1:' + PORT,
-                dtrace: helper.dtrace,
-                retry: false
-            });
-
+        }, function setupGlobal(client, server) {
+            CLIENT = client;
+            SERVER = server;
             done();
         });
     });
@@ -170,4 +178,50 @@ describe('throttle plugin', function () {
             });
         });
     });
+
+    it('should not expose rate limit headers per default', function (done) {
+        CLIENT.get('/test/throttleMe', function (err, _, res) {
+            assert.isUndefined(res.headers['x-ratelimit-limit']);
+            assert.isUndefined(res.headers['x-ratelimit-rate']);
+            assert.isUndefined(res.headers['x-ratelimit-rate']);
+
+            done();
+        });
+    });
+
+    it('should expose headers on options set', function (done) {
+        // setup a new server with headers set to true since we cant
+        // change throttle options after init
+        setupClientServer('127.0.0.2', {
+            burst: 17,
+            rate: 0.1,
+            username: true,
+            setHeaders: true
+        }, function setupWithHeaders (client, server) {
+            client.get('/test/throttleMe', function (err, req, res) {
+                assert.equal(res.headers['x-ratelimit-limit'], '17');
+                assert.equal(res.headers['x-ratelimit-rate'], '0.1');
+                assert.equal(res.headers['x-ratelimit-remaining'], '16');
+
+                // it should count down
+                client.get(
+                    '/test/throttleMe',
+                    function (nextErr, nextReq, nextRes) {
+                        assert.equal(
+                            nextRes.headers['x-ratelimit-limit'], '17'
+                        );
+                        assert.equal(
+                            nextRes.headers['x-ratelimit-rate'], '0.1'
+                        );
+                        assert.equal(
+                            nextRes.headers['x-ratelimit-remaining'], '15'
+                        );
+
+                        client.close();
+                        server.close(done);
+                    }
+                );
+            });
+        });
+    });
 });