-
-
Notifications
You must be signed in to change notification settings - Fork 167
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Added: RateLimiterMongo, tests, add attrs to RateLimiterRes
- Loading branch information
Showing
15 changed files
with
557 additions
and
52 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,10 +1,12 @@ | ||
| const RateLimiterRedis = require('./lib/RateLimiterRedis'); | ||
| const RateLimiterMongo = require('./lib/RateLimiterMongo'); | ||
| const { RateLimiterClusterMaster, RateLimiterCluster } = require('./lib/RateLimiterCluster'); | ||
| const RateLimiterMemory = require('./lib/RateLimiterMemory'); | ||
|
|
||
| module.exports = { | ||
| RateLimiterRedis, | ||
| RateLimiterMongo, | ||
| RateLimiterMemory, | ||
| RateLimiterClusterMaster, | ||
| RateLimiterCluster, | ||
| }; | ||
| }; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,135 @@ | ||
| const RateLimiterStoreAbstract = require('./RateLimiterStoreAbstract'); | ||
| const RateLimiterRes = require('./RateLimiterRes'); | ||
|
|
||
| const getRateLimiterRes = function (points, result) { | ||
| const res = new RateLimiterRes(); | ||
|
|
||
| res.isFirstInDuration = result.value === null; | ||
| res.consumedPoints = res.isFirstInDuration ? points : result.value.points; | ||
|
|
||
| res.remainingPoints = Math.max(this.points - res.consumedPoints, 0); | ||
| res.msBeforeNext = res.isFirstInDuration | ||
| ? this.duration * 1000 | ||
| : Math.max(new Date(result.value.expire).getTime() - Date.now(), 0); | ||
|
|
||
| return res; | ||
| }; | ||
|
|
||
| const afterConsume = function (resolve, reject, rlKey, points, result) { | ||
| const res = getRateLimiterRes.call(this, points, result); | ||
|
|
||
| if (res.consumedPoints > this.points) { | ||
| // Block key for this.blockDuration seconds | ||
| if (this.blockOnPointsConsumed > 0 && res.consumedPoints >= this.blockOnPointsConsumed) { | ||
| this._blockedKeys.add(rlKey, this.blockDuration); | ||
| res.msBeforeNext = this.msBlockDuration; | ||
| } | ||
|
|
||
| reject(res); | ||
| } else if (this.execEvenly && res.msBeforeNext > 0 && !res.isFirstInDuration) { | ||
| const delay = Math.ceil(res.msBeforeNext / (res.remainingPoints + 2)); | ||
| setTimeout(resolve, delay, res); | ||
| } else { | ||
| resolve(res); | ||
| } | ||
| }; | ||
|
|
||
| const update = function (key, points) { | ||
| return this._collection.findOneAndUpdate( | ||
| { | ||
| expire: { $gt: new Date() }, | ||
| key, | ||
| }, | ||
| { | ||
| $inc: { points }, | ||
| $setOnInsert: { expire: new Date(Date.now() + (this.duration * 1000)) }, | ||
| }, | ||
| { | ||
| upsert: true, | ||
| returnNewDocument: true, | ||
| } // eslint-disable-line comma-dangle | ||
| ); | ||
| }; | ||
|
|
||
| class RateLimiterMongo extends RateLimiterStoreAbstract { | ||
| /** | ||
| * | ||
| * @param {Object} opts | ||
| * Defaults { | ||
| * ... see other in RateLimiterStoreAbstract | ||
| * | ||
| * mongo: MongoClient | ||
| * } | ||
| */ | ||
| constructor(opts) { | ||
| super(opts); | ||
|
|
||
| this.mongo = opts.mongo; | ||
| this._collection = this.mongo.db('node-rate-limiter-flexible').collection(this.keyPrefix); | ||
| this._collection.ensureIndex({ expire: -1 }, { expireAfterSeconds: 0 }); | ||
| } | ||
|
|
||
| get mongo() { | ||
| return this._mongo; | ||
| } | ||
|
|
||
| set mongo(value) { | ||
| if (typeof value === 'undefined') { | ||
| throw new Error('mongo is not set'); | ||
| } | ||
| this._mongo = value; | ||
| } | ||
|
|
||
| /** | ||
| * | ||
| * @param key | ||
| * @param pointsToConsume | ||
| * @returns {Promise<any>} | ||
| */ | ||
| consume(key, pointsToConsume = 1) { | ||
| return new Promise((resolve, reject) => { | ||
| const rlKey = this.getKey(key); | ||
|
|
||
| const blockMsBeforeExpire = this.getBlockMsBeforeExpire(rlKey); | ||
| if (blockMsBeforeExpire > 0) { | ||
| return reject(new RateLimiterRes(0, blockMsBeforeExpire)); | ||
| } | ||
|
|
||
| update.call(this, rlKey, pointsToConsume) | ||
| .then((res) => { | ||
| afterConsume.call(this, resolve, reject, rlKey, pointsToConsume, res); | ||
| }) | ||
| .catch((err) => { | ||
| this.handleError(err, 'consume', resolve, reject, key, pointsToConsume); | ||
| }); | ||
| }); | ||
| } | ||
|
|
||
| penalty(key, points = 1) { | ||
| const rlKey = this.getKey(key); | ||
| return new Promise((resolve, reject) => { | ||
| update.call(this, rlKey, points) | ||
| .then((res) => { | ||
| resolve(getRateLimiterRes.call(this, points, res)); | ||
| }) | ||
| .catch((err) => { | ||
| this.handleError(err, 'penalty', resolve, reject, key, points); | ||
| }); | ||
| }); | ||
| } | ||
|
|
||
| reward(key, points = 1) { | ||
| const rlKey = this.getKey(key); | ||
| return new Promise((resolve, reject) => { | ||
| update.call(this, rlKey, -points) | ||
| .then((res) => { | ||
| resolve(getRateLimiterRes.call(this, points, res)); | ||
| }) | ||
| .catch((err) => { | ||
| this.handleError(err, 'reward', resolve, reject, key, points); | ||
| }); | ||
| }); | ||
| } | ||
| } | ||
|
|
||
| module.exports = RateLimiterMongo; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,82 @@ | ||
| const RateLimiterAbstract = require('./RateLimiterAbstract'); | ||
| const BlockedKeys = require('./component/BlockedKeys'); | ||
|
|
||
| module.exports = class RateLimiterStoreAbstract extends RateLimiterAbstract { | ||
| /** | ||
| * | ||
| * @param opts Object Defaults { | ||
| * ... see other in RateLimiterAbstract | ||
| * | ||
| * blockOnPointsConsumed: 40, // Number of points when key is blocked | ||
| * blockDuration: 10, // Block duration in seconds | ||
| * insuranceLimiter: RateLimiterAbstract | ||
| * } | ||
| */ | ||
| constructor(opts = {}) { | ||
| super(opts); | ||
|
|
||
| this.blockOnPointsConsumed = opts.blockOnPointsConsumed; | ||
| this.blockDuration = opts.blockDuration; | ||
| this.insuranceLimiter = opts.insuranceLimiter; | ||
| this._blockedKeys = new BlockedKeys(); | ||
| } | ||
|
|
||
| getBlockMsBeforeExpire(rlKey) { | ||
| if (this.blockOnPointsConsumed > 0) { | ||
| return this._blockedKeys.msBeforeExpire(rlKey); | ||
| } | ||
|
|
||
| return 0; | ||
| } | ||
|
|
||
| handleError(err, funcName, resolve, reject, key, points) { | ||
| if (!(this.insuranceLimiter instanceof RateLimiterAbstract)) { | ||
| reject(err); | ||
| } else { | ||
| this.insuranceLimiter[funcName](key, points) | ||
| .then((res) => { | ||
| resolve(res); | ||
| }) | ||
| .catch((res) => { | ||
| reject(res); | ||
| }); | ||
| } | ||
| } | ||
|
|
||
| get blockOnPointsConsumed() { | ||
| return this._blockOnPointsConsumed; | ||
| } | ||
|
|
||
| set blockOnPointsConsumed(value) { | ||
| this._blockOnPointsConsumed = value ? parseInt(value) : 0; | ||
| if (this.blockOnPointsConsumed > 0 && this.points >= this.blockOnPointsConsumed) { | ||
| throw new Error('blockOnPointsConsumed option must be more than points option'); | ||
| } | ||
| } | ||
|
|
||
| get blockDuration() { | ||
| return this._blockDuration; | ||
| } | ||
|
|
||
| get msBlockDuration() { | ||
| return this._blockDuration * 1000; | ||
| } | ||
|
|
||
| set blockDuration(value) { | ||
| this._blockDuration = value ? parseInt(value) : 0; | ||
| if (this.blockDuration > 0 && this.blockOnPointsConsumed === 0) { | ||
| throw new Error('blockOnPointsConsumed option must be set up'); | ||
| } | ||
| } | ||
|
|
||
| get insuranceLimiter() { | ||
| return this._insuranceLimiter; | ||
| } | ||
|
|
||
| set insuranceLimiter(value) { | ||
| if (typeof value !== 'undefined' && !(value instanceof RateLimiterAbstract)) { | ||
| throw new Error('insuranceLimiter must be instance of RateLimiterAbstract'); | ||
| } | ||
| this._insuranceLimiter = value; | ||
| } | ||
| }; |
Oops, something went wrong.