relu-core
Version:
301 lines (250 loc) • 8.78 kB
JavaScript
// 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('../errors');
///--- Globals
var TooManyRequestsError = errors.TooManyRequestsError;
var MESSAGE = 'You have exceeded your request rate of %s r/s.';
///--- Helpers
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
* @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
* @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
* @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
* @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
* @function put
* @param {String} key a name
* @param {TokenBucket} value a TokenBucket
* @returns {undefined}
*/
TokenTable.prototype.put = function put(key, value) {
this.table.set(key, value);
};
/**
* puts a value in the token table
* @private
* @function get
* @param {String} key a key
* @returns {TokenBucket}
*/
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.
*
* 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.
*
* An example options object with overrides:
*
* {
* 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
* }
* }
*
* @public
* @function throttle
* @throws {TooManyRequestsError}
* @param {Object} options required options with:
* - {Number} burst (required).
* - {Number} rate (required).
* - {Boolean} ip (optional).
* - {Boolean} username (optional).
* - {Boolean} xff (optional).
* - {Object} overrides (optional).
* - {Object} 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.
* - {Number} maxKeys: If using the default implementation,
* you can specify how large you want the table to
* be. Default is 10000.
* @returns {Function}
*/
function throttle(options) {
assert.object(options, 'options');
assert.number(options.burst, 'options.burst');
assert.number(options.rate, 'options.rate');
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);
if (!bucket.consume(1)) {
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;