// Copyright 2012 Mark Cavage <mcavage@gmail.com> All rights reserved.
|
|
'use strict';
|
|
var sprintf = require('util').format;
|
|
var assert = require('assert-plus');
|
var LRU = require('lru-cache');
|
var errors = require('restify-errors');
|
|
///--- Globals
|
|
var TooManyRequestsError = errors.TooManyRequestsError;
|
|
var MESSAGE = 'You have exceeded your request rate of %s r/s.';
|
|
///--- Helpers
|
|
/**
|
* @private
|
* @function xor
|
* @returns {undefined} no return value
|
*/
|
function xor() {
|
var x = false;
|
|
for (var i = 0; i < arguments.length; i++) {
|
if (arguments[i] && !x) {
|
x = true;
|
} else if (arguments[i] && x) {
|
return false;
|
}
|
}
|
return x;
|
}
|
|
///--- Internal Class (TokenBucket)
|
|
/**
|
* An implementation of the Token Bucket algorithm.
|
*
|
* Basically, in network throttling, there are two "mainstream"
|
* algorithms for throttling requests, Token Bucket and Leaky Bucket.
|
* For restify, I went with Token Bucket. For a good description of the
|
* algorithm, see: http://en.wikipedia.org/wiki/Token_bucket
|
*
|
* In the options object, you pass in the total tokens and the fill rate.
|
* Practically speaking, this means "allow `fill rate` requests/second,
|
* with bursts up to `total tokens`". Note that the bucket is initialized
|
* to full.
|
*
|
* Also, in googling, I came across a concise python implementation, so this
|
* is just a port of that. Thanks http://code.activestate.com/recipes/511490 !
|
*
|
* @private
|
* @class TokenBucket
|
* @param {Object} options - contains the parameters:
|
* - {Number} capacity the maximum burst.
|
* - {Number} fillRate the rate to refill tokens.
|
*/
|
function TokenBucket(options) {
|
assert.object(options, 'options');
|
assert.number(options.capacity, 'options.capacity');
|
assert.number(options.fillRate, 'options.fillRate');
|
|
this.tokens = this.capacity = options.capacity;
|
this.fillRate = options.fillRate;
|
this.time = Date.now();
|
}
|
|
/**
|
* Consume N tokens from the bucket.
|
*
|
* If there is not capacity, the tokens are not pulled from the bucket.
|
*
|
* @private
|
* @memberof TokenBucket
|
* @instance
|
* @function consume
|
* @param {Number} tokens - the number of tokens to pull out.
|
* @returns {Boolean} true if capacity, false otherwise.
|
*/
|
TokenBucket.prototype.consume = function consume(tokens) {
|
if (tokens <= this._fill()) {
|
this.tokens -= tokens;
|
return true;
|
}
|
|
return false;
|
};
|
|
/**
|
* Fills the bucket with more tokens.
|
*
|
* Rather than do some whacky setTimeout() deal, we just approximate refilling
|
* the bucket by tracking elapsed time from the last time we touched the bucket.
|
*
|
* Simply, we set the bucket size to min(totalTokens,
|
* current + (fillRate * elapsed time)).
|
*
|
* @private
|
* @memberof TokenBucket
|
* @instance
|
* @function _fill
|
* @returns {Number} the current number of tokens in the bucket.
|
*/
|
TokenBucket.prototype._fill = function _fill() {
|
var now = Date.now();
|
|
// reset account for clock drift (like DST)
|
if (now < this.time) {
|
this.time = now - 1000;
|
}
|
|
if (this.tokens < this.capacity) {
|
var delta = this.fillRate * ((now - this.time) / 1000);
|
this.tokens = Math.min(this.capacity, this.tokens + delta);
|
}
|
this.time = now;
|
|
return this.tokens;
|
};
|
|
///--- Internal Class (TokenTable)
|
/**
|
* Just a wrapper over LRU that supports put/get to store token -> bucket
|
* mappings.
|
*
|
* @private
|
* @class TokenTable
|
* @param {Object} options - an options object
|
* @param {Number} options.size - size of the LRU
|
*/
|
function TokenTable(options) {
|
assert.object(options, 'options');
|
|
this.table = new LRU(options.size || 10000);
|
}
|
|
/**
|
* Puts a value in the token table
|
*
|
* @private
|
* @memberof TokenTable
|
* @instance
|
* @function put
|
* @param {String} key - a name
|
* @param {TokenBucket} value - a TokenBucket
|
* @returns {undefined} no return value
|
*/
|
TokenTable.prototype.put = function put(key, value) {
|
this.table.set(key, value);
|
};
|
|
/**
|
* Puts a value in the token table
|
*
|
* @private
|
* @memberof TokenTable
|
* @instance
|
* @function get
|
* @param {String} key - a key
|
* @returns {TokenBucket} token bucket instance
|
*/
|
TokenTable.prototype.get = function get(key) {
|
return this.table.get(key);
|
};
|
|
///--- Exported API
|
|
/**
|
* Creates an API rate limiter that can be plugged into the standard
|
* restify request handling pipeline.
|
*
|
* `restify` ships with a fairly comprehensive implementation of
|
* [Token bucket](http://en.wikipedia.org/wiki/Token_bucket), with the ability
|
* to throttle on IP (or x-forwarded-for) and username (from `req.username`).
|
* You define "global" request rate and burst rate, and you can define
|
* overrides for specific keys.
|
* Note that you can always place this on per-URL routes to enable
|
* different request rates to different resources (if for example, one route,
|
* like `/my/slow/database` is much easier to overwhlem
|
* than `/my/fast/memcache`).
|
*
|
* If a client has consumed all of their available rate/burst, an HTTP response
|
* code of `429`
|
* [Too Many Requests]
|
* (http://tools.ietf.org/html/draft-nottingham-http-new-status-03#section-4)
|
* is returned.
|
*
|
* This throttle gives you three options on which to throttle:
|
* username, IP address and 'X-Forwarded-For'. IP/XFF is a /32 match,
|
* so keep that in mind if using it. Username takes the user specified
|
* on req.username (which gets automagically set for supported Authorization
|
* types; otherwise set it yourself with a filter that runs before this).
|
*
|
* In both cases, you can set a `burst` and a `rate` (in requests/seconds),
|
* as an integer/float. Those really translate to the `TokenBucket`
|
* algorithm, so read up on that (or see the comments above...).
|
*
|
* In either case, the top level options burst/rate set a blanket throttling
|
* rate, and then you can pass in an `overrides` object with rates for
|
* specific users/IPs. You should use overrides sparingly, as we make a new
|
* TokenBucket to track each.
|
*
|
* On the `options` object ip and username are treated as an XOR.
|
*
|
* @public
|
* @function throttle
|
* @throws {TooManyRequestsError}
|
* @param {Object} options - required options with:
|
* @param {Number} options.burst - burst
|
* @param {Number} options.rate - rate
|
* @param {Boolean} [options.ip] - ip
|
* @param {Boolean} [options.username] - username
|
* @param {Boolean} [options.xff] - xff
|
* @param {Boolean} [options.setHeaders=false] - Set response headers for rate,
|
* limit (burst) and remaining.
|
* @param {Object} [options.overrides] - overrides
|
* @param {Object} options.tokensTable - a storage engine this plugin will
|
* use to store throttling keys -> bucket mappings.
|
* If you don't specify this, the default is to
|
* use an in-memory O(1) LRU, with 10k distinct
|
* keys. Any implementation just needs to support
|
* put/get.
|
* @param {Number} [options.maxKeys=10000] - If using the default
|
* implementation, you can specify how large you
|
* want the table to be.
|
* @returns {Function} Handler
|
* @example
|
* <caption>
|
* An example options object with overrides:
|
* </caption>
|
* {
|
* burst: 10, // Max 10 concurrent requests (if tokens)
|
* rate: 0.5, // Steady state: 1 request / 2 seconds
|
* ip: true, // throttle per IP
|
* overrides: {
|
* '192.168.1.1': {
|
* burst: 0,
|
* rate: 0 // unlimited
|
* }
|
* }
|
*/
|
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)');
|
}
|
|
var table =
|
options.tokensTable || new TokenTable({ size: options.maxKeys });
|
|
function rateLimit(req, res, next) {
|
var attr;
|
var burst = options.burst;
|
var rate = options.rate;
|
|
if (options.ip) {
|
attr = req.connection.remoteAddress;
|
} else if (options.xff) {
|
attr = req.headers['x-forwarded-for'];
|
} else if (options.username) {
|
attr = req.username;
|
} else {
|
req.log.warn({ config: options }, 'Invalid throttle configuration');
|
return next();
|
}
|
|
// Before bothering with overrides, see if this request
|
// even matches
|
if (!attr) {
|
return next();
|
}
|
|
// Check the overrides
|
if (
|
options.overrides &&
|
options.overrides[attr] &&
|
options.overrides[attr].burst !== undefined &&
|
options.overrides[attr].rate !== undefined
|
) {
|
burst = options.overrides[attr].burst;
|
rate = options.overrides[attr].rate;
|
}
|
|
if (!rate || !burst) {
|
return next();
|
}
|
|
var bucket = table.get(attr);
|
|
if (!bucket) {
|
bucket = new TokenBucket({
|
capacity: burst,
|
fillRate: rate
|
});
|
table.put(attr, bucket);
|
}
|
|
req.log.trace('Throttle(%s): num_tokens= %d', attr, bucket.tokens);
|
|
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,
|
url: req.url,
|
user: req.username || '?'
|
},
|
'Throttling'
|
);
|
|
var msg = sprintf(MESSAGE, rate);
|
return next(new TooManyRequestsError(msg));
|
}
|
|
return next();
|
}
|
|
return rateLimit;
|
}
|
|
module.exports = throttle;
|
