@upstash/ratelimit 1.2.0-canary → 1.2.1-canary

package/dist/index.mjs

@@ -1,3 +1,9 @@
+var __defProp = Object.defineProperty;
+var __export = (target, all) => {
+  for (var name in all)
+    __defProp(target, name, { get: all[name], enumerable: true });
+};
+
 // src/analytics.ts
 import { Analytics as CoreAnalytics } from "@upstash/core-analytics";
 var Analytics = class {
@@ -266,7 +272,107 @@ var resetScript = `
       until cursor == "0"
     `;
 
-// src/deny-list.ts
+// src/types.ts
+var DenyListExtension = "denyList";
+var IpDenyListKey = "ipDenyList";
+var IpDenyListStatusKey = "ipDenyListStatus";
+
+// src/deny-list/scripts.ts
+var checkDenyListScript = `
+  -- Checks if values provideed in ARGV are present in the deny lists.
+  -- This is done using the allDenyListsKey below.
+
+  -- Additionally, checks the status of the ip deny list using the
+  -- ipDenyListStatusKey below. Here are the possible states of the
+  -- ipDenyListStatusKey key:
+  -- * status == -1: set to "disabled" with no TTL
+  -- * status == -2: not set, meaning that is was set before but expired
+  -- * status  >  0: set to "valid", with a TTL
+  --
+  -- In the case of status == -2, we set the status to "pending" with
+  -- 30 second ttl. During this time, the process which got status == -2
+  -- will update the ip deny list.
+
+  local allDenyListsKey     = KEYS[1]
+  local ipDenyListStatusKey = KEYS[2]
+
+  local results = redis.call('SMISMEMBER', allDenyListsKey, unpack(ARGV))
+  local status  = redis.call('TTL', ipDenyListStatusKey)
+  if status == -2 then
+    redis.call('SETEX', ipDenyListStatusKey, 30, "pending")
+  end
+
+  return { results, status }
+`;
+
+// src/deny-list/ip-deny-list.ts
+var ip_deny_list_exports = {};
+__export(ip_deny_list_exports, {
+  ThresholdError: () => ThresholdError,
+  disableIpDenyList: () => disableIpDenyList,
+  updateIpDenyList: () => updateIpDenyList
+});
+
+// src/deny-list/time.ts
+var MILLISECONDS_IN_HOUR = 60 * 60 * 1e3;
+var MILLISECONDS_IN_DAY = 24 * MILLISECONDS_IN_HOUR;
+var MILLISECONDS_TO_2AM = 2 * MILLISECONDS_IN_HOUR;
+var getIpListTTL = (time) => {
+  const now = time || Date.now();
+  const timeSinceLast2AM = (now - MILLISECONDS_TO_2AM) % MILLISECONDS_IN_DAY;
+  return MILLISECONDS_IN_DAY - timeSinceLast2AM;
+};
+
+// src/deny-list/ip-deny-list.ts
+var baseUrl = "https://raw.githubusercontent.com/stamparm/ipsum/master/levels";
+var ThresholdError = class extends Error {
+  constructor(threshold) {
+    super(`Allowed threshold values are from 1 to 8, 1 and 8 included. Received: ${threshold}`);
+    this.name = "ThresholdError";
+  }
+};
+var getIpDenyList = async (threshold) => {
+  if (typeof threshold !== "number" || threshold < 1 || threshold > 8) {
+    throw new ThresholdError(threshold);
+  }
+  try {
+    const response = await fetch(`${baseUrl}/${threshold}.txt`);
+    if (!response.ok) {
+      throw new Error(`Error fetching data: ${response.statusText}`);
+    }
+    const data = await response.text();
+    const lines = data.split("\n");
+    return lines.filter((value) => value.length > 0);
+  } catch (error) {
+    throw new Error(`Failed to fetch ip deny list: ${error}`);
+  }
+};
+var updateIpDenyList = async (redis, prefix, threshold, ttl) => {
+  const allIps = await getIpDenyList(threshold);
+  const allDenyLists = [prefix, DenyListExtension, "all"].join(":");
+  const ipDenyList = [prefix, DenyListExtension, IpDenyListKey].join(":");
+  const statusKey = [prefix, IpDenyListStatusKey].join(":");
+  const transaction = redis.multi();
+  transaction.sdiffstore(allDenyLists, allDenyLists, ipDenyList);
+  transaction.del(ipDenyList);
+  transaction.sadd(ipDenyList, ...allIps);
+  transaction.sdiffstore(ipDenyList, ipDenyList, allDenyLists);
+  transaction.sunionstore(allDenyLists, allDenyLists, ipDenyList);
+  transaction.set(statusKey, "valid", { px: ttl ?? getIpListTTL() });
+  return await transaction.exec();
+};
+var disableIpDenyList = async (redis, prefix) => {
+  const allDenyListsKey = [prefix, DenyListExtension, "all"].join(":");
+  const ipDenyListKey = [prefix, DenyListExtension, IpDenyListKey].join(":");
+  const statusKey = [prefix, IpDenyListStatusKey].join(":");
+  const transaction = redis.multi();
+  transaction.sdiffstore(allDenyListsKey, allDenyListsKey, ipDenyListKey);
+  transaction.del(ipDenyListKey);
+  transaction.set(statusKey, "disabled");
+  return await transaction.exec();
+};
+
+// src/deny-list/deny-list.ts
 var denyListCache = new Cache(/* @__PURE__ */ new Map());
 var checkDenyListCache = (members) => {
   return members.find(
@@ -279,25 +385,39 @@ var blockMember = (member) => {
   denyListCache.blockUntil(member, Date.now() + 6e4);
 };
 var checkDenyList = async (redis, prefix, members) => {
-  const deniedMembers = await redis.smismember(
-    [prefix, "denyList", "all"].join(":"),
+  const [deniedValues, ipDenyListStatus] = await redis.eval(
+    checkDenyListScript,
+    [
+      [prefix, DenyListExtension, "all"].join(":"),
+      [prefix, IpDenyListStatusKey].join(":")
+    ],
     members
   );
-  let deniedMember = void 0;
-  deniedMembers.map((memberDenied, index) => {
+  let deniedValue = void 0;
+  deniedValues.map((memberDenied, index) => {
     if (memberDenied) {
       blockMember(members[index]);
-      deniedMember = members[index];
+      deniedValue = members[index];
     }
   });
-  return deniedMember;
+  return {
+    deniedValue,
+    invalidIpDenyList: ipDenyListStatus === -2
+  };
 };
-var resolveResponses = ([ratelimitResponse, denyListResponse]) => {
-  if (denyListResponse) {
+var resolveLimitPayload = (redis, prefix, [ratelimitResponse, denyListResponse], threshold) => {
+  if (denyListResponse.deniedValue) {
     ratelimitResponse.success = false;
     ratelimitResponse.remaining = 0;
     ratelimitResponse.reason = "denyList";
-    ratelimitResponse.deniedValue = denyListResponse;
+    ratelimitResponse.deniedValue = denyListResponse.deniedValue;
+  }
+  if (denyListResponse.invalidIpDenyList) {
+    const updatePromise = updateIpDenyList(redis, prefix, threshold);
+    ratelimitResponse.pending = Promise.all([
+      ratelimitResponse.pending,
+      updatePromise
+    ]);
   }
   return ratelimitResponse;
 };
@@ -322,12 +442,14 @@ var Ratelimit = class {
   primaryRedis;
   analytics;
   enableProtection;
+  denyListThreshold;
   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.denyListThreshold = config.denyListThreshold ?? 6;
     this.primaryRedis = "redis" in this.ctx ? this.ctx.redis : this.ctx.regionContexts[0].redis;
     this.analytics = config.analytics ? new Analytics({
       redis: this.primaryRedis,
@@ -457,21 +579,17 @@ var Ratelimit = class {
   getRatelimitResponse = async (identifier, req) => {
     const key = this.getKey(identifier);
     const definedMembers = this.getDefinedMembers(identifier, req);
-    const deniedMember = checkDenyListCache(definedMembers);
+    const deniedValue = checkDenyListCache(definedMembers);
     let result;
-    if (deniedMember) {
-      result = [defaultDeniedResponse(deniedMember), deniedMember];
+    if (deniedValue) {
+      result = [defaultDeniedResponse(deniedValue), { deniedValue, invalidIpDenyList: false }];
     } else {
       result = await Promise.all([
         this.limiter().limit(this.ctx, key, req?.rate),
-        checkDenyList(
-          this.primaryRedis,
-          this.prefix,
-          definedMembers
-        )
+        this.enableProtection ? checkDenyList(this.primaryRedis, this.prefix, definedMembers) : { deniedValue: void 0, invalidIpDenyList: false }
       ]);
     }
-    return resolveResponses(result);
+    return resolveLimitPayload(this.primaryRedis, this.prefix, result, this.denyListThreshold);
   };
   /**
    * Creates an array with the original response promise and a timeout promise
@@ -1081,7 +1199,8 @@ var RegionRatelimit = class extends Ratelimit {
         cacheScripts: config.cacheScripts ?? true
       },
       ephemeralCache: config.ephemeralCache,
-      enableProtection: config.enableProtection
+      enableProtection: config.enableProtection,
+      denyListThreshold: config.denyListThreshold
     });
   }
   /**
@@ -1447,6 +1566,7 @@ var RegionRatelimit = class extends Ratelimit {
 };
 export {
   Analytics,
+  ip_deny_list_exports as IpDenyList,
   MultiRegionRatelimit,
   RegionRatelimit as Ratelimit
 };