koa2-ratelimit/src/RedisStore.js at 7cc8cca6c7123c17f16d992943e3440bc1a0a152 · jmacpherson/koa2-ratelimit · GitHub

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
/**
 * RedisStore
 *
 * RedisStore for koa2-ratelimit
 *
 * @author Ashok Vishwakarma <akvlko@gmail.com>
 */

/**
 * Store
 *
 * Existing Store class
 */
const Store = require('./Store.js');

/**
 * redis
 *
 * node-redis module
 */
const redis = require('redis');

/**
 * RedisStore
 *
 * Class RedisStore
 */
class RedisStore extends Store {
  /**
   * constructor
   * @param {*} config
   *
   * config is redis config
   */
  constructor(config){
    super();
    this.client = redis.createClient(config);
    this.client.on('error', (err) => console.log('Redis Client Error', err));
    this.client.connect()


  }

  /**
   * _hit
   * @access private
   * @param {*} key
   * @param {*} options
   * @param {*} weight
   */
  async _hit(key, options, weight) {

    let [counter, dateEnd] = await this.client.multi().get(key).ttl(key).exec();

    if(counter === null) {
      counter = weight;
      dateEnd = Date.now() + options.interval;

      const seconds = Math.ceil(options.interval / 1000);
      await this.client.setEx(key, seconds.toString(), counter.toString());
    } else if (dateEnd === -2 || dateEnd === -1) {
      counter = counter + weight;
      dateEnd = Date.now() + options.interval;

      const seconds = Math.ceil(options.interval / 1000);
      await this.client.setEx(key, seconds.toString(), counter.toString());
    } else {
      counter = await this.client.incrBy(key, weight);
    }

    return {
      counter,
      dateEnd
    }
  }

  /**
   * incr
   *
   * Override incr method from Store class
   * @param {*} key
   * @param {*} options
   * @param {*} weight
   */
  async incr(key, options, weight) {
    return await this._hit(key, options, weight);
  }

  /**
   * decrement
   *
   * Override decrement method from Store class
   * @param {*} key
   * @param {*} options
   * @param {*} weight
   */
  async decrement(key, options, weight) {
    await this.client.decrBy(key, weight);
  }

  /**
   * saveAbuse
   *
   * Override saveAbuse method from Store class
   */
  saveAbuse() {}
}

module.exports = RedisStore;