/
throttle.js
311 lines (260 loc) · 9.13 KB
/
throttle.js
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
// 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
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).
* - {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.
* 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');
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;