index.mjs 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175
  1. const scripts = {
  2. increment: `
  3. local totalHits = redis.call("INCR", KEYS[1])
  4. local timeToExpire = redis.call("PTTL", KEYS[1])
  5. if timeToExpire <= 0 or ARGV[1] == "1"
  6. then
  7. redis.call("PEXPIRE", KEYS[1], tonumber(ARGV[2]))
  8. timeToExpire = tonumber(ARGV[2])
  9. end
  10. return { totalHits, timeToExpire }
  11. `.replaceAll(/^\s+/gm, "").trim(),
  12. get: `
  13. local totalHits = redis.call("GET", KEYS[1])
  14. local timeToExpire = redis.call("PTTL", KEYS[1])
  15. return { totalHits, timeToExpire }
  16. `.replaceAll(/^\s+/gm, "").trim()
  17. };
  18. const toInt = (input) => {
  19. if (typeof input === "number")
  20. return input;
  21. return Number.parseInt((input ?? "").toString(), 10);
  22. };
  23. const parseScriptResponse = (results) => {
  24. if (!Array.isArray(results))
  25. throw new TypeError("Expected result to be array of values");
  26. if (results.length !== 2)
  27. throw new Error(`Expected 2 replies, got ${results.length}`);
  28. const totalHits = results[0] === false ? 0 : toInt(results[0]);
  29. const timeToExpire = toInt(results[1]);
  30. const resetTime = new Date(Date.now() + timeToExpire);
  31. return { totalHits, resetTime };
  32. };
  33. class RedisStore {
  34. /**
  35. * The function used to send raw commands to Redis.
  36. */
  37. sendCommand;
  38. /**
  39. * The text to prepend to the key in Redis.
  40. */
  41. prefix;
  42. /**
  43. * Whether to reset the expiry for a particular key whenever its hit count
  44. * changes.
  45. */
  46. resetExpiryOnChange;
  47. /**
  48. * Stores the loaded SHA1s of the LUA scripts used for executing the increment
  49. * and get key operations.
  50. */
  51. incrementScriptSha;
  52. getScriptSha;
  53. /**
  54. * The number of milliseconds to remember that user's requests.
  55. */
  56. windowMs;
  57. /**
  58. * @constructor for `RedisStore`.
  59. *
  60. * @param options {Options} - The configuration options for the store.
  61. */
  62. constructor(options) {
  63. this.sendCommand = options.sendCommand;
  64. this.prefix = options.prefix ?? "rl:";
  65. this.resetExpiryOnChange = options.resetExpiryOnChange ?? false;
  66. this.incrementScriptSha = this.loadIncrementScript();
  67. this.getScriptSha = this.loadGetScript();
  68. }
  69. /**
  70. * Loads the script used to increment a client's hit count.
  71. */
  72. async loadIncrementScript() {
  73. const result = await this.sendCommand("SCRIPT", "LOAD", scripts.increment);
  74. if (typeof result !== "string") {
  75. throw new TypeError("unexpected reply from redis client");
  76. }
  77. return result;
  78. }
  79. /**
  80. * Loads the script used to fetch a client's hit count and expiry time.
  81. */
  82. async loadGetScript() {
  83. const result = await this.sendCommand("SCRIPT", "LOAD", scripts.get);
  84. if (typeof result !== "string") {
  85. throw new TypeError("unexpected reply from redis client");
  86. }
  87. return result;
  88. }
  89. /**
  90. * Runs the increment command, and retries it if the script is not loaded.
  91. */
  92. async retryableIncrement(key) {
  93. const evalCommand = async () => this.sendCommand(
  94. "EVALSHA",
  95. await this.incrementScriptSha,
  96. "1",
  97. this.prefixKey(key),
  98. this.resetExpiryOnChange ? "1" : "0",
  99. this.windowMs.toString()
  100. );
  101. try {
  102. const result = await evalCommand();
  103. return result;
  104. } catch {
  105. this.incrementScriptSha = this.loadIncrementScript();
  106. return evalCommand();
  107. }
  108. }
  109. /**
  110. * Method to prefix the keys with the given text.
  111. *
  112. * @param key {string} - The key.
  113. *
  114. * @returns {string} - The text + the key.
  115. */
  116. prefixKey(key) {
  117. return `${this.prefix}${key}`;
  118. }
  119. /**
  120. * Method that actually initializes the store.
  121. *
  122. * @param options {RateLimitConfiguration} - The options used to setup the middleware.
  123. */
  124. init(options) {
  125. this.windowMs = options.windowMs;
  126. }
  127. /**
  128. * Method to fetch a client's hit count and reset time.
  129. *
  130. * @param key {string} - The identifier for a client.
  131. *
  132. * @returns {ClientRateLimitInfo | undefined} - The number of hits and reset time for that client.
  133. */
  134. async get(key) {
  135. const results = await this.sendCommand(
  136. "EVALSHA",
  137. await this.getScriptSha,
  138. "1",
  139. this.prefixKey(key)
  140. );
  141. return parseScriptResponse(results);
  142. }
  143. /**
  144. * Method to increment a client's hit counter.
  145. *
  146. * @param key {string} - The identifier for a client
  147. *
  148. * @returns {IncrementResponse} - The number of hits and reset time for that client
  149. */
  150. async increment(key) {
  151. const results = await this.retryableIncrement(key);
  152. return parseScriptResponse(results);
  153. }
  154. /**
  155. * Method to decrement a client's hit counter.
  156. *
  157. * @param key {string} - The identifier for a client
  158. */
  159. async decrement(key) {
  160. await this.sendCommand("DECR", this.prefixKey(key));
  161. }
  162. /**
  163. * Method to reset a client's hit counter.
  164. *
  165. * @param key {string} - The identifier for a client
  166. */
  167. async resetKey(key) {
  168. await this.sendCommand("DEL", this.prefixKey(key));
  169. }
  170. }
  171. export { RedisStore, RedisStore as default };