kopia lustrzana https://github.com/c9/core
Switch rate limited to use one setInterval instead of many setTimeouts to clean itself up.
rodzic
07e7c55b44
commit
6c0737c3c0
|
@ -1,4 +1,5 @@
|
|||
var error = require("http-error");
|
||||
var _ = require("lodash");
|
||||
|
||||
/**
|
||||
* In memory rate limiter as connect middleware
|
||||
|
@ -6,25 +7,34 @@ var error = require("http-error");
|
|||
module.exports = ratelimit;
|
||||
|
||||
function ratelimit(key, duration, max) {
|
||||
var counts = {};
|
||||
var requests = {};
|
||||
setInterval(function() {
|
||||
_.forIn(requests, function (requestsForHandle, key) {
|
||||
var totalToSplice = 0;
|
||||
var expireTime = Date.now() - duration;
|
||||
/* Requests are already sorted by date as they are appended, so we just loop
|
||||
until we find one that shouldn't have expired and splice them from the list */
|
||||
for (var i = 0; i < requestsForHandle.length; i++) {
|
||||
if (requestsForHandle[i] >= expireTime) break;
|
||||
totalToSplice++;
|
||||
}
|
||||
requests[key].splice(0, totalToSplice);
|
||||
if (requests[key].length == 0) {
|
||||
delete requests[key];
|
||||
}
|
||||
});
|
||||
}, 10);
|
||||
return function(req, res, next) {
|
||||
var handle = req.params[key];
|
||||
|
||||
counts[handle] = counts[handle] || 0;
|
||||
if (counts[handle] >= max) {
|
||||
requests[handle] = requests[handle] || [];
|
||||
if (requests[handle].length >= max) {
|
||||
var err = new error.TooManyRequests("Rate limit exceeded");
|
||||
err.retryIn = Math.min(duration, 5000);
|
||||
return next(err);
|
||||
}
|
||||
|
||||
counts[handle]++;
|
||||
setTimeout(function () {
|
||||
counts[handle]--;
|
||||
if (counts[handle] == 0) {
|
||||
delete counts[handle];
|
||||
}
|
||||
}, duration);
|
||||
|
||||
requests[handle].push(Date.now());
|
||||
return next();
|
||||
};
|
||||
}
|
||||
|
|
|
@ -4,13 +4,14 @@ require("c9/inline-mocha")(module);
|
|||
|
||||
var ratelimit = require("./ratelimit");
|
||||
var assert = require("assert");
|
||||
var async = require("async");
|
||||
|
||||
describe("ratelimit", function() {
|
||||
it("Should limit based on key", function (done) {
|
||||
var limiter = ratelimit("test", 10, 1);
|
||||
limiter({params: {test: 1}}, null, function (err) {
|
||||
var limiter = ratelimit("username", 10, 1);
|
||||
limiter({params: {username: "super"}}, null, function (err) {
|
||||
assert(!err, err);
|
||||
limiter({params: {test: 1}}, null, function (err) {
|
||||
limiter({params: {username: "super"}}, null, function (err) {
|
||||
assert(err);
|
||||
assert.equal(err.code, 429);
|
||||
done();
|
||||
|
@ -18,16 +19,39 @@ describe("ratelimit", function() {
|
|||
});
|
||||
});
|
||||
|
||||
it("Should work again after a delay", function (done) {
|
||||
var limiter = ratelimit("test", 10, 1);
|
||||
limiter({params: {test: 1}}, null, function (err) {
|
||||
it("Should work again after a delay", function (done) {
|
||||
var limiter = ratelimit("username", 10, 1);
|
||||
limiter({params: {username: "super"}}, null, function (err) {
|
||||
assert(!err, err);
|
||||
setTimeout(function() {
|
||||
limiter({params: {test: 1}}, null, function (err) {
|
||||
limiter({params: {username: "super"}}, null, function (err) {
|
||||
assert(!err, err);
|
||||
done();
|
||||
});
|
||||
}, 15);
|
||||
}, 25);
|
||||
});
|
||||
});
|
||||
|
||||
it("Should work with many requests", function (done) {
|
||||
var MAX_REQUESTS = 5;
|
||||
var limiter = ratelimit("username", 10, MAX_REQUESTS);
|
||||
var successfulRequests = 0;
|
||||
async.times(10, function(n, next) {
|
||||
limiter({params: {username: "super"}}, null, function (err) {
|
||||
if (err) return next(err);
|
||||
successfulRequests++;
|
||||
next();
|
||||
});
|
||||
}, function (err) {
|
||||
assert.equal(successfulRequests, MAX_REQUESTS);
|
||||
setTimeout(function() {
|
||||
limiter({params: {username: "super"}}, null, function (err) {
|
||||
assert(!err, err);
|
||||
done();
|
||||
});
|
||||
}, 25);
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
});
|
Ładowanie…
Reference in New Issue