diff --git a/lib/plugins/throttle.js b/lib/plugins/throttle.js index a1c96729d..c27b0a4a5 100644 --- a/lib/plugins/throttle.js +++ b/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, diff --git a/test/plugins/throttle.test.js b/test/plugins/throttle.test.js index 1237ac2b7..a9dde7ff7 100644 --- a/test/plugins/throttle.test.js +++ b/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); + } + ); + }); + }); + }); });