diff --git a/lib/middleware.js b/lib/middleware.js index a95aa96..3bfb769 100644 --- a/lib/middleware.js +++ b/lib/middleware.js @@ -1,4 +1,6 @@ -var rateLimiter = require('./rate-limiter'); +var rateLimiter = require('./rate-limiter'); + +var defaultFailMessage = 'Rate limit exceeded, retry after {after}ms'; module.exports = function(opts) { var limiter = rateLimiter(opts); @@ -7,8 +9,24 @@ module.exports = function(opts) { if (err) { next(); } else { + res.set({ + 'X-RateLimit-Limit': rate.limit, + 'X-RateLimit-Remaining': rate.current > rate.limit ? 0 : rate.limit - rate.current, + 'X-RateLimit-Reset': (rate.reset / 1000) | 0 + }); if (rate.current > rate.limit) { - res.writeHead(429); + var after = rate.reset - Date.now(); + var failMessage = defaultFailMessage; + res.set('Retry-After', (after / 1000) | 0); + if(opts.failMessage) { + if(typeof opts.failMessage === 'function') { + failMessage = opts.failMessage(req); + } + else { + failMessage = opts.failMessage; + } + } + res.status(429).send(failMessage.replace(/{after}/g, after)); res.end(); } else { next(); diff --git a/lib/rate-limiter.js b/lib/rate-limiter.js index e022841..e6b009f 100644 --- a/lib/rate-limiter.js +++ b/lib/rate-limiter.js @@ -10,13 +10,18 @@ module.exports = function(opts) { .setex(tempKey, opts.window(), 0) .renamenx(tempKey, realKey) .incr(realKey) - .ttl(realKey) + .pttl(realKey) .exec(function(err, results) { if(err) { callback(err); } else { - if (results[3] == -1) { // automatically recover from possible race condition - opts.redis.expire(realKey,opts.window()); + var reset, ttl = results[3]; + if(ttl == -1) { // automatically recover from possible race condition + opts.redis.expire(realKey, opts.window()); + reset = Date.now() + opts.window() * 1000; + } + else { + reset = Date.now() + ttl; } var current = results[2]; callback(null, { @@ -24,7 +29,8 @@ module.exports = function(opts) { current: current, limit: opts.limit(), window: opts.window(), - over: (current > opts.limit()) + over: (current > opts.limit()), + reset: reset }); } });