@depup/rate-limiter-flexible 9.1.1-depup.0

package/lib/RateLimiterPostgres.js

@@ -0,0 +1,351 @@
+const RateLimiterStoreAbstract = require('./RateLimiterStoreAbstract');
+const RateLimiterRes = require('./RateLimiterRes');
+
+class RateLimiterPostgres extends RateLimiterStoreAbstract {
+  /**
+   * @callback callback
+   * @param {Object} err
+   *
+   * @param {Object} opts
+   * @param {callback} cb
+   * Defaults {
+   *   ... see other in RateLimiterStoreAbstract
+   *
+   *   storeClient: postgresClient,
+   *   storeType: 'knex', // required only for Knex instance
+   *   tableName: 'string',
+   *   schemaName: 'string', // optional
+   * }
+   */
+  constructor(opts, cb = null) {
+    super(opts);
+
+    this.client = opts.storeClient;
+    this.clientType = opts.storeType;
+
+    this.tableName = opts.tableName;
+    this.schemaName = opts.schemaName;
+
+    this.clearExpiredByTimeout = opts.clearExpiredByTimeout;
+
+    this.tableCreated = opts.tableCreated;
+    if (!this.tableCreated) {
+      this._createTable()
+        .then(() => {
+          this.tableCreated = true;
+          if (this.clearExpiredByTimeout) {
+            this._clearExpiredHourAgo();
+          }
+          if (typeof cb === 'function') {
+            cb();
+          }
+        })
+        .catch((err) => {
+          if (typeof cb === 'function') {
+            cb(err);
+          } else {
+            throw err;
+          }
+        });
+    } else {
+      if (this.clearExpiredByTimeout) {
+        this._clearExpiredHourAgo();
+      }
+      if (typeof cb === 'function') {
+        cb();
+      }
+    }
+  }
+
+  _getTableIdentifier() {
+    return this.schemaName ? `"${this.schemaName}"."${this.tableName}"` : `"${this.tableName}"`;
+  }
+
+  clearExpired(expire) {
+    return new Promise((resolve) => {
+      const q = {
+        name: 'rlflx-clear-expired',
+        text: `DELETE FROM ${this._getTableIdentifier()} WHERE expire < $1`,
+        values: [expire],
+      };
+      this._query(q)
+        .then(() => {
+          resolve();
+        })
+        .catch(() => {
+          // Deleting expired query is not critical
+          resolve();
+        });
+    });
+  }
+
+  /**
+   * Delete all rows expired 1 hour ago once per 5 minutes
+   *
+   * @private
+   */
+  _clearExpiredHourAgo() {
+    if (this._clearExpiredTimeoutId) {
+      clearTimeout(this._clearExpiredTimeoutId);
+    }
+    this._clearExpiredTimeoutId = setTimeout(() => {
+      this.clearExpired(Date.now() - 3600000) // Never rejected
+        .then(() => {
+          this._clearExpiredHourAgo();
+        });
+    }, 300000);
+    this._clearExpiredTimeoutId.unref();
+  }
+
+  /**
+   *
+   * @return Promise<any>
+   * @private
+   */
+  _getConnection() {
+    switch (this.clientType) {
+      case 'pool':
+        return Promise.resolve(this.client);
+      case 'sequelize':
+        return this._getSequelizeConnectionManager().getConnection();
+      case 'knex':
+        return this.client.client.acquireConnection();
+      case 'typeorm':
+        return Promise.resolve(this.client.driver.master);
+      default:
+        return Promise.resolve(this.client);
+    }
+  }
+
+  _releaseConnection(conn) {
+    switch (this.clientType) {
+      case 'pool':
+        return true;
+      case 'sequelize':
+        return this._getSequelizeConnectionManager().releaseConnection(conn);
+      case 'knex':
+        return this.client.client.releaseConnection(conn);
+      case 'typeorm':
+        return true;
+      default:
+        return true;
+    }
+  }
+
+  _getSequelizeConnectionManager() {
+    let connectionManager;
+    let accessError;
+    try {
+      connectionManager = this.client.connectionManager;
+    } catch (err) {
+      // Accessing connectionManager can throw in Sequelize version 7 and higher.
+      accessError = err;
+    }
+    if (connectionManager) {
+      return connectionManager;
+    }
+    if (this.client.dialect && this.client.dialect.connectionManager) {
+      return this.client.dialect.connectionManager;
+    }
+    // Rethrow the original error if it exists, otherwise throw a generic error
+    if (accessError) {
+      throw accessError;
+    }
+    throw new Error('Sequelize connection manager is not available');
+  }
+
+  /**
+   *
+   * @returns {Promise<any>}
+   * @private
+   */
+  _createTable() {
+    return new Promise((resolve, reject) => {
+      this._query({
+        text: this._getCreateTableStmt(),
+      })
+        .then(() => {
+          resolve();
+        })
+        .catch((err) => {
+          if (err.code === '23505') {
+            // Error: duplicate key value violates unique constraint "pg_type_typname_nsp_index"
+            // Postgres doesn't handle concurrent table creation
+            // It is supposed, that table is created by another worker
+            resolve();
+          } else {
+            reject(err);
+          }
+        });
+    });
+  }
+
+  _getCreateTableStmt() {
+    return `CREATE TABLE IF NOT EXISTS ${this._getTableIdentifier()} (
+      key varchar(255) PRIMARY KEY,
+      points integer NOT NULL DEFAULT 0,
+      expire bigint
+    );`;
+  }
+
+  get clientType() {
+    return this._clientType;
+  }
+
+  set clientType(value) {
+    const constructorName = this.client.constructor.name;
+
+    if (typeof value === 'undefined') {
+      if (constructorName === 'Client') {
+        value = 'client';
+      } else if (
+        constructorName === 'Pool' ||
+        constructorName === 'BoundPool'
+      ) {
+        value = 'pool';
+      } else if (constructorName === 'Sequelize') {
+        value = 'sequelize';
+      } else {
+        throw new Error('storeType is not defined');
+      }
+    }
+
+    this._clientType = value.toLowerCase();
+  }
+
+  get tableName() {
+    return this._tableName;
+  }
+
+  set tableName(value) {
+    this._tableName = typeof value === 'undefined' ? this.keyPrefix : value;
+  }
+
+  get schemaName() {
+    return this._schemaName;
+  }
+
+  set schemaName(value) {
+    this._schemaName = value;
+  }
+
+  get tableCreated() {
+    return this._tableCreated;
+  }
+
+  set tableCreated(value) {
+    this._tableCreated = typeof value === 'undefined' ? false : !!value;
+  }
+
+  get clearExpiredByTimeout() {
+    return this._clearExpiredByTimeout;
+  }
+
+  set clearExpiredByTimeout(value) {
+    this._clearExpiredByTimeout = typeof value === 'undefined' ? true : Boolean(value);
+  }
+
+  _getRateLimiterRes(rlKey, changedPoints, result) {
+    const res = new RateLimiterRes();
+    const row = result.rows[0];
+
+    res.isFirstInDuration = changedPoints === row.points;
+    res.consumedPoints = res.isFirstInDuration ? changedPoints : row.points;
+
+    res.remainingPoints = Math.max(this.points - res.consumedPoints, 0);
+    res.msBeforeNext = row.expire
+      ? Math.max(row.expire - Date.now(), 0)
+      : -1;
+
+    return res;
+  }
+
+  _query(q) {
+    const prefix = this.tableName.toLowerCase();
+    const queryObj = { name: `${prefix}:${q.name}`, text: q.text, values: q.values };
+    return new Promise((resolve, reject) => {
+      this._getConnection()
+        .then((conn) => {
+          conn.query(queryObj)
+            .then((res) => {
+              resolve(res);
+              this._releaseConnection(conn);
+            })
+            .catch((err) => {
+              reject(err);
+              this._releaseConnection(conn);
+            });
+        })
+        .catch((err) => {
+          reject(err);
+        });
+    });
+  }
+
+  _upsert(key, points, msDuration, forceExpire = false) {
+    if (!this.tableCreated) {
+      return Promise.reject(Error('Table is not created yet'));
+    }
+
+    const newExpire = msDuration > 0 ? Date.now() + msDuration : null;
+    const expireQ = forceExpire
+      ? ' $3 '
+      : ` CASE
+             WHEN ${this._getTableIdentifier()}.expire <= $4 THEN $3
+             ELSE ${this._getTableIdentifier()}.expire
+            END `;
+
+    return this._query({
+      name: forceExpire ? 'rlflx-upsert-force' : 'rlflx-upsert',
+      text: `
+            INSERT INTO ${this._getTableIdentifier()} VALUES ($1, $2, $3)
+              ON CONFLICT(key) DO UPDATE SET
+                points = CASE
+                          WHEN (${this._getTableIdentifier()}.expire <= $4 OR 1=${forceExpire ? 1 : 0}) THEN $2
+                          ELSE ${this._getTableIdentifier()}.points + ($2)
+                         END,
+                expire = ${expireQ}
+            RETURNING points, expire;`,
+      values: [key, points, newExpire, Date.now()],
+    });
+  }
+
+  _get(rlKey) {
+    if (!this.tableCreated) {
+      return Promise.reject(Error('Table is not created yet'));
+    }
+
+    return new Promise((resolve, reject) => {
+      this._query({
+        name: 'rlflx-get',
+        text: `
+            SELECT points, expire FROM ${this._getTableIdentifier()} WHERE key = $1 AND (expire > $2 OR expire IS NULL);`,
+        values: [rlKey, Date.now()],
+      })
+        .then((res) => {
+          if (res.rowCount === 0) {
+            res = null;
+          }
+          resolve(res);
+        })
+        .catch((err) => {
+          reject(err);
+        });
+    });
+  }
+
+  _delete(rlKey) {
+    if (!this.tableCreated) {
+      return Promise.reject(Error('Table is not created yet'));
+    }
+
+    return this._query({
+      name: 'rlflx-delete',
+      text: `DELETE FROM ${this._getTableIdentifier()} WHERE key = $1`,
+      values: [rlKey],
+    })
+      .then(res => res.rowCount > 0);
+  }
+}
+
+module.exports = RateLimiterPostgres;

package/lib/RateLimiterPrisma.js

@@ -0,0 +1,127 @@
+const RateLimiterStoreAbstract = require('./RateLimiterStoreAbstract');
+const RateLimiterRes = require('./RateLimiterRes');
+
+class RateLimiterPrisma extends RateLimiterStoreAbstract {
+  /**
+   * Constructor for the rate limiter
+   * @param {Object} opts - Options for the rate limiter
+   */
+  constructor(opts) {
+    super(opts);
+
+    this.modelName = opts.tableName || 'RateLimiterFlexible';
+    this.prismaClient = opts.storeClient;
+    this.clearExpiredByTimeout = opts.clearExpiredByTimeout || true;
+
+    if (!this.prismaClient) {
+      throw new Error('Prisma client is not provided');
+    }
+
+    if (this.clearExpiredByTimeout) {
+      this._clearExpiredHourAgo();
+    }
+  }
+
+  _getRateLimiterRes(rlKey, changedPoints, result) {
+    const res = new RateLimiterRes();
+
+    let doc = result;
+
+    res.isFirstInDuration = doc.points === changedPoints;
+    res.consumedPoints = doc.points;
+
+    res.remainingPoints = Math.max(this.points - res.consumedPoints, 0);
+    res.msBeforeNext = doc.expire !== null
+      ? Math.max(new Date(doc.expire).getTime() - Date.now(), 0)
+      : -1;
+
+    return res;
+  }
+
+  _upsert(key, points, msDuration, forceExpire = false) {
+    if (!this.prismaClient) {
+      return Promise.reject(new Error('Prisma client is not established'));
+    }
+
+    const now = new Date();
+    const newExpire = msDuration > 0 ? new Date(now.getTime() + msDuration) : null;
+
+    return this.prismaClient.$transaction(async (prisma) => {
+      const existingRecord = await prisma[this.modelName].findFirst({
+        where: { key: key },
+      });
+
+      if (existingRecord) {
+        // Determine if we should update the expire field
+        const shouldUpdateExpire = forceExpire || !existingRecord.expire || existingRecord.expire <= now || newExpire === null;
+
+        return prisma[this.modelName].update({
+          where: { key: key },
+          data: {
+            points: !shouldUpdateExpire ? existingRecord.points + points : points,
+            ...(shouldUpdateExpire && { expire: newExpire }),
+          },
+        });
+      } else {
+        return prisma[this.modelName].create({
+          data: {
+            key: key,
+            points: points,
+            expire: newExpire,
+          },
+        });
+      }
+    });
+  }
+
+  _get(rlKey) {
+    if (!this.prismaClient) {
+      return Promise.reject(new Error('Prisma client is not established'));
+    }
+
+    return this.prismaClient[this.modelName].findFirst({
+      where: {
+        AND: [
+          { key: rlKey },
+          {
+            OR: [
+              { expire: { gt: new Date() } },
+              { expire: null },
+            ],
+          },
+        ],
+      },
+    });
+  }
+
+  _delete(rlKey) {
+    if (!this.prismaClient) {
+      return Promise.reject(new Error('Prisma client is not established'));
+    }
+
+    return this.prismaClient[this.modelName].deleteMany({
+      where: {
+        key: rlKey,
+      },
+    }).then(res => res.count > 0);
+  }
+
+  _clearExpiredHourAgo() {
+    if (this._clearExpiredTimeoutId) {
+      clearTimeout(this._clearExpiredTimeoutId);
+    }
+    this._clearExpiredTimeoutId = setTimeout(async () => {
+      await this.prismaClient[this.modelName].deleteMany({
+        where: {
+          expire: {
+            lt: new Date(Date.now() - 3600000),
+          },
+        },
+      });
+      this._clearExpiredHourAgo();
+    }, 300000); // Clear every 5 minutes
+    this._clearExpiredTimeoutId.unref();
+  }
+}
+
+module.exports = RateLimiterPrisma;

package/lib/RateLimiterQueue.js

@@ -0,0 +1,131 @@
+const RateLimiterQueueError = require('./component/RateLimiterQueueError')
+const MAX_QUEUE_SIZE = 4294967295;
+const KEY_DEFAULT = 'limiter';
+
+module.exports = class RateLimiterQueue {
+  constructor(limiterFlexible, opts = {}) {
+    const maxQueueSize =
+      opts.maxQueueSize !== undefined ? opts.maxQueueSize : MAX_QUEUE_SIZE;
+    this._queueLimiters = {
+      KEY_DEFAULT: new RateLimiterQueueInternal(limiterFlexible, {
+        ...opts,
+        maxQueueSize,
+        key: KEY_DEFAULT
+      }),
+    };
+    this._limiterFlexible = limiterFlexible;
+    this._maxQueueSize = maxQueueSize;
+  }
+
+  getTokensRemaining(key = KEY_DEFAULT) {
+    if (this._queueLimiters[key]) {
+      return this._queueLimiters[key].getTokensRemaining()
+    } else {
+      return Promise.resolve(this._limiterFlexible.points)
+    }
+  }
+
+  removeTokens(tokens, key = KEY_DEFAULT) {
+    if (!this._queueLimiters[key]) {
+      this._queueLimiters[key] = new RateLimiterQueueInternal(
+        this._limiterFlexible, {
+          key,
+          maxQueueSize: this._maxQueueSize,
+        })
+    }
+
+    return this._queueLimiters[key].removeTokens(tokens)
+  }
+};
+
+class RateLimiterQueueInternal {
+
+  constructor(limiterFlexible, opts = {
+    maxQueueSize: MAX_QUEUE_SIZE,
+    key: KEY_DEFAULT,
+  }) {
+    this._key = opts.key;
+    this._waitTimeout = null;
+    this._queue = [];
+    this._limiterFlexible = limiterFlexible;
+
+    this._maxQueueSize = opts.maxQueueSize
+  }
+
+  getTokensRemaining() {
+    return this._limiterFlexible.get(this._key)
+      .then((rlRes) => {
+        return rlRes !== null ? rlRes.remainingPoints : this._limiterFlexible.points;
+      })
+  }
+
+  removeTokens(tokens) {
+    const _this = this;
+
+    return new Promise((resolve, reject) => {
+      if (tokens > _this._limiterFlexible.points) {
+        reject(new RateLimiterQueueError(`Requested tokens ${tokens} exceeds maximum ${_this._limiterFlexible.points} tokens per interval`));
+        return
+      }
+
+      if (_this._queue.length > 0) {
+        _this._queueRequest.call(_this, resolve, reject, tokens);
+      } else {
+        _this._limiterFlexible.consume(_this._key, tokens)
+          .then((res) => {
+            resolve(res.remainingPoints);
+          })
+          .catch((rej) => {
+            if (rej instanceof Error) {
+              reject(rej);
+            } else {
+              _this._queueRequest.call(_this, resolve, reject, tokens);
+              if (_this._waitTimeout === null) {
+                _this._waitTimeout = setTimeout(_this._processFIFO.bind(_this), rej.msBeforeNext);
+              }
+            }
+          });
+      }
+    })
+  }
+
+  _queueRequest(resolve, reject, tokens) {
+    const _this = this;
+    if (_this._queue.length < _this._maxQueueSize) {
+      _this._queue.push({resolve, reject, tokens});
+    } else {
+      reject(new RateLimiterQueueError(`Number of requests reached it's maximum ${_this._maxQueueSize}`))
+    }
+  }
+
+  _processFIFO() {
+    const _this = this;
+
+    if (_this._waitTimeout !== null) {
+      clearTimeout(_this._waitTimeout);
+      _this._waitTimeout = null;
+    }
+
+    if (_this._queue.length === 0) {
+      return;
+    }
+
+    const item = _this._queue.shift();
+    _this._limiterFlexible.consume(_this._key, item.tokens)
+      .then((res) => {
+        item.resolve(res.remainingPoints);
+        _this._processFIFO.call(_this);
+      })
+      .catch((rej) => {
+        if (rej instanceof Error) {
+          item.reject(rej);
+          _this._processFIFO.call(_this);
+        } else {
+          _this._queue.unshift(item);
+          if (_this._waitTimeout === null) {
+            _this._waitTimeout = setTimeout(_this._processFIFO.bind(_this), rej.msBeforeNext);
+          }
+        }
+      });
+  }
+}