ratelimit.js/lib/rate_limit.js

A NodeJS library for efficient rate limiting using sliding windows stored in Redis.

// Generated by CoffeeScript 1.8.0
var EvalSha, RateLimit, async, fs, _,
  __indexOf = [].indexOf || function(item) { for (var i = 0, l = this.length; i < l; i++) { if (i in this && this[i] === item) return i; } return -1; };

EvalSha = require('redis-evalsha');

_ = require('underscore');

async = require('async');

fs = require('fs');

module.exports = RateLimit = (function() {
  RateLimit.DEFAULT_PREFIX = 'ratelimit';

  RateLimit.DENIED_NUMS = [1, 2];

  function RateLimit(redisClient, rules, options) {
    var _base, _base1;
    this.redisClient = redisClient;
    this.options = options != null ? options : {};
    if (_.isString(this.options)) {
      this.options = {
        prefix: this.options
      };
    }
    if ((_base = this.options).prefix == null) {
      _base.prefix = this.constructor.DEFAULT_PREFIX;
    }
    if ((_base1 = this.options).clientPrefix == null) {
      _base1.clientPrefix = false;
    }
    this.checkFn = 'check_rate_limit';
    this.checkIncrFn = 'check_incr_rate_limit';
    this["eval"] = new EvalSha(this.redisClient);
    this["eval"].add(this.checkFn, this.checkLimitScript());
    this["eval"].add(this.checkIncrFn, this.checkLimitIncrScript());
    this.rules = this.convertRules(rules);
  }

  RateLimit.prototype.readLua = function(filename) {
    return fs.readFileSync("" + __dirname + "/../lua/" + filename + ".lua", 'utf8');
  };

  RateLimit.prototype.checkLimitScript = function() {
    return [this.readLua('unpack_args'), this.readLua('check_whitelist_blacklist'), this.readLua('check_limit'), 'return 0'].join('\n');
  };

  RateLimit.prototype.checkLimitIncrScript = function() {
    return [this.readLua('unpack_args'), this.readLua('check_whitelist_blacklist'), this.readLua('check_limit'), this.readLua('check_incr_limit')].join('\n');
  };

  RateLimit.prototype.convertRules = function(rules) {
    var interval, limit, precision, _i, _len, _ref, _results;
    _results = [];
    for (_i = 0, _len = rules.length; _i < _len; _i++) {
      _ref = rules[_i], interval = _ref.interval, limit = _ref.limit, precision = _ref.precision;
      if (interval && limit) {
        _results.push(_.compact([interval, limit, precision]));
      }
    }
    return _results;
  };

  RateLimit.prototype.prefixKey = function(key, force) {
    var parts;
    if (force == null) {
      force = false;
    }
    parts = [key];
    if (force || !this.options.clientPrefix) {
      parts.unshift(this.options.prefix);
    }
    return _.compact(parts).join(':');
  };

  RateLimit.prototype.whitelistKey = function() {
    return this.prefixKey('whitelist', true);
  };

  RateLimit.prototype.blacklistKey = function() {
    return this.prefixKey('blacklist', true);
  };

  RateLimit.prototype.scriptArgs = function(keys, weight) {
    var adjustedKeys, rules, ts;
    if (weight == null) {
      weight = 1;
    }
    adjustedKeys = _.chain([keys]).flatten().compact().filter(function(key) {
      return _.isString(key) && key.length;
    }).map((function(_this) {
      return function(key) {
        return _this.prefixKey(key);
      };
    })(this)).value();
    if (!adjustedKeys.length) {
      throw new Error("Bad keys: " + keys);
    }
    rules = JSON.stringify(this.rules);
    ts = Math.floor(Date.now() / 1000);
    weight = Math.max(weight, 1);
    return [adjustedKeys, [rules, ts, weight, this.whitelistKey(), this.blacklistKey()]];
  };

  RateLimit.prototype.check = function(keys, callback) {
    var args, err, _ref;
    try {
      _ref = this.scriptArgs(keys), keys = _ref[0], args = _ref[1];
    } catch (_error) {
      err = _error;
      return callback(err);
    }
    return this["eval"].exec(this.checkFn, keys, args, (function(_this) {
      return function(err, result) {
        return callback(err, __indexOf.call(_this.constructor.DENIED_NUMS, result) >= 0);
      };
    })(this));
  };

  RateLimit.prototype.incr = function(keys, weight, callback) {
    var args, err, _ref, _ref1;
    if (arguments.length === 2) {
      _ref = [1, weight], weight = _ref[0], callback = _ref[1];
    }
    try {
      _ref1 = this.scriptArgs(keys, weight), keys = _ref1[0], args = _ref1[1];
    } catch (_error) {
      err = _error;
      return callback(err);
    }
    return this["eval"].exec(this.checkIncrFn, keys, args, (function(_this) {
      return function(err, result) {
        return callback(err, __indexOf.call(_this.constructor.DENIED_NUMS, result) >= 0);
      };
    })(this));
  };

  RateLimit.prototype.keys = function(callback) {
    return this.redisClient.keys(this.prefixKey('*'), (function(_this) {
      return function(err, results) {
        var key, keys, re;
        if (err) {
          return callback(err);
        }
        re = new RegExp(_this.prefixKey('(.+)'));
        keys = (function() {
          var _i, _len, _results;
          _results = [];
          for (_i = 0, _len = results.length; _i < _len; _i++) {
            key = results[_i];
            _results.push(re.exec(key)[1]);
          }
          return _results;
        })();
        return callback(null, keys);
      };
    })(this));
  };

  RateLimit.prototype.violatedRules = function(keys, callback) {
    var checkKey;
    checkKey = (function(_this) {
      return function(key, callback) {
        var checkRule;
        checkRule = function(rule, callback) {
          var countKey, interval, limit, precision;
          interval = rule[0], limit = rule[1], precision = rule[2];
          precision = Math.min(precision != null ? precision : interval, interval);
          countKey = "" + interval + ":" + precision + ":";
          return _this.redisClient.hget(_this.prefixKey(key), countKey, function(err, count) {
            if (count == null) {
              count = -1;
            }
            if (!(count >= limit)) {
              return callback();
            }
            return callback(null, {
              interval: interval,
              limit: limit
            });
          });
        };
        return async.map(_this.rules, checkRule, function(err, violatedRules) {
          return callback(err, _.compact(violatedRules));
        });
      };
    })(this);
    return async.concat(_.flatten([keys]), checkKey, callback);
  };

  RateLimit.prototype.limitedKeys = function(callback) {
    return this.keys((function(_this) {
      return function(err, keys) {
        var fn;
        if (err) {
          return callback(err);
        }
        fn = function(key, callback) {
          return _this.check(key, function(err, limited) {
            return callback(limited);
          });
        };
        return async.filter(keys, fn, function(results) {
          return callback(null, results);
        });
      };
    })(this));
  };

  RateLimit.prototype.whitelist = function(keys, callback) {
    var whitelist;
    whitelist = (function(_this) {
      return function(key, callback) {
        key = _this.prefixKey(key);
        return async.series([
          function(callback) {
            return _this.redisClient.srem(_this.blacklistKey(), key, callback);
          }, function(callback) {
            return _this.redisClient.sadd(_this.whitelistKey(), key, callback);
          }
        ], callback);
      };
    })(this);
    return async.each(keys, whitelist, callback);
  };

  RateLimit.prototype.unwhitelist = function(keys, callback) {
    var unwhitelist;
    unwhitelist = (function(_this) {
      return function(key, callback) {
        key = _this.prefixKey(key);
        return _this.redisClient.srem(_this.whitelistKey(), key, callback);
      };
    })(this);
    return async.each(keys, unwhitelist, callback);
  };

  RateLimit.prototype.blacklist = function(keys, callback) {
    var blacklist;
    blacklist = (function(_this) {
      return function(key, callback) {
        key = _this.prefixKey(key);
        return async.series([
          function(callback) {
            return _this.redisClient.srem(_this.whitelistKey(), key, callback);
          }, function(callback) {
            return _this.redisClient.sadd(_this.blacklistKey(), key, callback);
          }
        ], callback);
      };
    })(this);
    return async.each(keys, blacklist, callback);
  };

  RateLimit.prototype.unblacklist = function(keys, callback) {
    var unblacklist;
    unblacklist = (function(_this) {
      return function(key, callback) {
        key = _this.prefixKey(key);
        return _this.redisClient.srem(_this.blacklistKey(), key, callback);
      };
    })(this);
    return async.each(keys, unblacklist, callback);
  };

  return RateLimit;

})();