@upstash/ratelimit 1.1.3 → 1.2.0-canary

package/dist/index.mjs

@@ -97,6 +97,9 @@ var Cache = class {
   empty() {
     this.cache.clear();
   }
+  size() {
+    return this.cache.size;
+  }
 };
 
 // src/duration.ts
@@ -123,6 +126,37 @@ function ms(d) {
   }
 }
 
+// src/hash.ts
+var setHash = async (ctx, script, kind) => {
+  const regionContexts = "redis" in ctx ? [ctx] : ctx.regionContexts;
+  const hashSample = regionContexts[0].scriptHashes[kind];
+  if (!hashSample) {
+    await Promise.all(regionContexts.map(async (context) => {
+      context.scriptHashes[kind] = await context.redis.scriptLoad(script);
+    }));
+  }
+  ;
+};
+var safeEval = async (ctx, script, kind, keys, args) => {
+  if (!ctx.cacheScripts) {
+    return await ctx.redis.eval(script, keys, args);
+  }
+  ;
+  await setHash(ctx, script, kind);
+  try {
+    return await ctx.redis.evalsha(ctx.scriptHashes[kind], keys, args);
+  } catch (error) {
+    if (`${error}`.includes("NOSCRIPT")) {
+      console.log("Script with the expected hash was not found in redis db. It is probably flushed. Will load another scipt before continuing.");
+      ctx.scriptHashes[kind] = void 0;
+      await setHash(ctx, script, kind);
+      console.log("  New script successfully loaded.");
+      return await ctx.redis.evalsha(ctx.scriptHashes[kind], keys, args);
+    }
+    throw error;
+  }
+};
+
 // src/lua-scripts/multi.ts
 var fixedWindowLimitScript = `
 	local key           = KEYS[1]
@@ -232,20 +266,71 @@ var resetScript = `
       until cursor == "0"
     `;
 
+// src/deny-list.ts
+var denyListCache = new Cache(/* @__PURE__ */ new Map());
+var checkDenyListCache = (members) => {
+  return members.find(
+    (member) => denyListCache.isBlocked(member).blocked
+  );
+};
+var blockMember = (member) => {
+  if (denyListCache.size() > 1e3)
+    denyListCache.empty();
+  denyListCache.blockUntil(member, Date.now() + 6e4);
+};
+var checkDenyList = async (redis, prefix, members) => {
+  const deniedMembers = await redis.smismember(
+    [prefix, "denyList", "all"].join(":"),
+    members
+  );
+  let deniedMember = void 0;
+  deniedMembers.map((memberDenied, index) => {
+    if (memberDenied) {
+      blockMember(members[index]);
+      deniedMember = members[index];
+    }
+  });
+  return deniedMember;
+};
+var resolveResponses = ([ratelimitResponse, denyListResponse]) => {
+  if (denyListResponse) {
+    ratelimitResponse.success = false;
+    ratelimitResponse.remaining = 0;
+    ratelimitResponse.reason = "denyList";
+    ratelimitResponse.deniedValue = denyListResponse;
+  }
+  return ratelimitResponse;
+};
+var defaultDeniedResponse = (deniedValue) => {
+  return {
+    success: false,
+    limit: 0,
+    remaining: 0,
+    reset: 0,
+    pending: Promise.resolve(),
+    reason: "denyList",
+    deniedValue
+  };
+};
+
 // src/ratelimit.ts
 var Ratelimit = class {
   limiter;
   ctx;
   prefix;
   timeout;
+  primaryRedis;
   analytics;
+  enableProtection;
   constructor(config) {
     this.ctx = config.ctx;
     this.limiter = config.limiter;
     this.timeout = config.timeout ?? 5e3;
     this.prefix = config.prefix ?? "@upstash/ratelimit";
+    this.enableProtection = config.enableProtection ?? false;
+    this.primaryRedis = "redis" in this.ctx ? this.ctx.redis : this.ctx.regionContexts[0].redis;
     this.analytics = config.analytics ? new Analytics({
-      redis: Array.isArray(this.ctx.redis) ? this.ctx.redis[0] : this.ctx.redis,
+      redis: this.primaryRedis,
       prefix: this.prefix
     }) : void 0;
     if (config.ephemeralCache instanceof Map) {
@@ -291,57 +376,14 @@ var Ratelimit = class {
    * ```
    */
   limit = async (identifier, req) => {
-    const key = [this.prefix, identifier].join(":");
     let timeoutId = null;
     try {
-      const arr = [this.limiter().limit(this.ctx, key, req?.rate)];
-      if (this.timeout > 0) {
-        arr.push(
-          new Promise((resolve) => {
-            timeoutId = setTimeout(() => {
-              resolve({
-                success: true,
-                limit: 0,
-                remaining: 0,
-                reset: 0,
-                pending: Promise.resolve()
-              });
-            }, this.timeout);
-          })
-        );
-      }
-      const res = await Promise.race(arr);
-      if (this.analytics) {
-        try {
-          const geo = req ? this.analytics.extractGeo(req) : void 0;
-          const analyticsP = this.analytics.record({
-            identifier,
-            time: Date.now(),
-            success: res.success,
-            ...geo
-          }).catch((err) => {
-            let errorMessage = "Failed to record analytics";
-            if (`${err}`.includes("WRONGTYPE")) {
-              errorMessage = `
-Failed to record analytics. See the information below:
-
-This can occur when you uprade to Ratelimit version 1.1.2
-or later from an earlier version.
-
-This occurs simply because the way we store analytics data
-has changed. To avoid getting this error, disable analytics
-for *an hour*, then simply enable it back.
-
-`;
-            }
-            console.warn(errorMessage, err);
-          });
-          res.pending = Promise.all([res.pending, analyticsP]);
-        } catch (err) {
-          console.warn("Failed to record analytics", err);
-        }
-      }
-      return res;
+      const response = this.getRatelimitResponse(identifier, req);
+      const { responseArray, newTimeoutId } = this.applyTimeout(response);
+      timeoutId = newTimeoutId;
+      const timedResponse = await Promise.race(responseArray);
+      const finalResponse = this.submitAnalytics(timedResponse, identifier, req);
+      return finalResponse;
     } finally {
       if (timeoutId) {
         clearTimeout(timeoutId);
@@ -400,6 +442,125 @@ for *an hour*, then simply enable it back.
     const pattern = [this.prefix, identifier].join(":");
     return await this.limiter().getRemaining(this.ctx, pattern);
   };
+  /**
+   * Checks if the identifier or the values in req are in the deny list cache.
+   * If so, returns the default denied response.
+   * 
+   * Otherwise, calls redis to check the rate limit and deny list. Returns after
+   * resolving the result. Resolving is overriding the rate limit result if
+   * the some value is in deny list.
+   * 
+   * @param identifier identifier to block
+   * @param req options with ip, user agent, country, rate and geo info
+   * @returns rate limit response
+   */
+  getRatelimitResponse = async (identifier, req) => {
+    const key = this.getKey(identifier);
+    const definedMembers = this.getDefinedMembers(identifier, req);
+    const deniedMember = checkDenyListCache(definedMembers);
+    let result;
+    if (deniedMember) {
+      result = [defaultDeniedResponse(deniedMember), deniedMember];
+    } else {
+      result = await Promise.all([
+        this.limiter().limit(this.ctx, key, req?.rate),
+        checkDenyList(
+          this.primaryRedis,
+          this.prefix,
+          definedMembers
+        )
+      ]);
+    }
+    return resolveResponses(result);
+  };
+  /**
+   * Creates an array with the original response promise and a timeout promise
+   * if this.timeout > 0.
+   * 
+   * @param response Ratelimit response promise
+   * @returns array with the response and timeout promise. also includes the timeout id
+   */
+  applyTimeout = (response) => {
+    let newTimeoutId = null;
+    const responseArray = [response];
+    if (this.timeout > 0) {
+      const timeoutResponse = new Promise((resolve) => {
+        newTimeoutId = setTimeout(() => {
+          resolve({
+            success: true,
+            limit: 0,
+            remaining: 0,
+            reset: 0,
+            pending: Promise.resolve(),
+            reason: "timeout"
+          });
+        }, this.timeout);
+      });
+      responseArray.push(timeoutResponse);
+    }
+    return {
+      responseArray,
+      newTimeoutId
+    };
+  };
+  /**
+   * submits analytics if this.analytics is set
+   * 
+   * @param ratelimitResponse final rate limit response
+   * @param identifier identifier to submit
+   * @param req limit options
+   * @returns rate limit response after updating the .pending field
+   */
+  submitAnalytics = (ratelimitResponse, identifier, req) => {
+    if (this.analytics) {
+      try {
+        const geo = req ? this.analytics.extractGeo(req) : void 0;
+        const analyticsP = this.analytics.record({
+          identifier: ratelimitResponse.reason === "denyList" ? ratelimitResponse.deniedValue : identifier,
+          time: Date.now(),
+          success: ratelimitResponse.reason === "denyList" ? "denied" : ratelimitResponse.success,
+          ...geo
+        }).catch((err) => {
+          let errorMessage = "Failed to record analytics";
+          if (`${err}`.includes("WRONGTYPE")) {
+            errorMessage = `
+    Failed to record analytics. See the information below:
+
+    This can occur when you uprade to Ratelimit version 1.1.2
+    or later from an earlier version.
+
+    This occurs simply because the way we store analytics data
+    has changed. To avoid getting this error, disable analytics
+    for *an hour*, then simply enable it back.
+
+    `;
+          }
+          console.warn(errorMessage, err);
+        });
+        ratelimitResponse.pending = Promise.all([ratelimitResponse.pending, analyticsP]);
+      } catch (err) {
+        console.warn("Failed to record analytics", err);
+      }
+      ;
+    }
+    ;
+    return ratelimitResponse;
+  };
+  getKey = (identifier) => {
+    return [this.prefix, identifier].join(":");
+  };
+  /**
+   * returns a list of defined values from
+   * [identifier, req.ip, req.userAgent, req.country]
+   * 
+   * @param identifier identifier
+   * @param req limit options
+   * @returns list of defined values
+   */
+  getDefinedMembers = (identifier, req) => {
+    const members = [identifier, req?.ip, req?.userAgent, req?.country];
+    return members.filter((item) => Boolean(item));
+  };
 };
 
 // src/multi.ts
@@ -423,7 +584,11 @@ var MultiRegionRatelimit = class extends Ratelimit {
       timeout: config.timeout,
       analytics: config.analytics,
       ctx: {
-        redis: config.redis,
+        regionContexts: config.redis.map((redis) => ({
+          redis,
+          scriptHashes: {},
+          cacheScripts: config.cacheScripts ?? true
+        })),
         cache: config.ephemeralCache ? new Cache(config.ephemeralCache) : void 0
       }
     });
@@ -458,7 +623,8 @@ var MultiRegionRatelimit = class extends Ratelimit {
               limit: tokens,
               remaining: 0,
               reset: reset2,
-              pending: Promise.resolve()
+              pending: Promise.resolve(),
+              reason: "cacheBlock"
             };
           }
         }
@@ -466,10 +632,12 @@ var MultiRegionRatelimit = class extends Ratelimit {
         const bucket = Math.floor(Date.now() / windowDuration);
         const key = [identifier, bucket].join(":");
         const incrementBy = rate ? Math.max(1, rate) : 1;
-        const dbs = ctx.redis.map((redis) => ({
-          redis,
-          request: redis.eval(
+        const dbs = ctx.regionContexts.map((regionContext) => ({
+          redis: regionContext.redis,
+          request: safeEval(
+            regionContext,
             fixedWindowLimitScript,
+            "limitHash",
             [key],
             [requestId, windowDuration, incrementBy]
           )
@@ -540,9 +708,15 @@ var MultiRegionRatelimit = class extends Ratelimit {
       async getRemaining(ctx, identifier) {
         const bucket = Math.floor(Date.now() / windowDuration);
         const key = [identifier, bucket].join(":");
-        const dbs = ctx.redis.map((redis) => ({
-          redis,
-          request: redis.eval(fixedWindowRemainingTokensScript, [key], [null])
+        const dbs = ctx.regionContexts.map((regionContext) => ({
+          redis: regionContext.redis,
+          request: safeEval(
+            regionContext,
+            fixedWindowRemainingTokensScript,
+            "getRemainingHash",
+            [key],
+            [null]
+          )
         }));
         const firstResponse = await Promise.any(dbs.map((s) => s.request));
         const usedTokens = firstResponse.reduce((accTokens, usedToken, index) => {
@@ -559,9 +733,15 @@ var MultiRegionRatelimit = class extends Ratelimit {
         if (ctx.cache) {
           ctx.cache.pop(identifier);
         }
-        for (const db of ctx.redis) {
-          await db.eval(resetScript, [pattern], [null]);
-        }
+        await Promise.all(ctx.regionContexts.map((regionContext) => {
+          safeEval(
+            regionContext,
+            resetScript,
+            "resetHash",
+            [pattern],
+            [null]
+          );
+        }));
       }
     });
   }
@@ -586,6 +766,19 @@ var MultiRegionRatelimit = class extends Ratelimit {
     const windowDuration = ms(window);
     return () => ({
       async limit(ctx, identifier, rate) {
+        if (ctx.cache) {
+          const { blocked, reset: reset2 } = ctx.cache.isBlocked(identifier);
+          if (blocked) {
+            return {
+              success: false,
+              limit: tokens,
+              remaining: 0,
+              reset: reset2,
+              pending: Promise.resolve(),
+              reason: "cacheBlock"
+            };
+          }
+        }
         const requestId = randomId();
         const now = Date.now();
         const currentWindow = Math.floor(now / windowSize);
@@ -593,10 +786,12 @@ var MultiRegionRatelimit = class extends Ratelimit {
         const previousWindow = currentWindow - 1;
         const previousKey = [identifier, previousWindow].join(":");
         const incrementBy = rate ? Math.max(1, rate) : 1;
-        const dbs = ctx.redis.map((redis) => ({
-          redis,
-          request: redis.eval(
+        const dbs = ctx.regionContexts.map((regionContext) => ({
+          redis: regionContext.redis,
+          request: safeEval(
+            regionContext,
             slidingWindowLimitScript,
+            "limitHash",
             [currentKey, previousKey],
             [tokens, now, windowDuration, requestId, incrementBy]
             // lua seems to return `1` for true and `null` for false
@@ -681,10 +876,12 @@ var MultiRegionRatelimit = class extends Ratelimit {
         const currentKey = [identifier, currentWindow].join(":");
         const previousWindow = currentWindow - 1;
         const previousKey = [identifier, previousWindow].join(":");
-        const dbs = ctx.redis.map((redis) => ({
-          redis,
-          request: redis.eval(
+        const dbs = ctx.regionContexts.map((regionContext) => ({
+          redis: regionContext.redis,
+          request: safeEval(
+            regionContext,
             slidingWindowRemainingTokensScript,
+            "getRemainingHash",
             [currentKey, previousKey],
             [now, windowSize]
             // lua seems to return `1` for true and `null` for false
@@ -698,9 +895,15 @@ var MultiRegionRatelimit = class extends Ratelimit {
         if (ctx.cache) {
           ctx.cache.pop(identifier);
         }
-        for (const db of ctx.redis) {
-          await db.eval(resetScript, [pattern], [null]);
-        }
+        await Promise.all(ctx.regionContexts.map((regionContext) => {
+          safeEval(
+            regionContext,
+            resetScript,
+            "resetHash",
+            [pattern],
+            [null]
+          );
+        }));
       }
     });
   }
@@ -873,9 +1076,12 @@ var RegionRatelimit = class extends Ratelimit {
       timeout: config.timeout,
       analytics: config.analytics,
       ctx: {
-        redis: config.redis
+        redis: config.redis,
+        scriptHashes: {},
+        cacheScripts: config.cacheScripts ?? true
       },
-      ephemeralCache: config.ephemeralCache
+      ephemeralCache: config.ephemeralCache,
+      enableProtection: config.enableProtection
     });
   }
   /**
@@ -910,13 +1116,16 @@ var RegionRatelimit = class extends Ratelimit {
               limit: tokens,
               remaining: 0,
               reset: reset2,
-              pending: Promise.resolve()
+              pending: Promise.resolve(),
+              reason: "cacheBlock"
             };
           }
         }
         const incrementBy = rate ? Math.max(1, rate) : 1;
-        const usedTokensAfterUpdate = await ctx.redis.eval(
+        const usedTokensAfterUpdate = await safeEval(
+          ctx,
           fixedWindowLimitScript2,
+          "limitHash",
           [key],
           [windowDuration, incrementBy]
         );
@@ -937,8 +1146,10 @@ var RegionRatelimit = class extends Ratelimit {
       async getRemaining(ctx, identifier) {
         const bucket = Math.floor(Date.now() / windowDuration);
         const key = [identifier, bucket].join(":");
-        const usedTokens = await ctx.redis.eval(
+        const usedTokens = await safeEval(
+          ctx,
           fixedWindowRemainingTokensScript2,
+          "getRemainingHash",
           [key],
           [null]
         );
@@ -949,7 +1160,13 @@ var RegionRatelimit = class extends Ratelimit {
         if (ctx.cache) {
           ctx.cache.pop(identifier);
         }
-        await ctx.redis.eval(resetScript, [pattern], [null]);
+        await safeEval(
+          ctx,
+          resetScript,
+          "resetHash",
+          [pattern],
+          [null]
+        );
       }
     });
   }
@@ -986,13 +1203,16 @@ var RegionRatelimit = class extends Ratelimit {
               limit: tokens,
               remaining: 0,
               reset: reset2,
-              pending: Promise.resolve()
+              pending: Promise.resolve(),
+              reason: "cacheBlock"
             };
           }
         }
         const incrementBy = rate ? Math.max(1, rate) : 1;
-        const remainingTokens = await ctx.redis.eval(
+        const remainingTokens = await safeEval(
+          ctx,
           slidingWindowLimitScript2,
+          "limitHash",
           [currentKey, previousKey],
           [tokens, now, windowSize, incrementBy]
         );
@@ -1015,8 +1235,10 @@ var RegionRatelimit = class extends Ratelimit {
         const currentKey = [identifier, currentWindow].join(":");
         const previousWindow = currentWindow - 1;
         const previousKey = [identifier, previousWindow].join(":");
-        const usedTokens = await ctx.redis.eval(
+        const usedTokens = await safeEval(
+          ctx,
           slidingWindowRemainingTokensScript2,
+          "getRemainingHash",
           [currentKey, previousKey],
           [now, windowSize]
         );
@@ -1027,7 +1249,13 @@ var RegionRatelimit = class extends Ratelimit {
         if (ctx.cache) {
           ctx.cache.pop(identifier);
         }
-        await ctx.redis.eval(resetScript, [pattern], [null]);
+        await safeEval(
+          ctx,
+          resetScript,
+          "resetHash",
+          [pattern],
+          [null]
+        );
       }
     });
   }
@@ -1056,14 +1284,17 @@ var RegionRatelimit = class extends Ratelimit {
               limit: maxTokens,
               remaining: 0,
               reset: reset2,
-              pending: Promise.resolve()
+              pending: Promise.resolve(),
+              reason: "cacheBlock"
             };
           }
         }
         const now = Date.now();
         const incrementBy = rate ? Math.max(1, rate) : 1;
-        const [remaining, reset] = await ctx.redis.eval(
+        const [remaining, reset] = await safeEval(
+          ctx,
           tokenBucketLimitScript,
+          "limitHash",
           [identifier],
           [maxTokens, intervalDuration, refillRate, now, incrementBy]
         );
@@ -1080,8 +1311,10 @@ var RegionRatelimit = class extends Ratelimit {
         };
       },
       async getRemaining(ctx, identifier) {
-        const remainingTokens = await ctx.redis.eval(
+        const remainingTokens = await safeEval(
+          ctx,
           tokenBucketRemainingTokensScript,
+          "getRemainingHash",
           [identifier],
           [maxTokens]
         );
@@ -1092,7 +1325,13 @@ var RegionRatelimit = class extends Ratelimit {
         if (ctx.cache) {
           ctx.cache.pop(identifier);
         }
-        await ctx.redis.eval(resetScript, [pattern], [null]);
+        await safeEval(
+          ctx,
+          resetScript,
+          "resetHash",
+          [pattern],
+          [null]
+        );
       }
     });
   }
@@ -1135,9 +1374,13 @@ var RegionRatelimit = class extends Ratelimit {
         if (hit) {
           const cachedTokensAfterUpdate = ctx.cache.incr(key);
           const success = cachedTokensAfterUpdate < tokens;
-          const pending = success ? ctx.redis.eval(cachedFixedWindowLimitScript, [key], [windowDuration, incrementBy]).then((t) => {
-            ctx.cache.set(key, t);
-          }) : Promise.resolve();
+          const pending = success ? safeEval(
+            ctx,
+            cachedFixedWindowLimitScript,
+            "limitHash",
+            [key],
+            [windowDuration, incrementBy]
+          ) : Promise.resolve();
           return {
             success,
             limit: tokens,
@@ -1146,8 +1389,10 @@ var RegionRatelimit = class extends Ratelimit {
             pending
           };
         }
-        const usedTokensAfterUpdate = await ctx.redis.eval(
+        const usedTokensAfterUpdate = await safeEval(
+          ctx,
           cachedFixedWindowLimitScript,
+          "limitHash",
           [key],
           [windowDuration, incrementBy]
         );
@@ -1172,8 +1417,10 @@ var RegionRatelimit = class extends Ratelimit {
           const cachedUsedTokens = ctx.cache.get(key) ?? 0;
           return Math.max(0, tokens - cachedUsedTokens);
         }
-        const usedTokens = await ctx.redis.eval(
+        const usedTokens = await safeEval(
+          ctx,
           cachedFixedWindowRemainingTokenScript,
+          "getRemainingHash",
           [key],
           [null]
         );
@@ -1187,7 +1434,13 @@ var RegionRatelimit = class extends Ratelimit {
         const key = [identifier, bucket].join(":");
         ctx.cache.pop(key);
         const pattern = [identifier, "*"].join(":");
-        await ctx.redis.eval(resetScript, [pattern], [null]);
+        await safeEval(
+          ctx,
+          resetScript,
+          "resetHash",
+          [pattern],
+          [null]
+        );
       }
     });
   }