Merge pull request #112 from upstash/DX-960

DX-960: Update Deny List from SDK

Author: CahidArda

Diff for: package.json

@@ -12,7 +12,7 @@
     "fmt": "bunx @biomejs/biome check --apply ./src"
   },
   "devDependencies": {
-    "@upstash/redis": "^1.28.3",
+    "@upstash/redis": "^1.31.5",
     "bun-types": "latest",
     "rome": "^11.0.0",
     "tsup": "^7.2.0",

Diff for: src/deny-list.test.ts renamed to src/deny-list/deny-list.test.ts

@@ -1,8 +1,8 @@
-import { expect, test, describe, afterAll } from "bun:test";
+import { expect, test, describe, afterAll, beforeAll } from "bun:test";
 import { Redis } from "@upstash/redis";
-import { Ratelimit } from "./index";
-import { checkDenyListCache, defaultDeniedResponse, resolveResponses } from "./deny-list";
-import { RatelimitResponseType } from "./types";
+import { Ratelimit } from "../index";
+import { checkDenyListCache, defaultDeniedResponse, resolveLimitPayload } from "./deny-list";
+import { DenyListResponse, RatelimitResponseType } from "../types";
 
 
 test("should get expected response from defaultDeniedResponse", () => {
@@ -20,8 +20,18 @@ test("should get expected response from defaultDeniedResponse", () => {
   });
 });
 
+describe("should resolve ratelimit and deny list response", async () => {
+  const redis = Redis.fromEnv();
+  const prefix = `test-resolve-prefix`;
+
+  let callCount = 0;
+  const spyRedis = {
+    multi: () => {
+      callCount += 1;
+      return redis.multi();
+    }
+  }
 
-test.only("should override response in resolveResponses correctly", () => {
   const initialResponse = {
     success: true,
     limit: 100,
@@ -31,40 +41,83 @@ test.only("should override response in resolveResponses correctly", () => {
     reason: undefined,
     deniedValue: undefined
   };
-  
-  const denyListResponse = "testValue";
+
   const expectedResponse = {
     success: false,
     limit: 100,
     remaining: 0,
     reset: 60,
     pending: Promise.resolve(),
     reason: "denyList" as RatelimitResponseType,
-    deniedValue: denyListResponse
+    deniedValue: "testValue"
   };
 
-  const response = resolveResponses([initialResponse, denyListResponse]);
-  expect(response).toEqual(expectedResponse);
-});
+  test("should update ip deny list when invalidIpDenyList is true", async () => {
+    let callCount = 0;
+    const spyRedis = {
+      multi: () => {
+        callCount += 1;
+        return redis.multi();
+      }
+    }
+  
+    const denyListResponse: DenyListResponse = {
+      deniedValue: "testValue",
+      invalidIpDenyList: true
+    };
+
+    const response = resolveLimitPayload(spyRedis as Redis, prefix, [initialResponse, denyListResponse], 8);
+    await response.pending;
+  
+    expect(response).toEqual(expectedResponse);
+    expect(callCount).toBe(1) // calls multi once to store ips
+  });
+
+  test("should update ip deny list when invalidIpDenyList is true", async () => {
+
+    let callCount = 0;
+    const spyRedis = {
+      multi: () => {
+        callCount += 1;
+        return redis.multi();
+      }
+    }
+  
+    const denyListResponse: DenyListResponse = {
+      deniedValue: "testValue",
+      invalidIpDenyList: false
+    };
+
+    const response = resolveLimitPayload(spyRedis as Redis, prefix, [initialResponse, denyListResponse], 8);
+    await response.pending;
+  
+    expect(response).toEqual(expectedResponse);
+    expect(callCount).toBe(0) // doesn't call multi to update deny list
+  });
+})
 
 
 describe("should reject in deny list", async () => {
   const redis = Redis.fromEnv();
   const prefix = `test-prefix`;
   const denyListKey = [prefix, "denyList", "all"].join(":");
 
-  // Insert a value into the deny list
-  await redis.sadd(denyListKey, "denyIdentifier", "denyIp", "denyAgent", "denyCountry");
 
   const ratelimit = new Ratelimit({
     redis,
     limiter: Ratelimit.tokenBucket(10, "5 s", 10),
     prefix,
-    enableProtection: true
+    enableProtection: true,
+    denyListThreshold: 8
   });
 
   afterAll(async () => {
-    redis.del(denyListKey)
+    await redis.del(denyListKey)
+  })
+
+  // Insert a value into the deny list
+  beforeAll(async () => {
+    await redis.sadd(denyListKey, "denyIdentifier", "denyIp", "denyAgent", "denyCountry");
   })
 
   test("should allow with values not in the deny list", async () => {

Diff for: src/deny-list.ts renamed to src/deny-list/deny-list.ts

@@ -1,6 +1,8 @@
-import { DeniedValue, LimitPayload, Redis } from "./types"
-import { RatelimitResponse } from "./types"
-import { Cache } from "./cache";
+import { DeniedValue, DenyListResponse, DenyListExtension, LimitPayload, IpDenyListStatusKey } from "../types"
+import { RatelimitResponse, Redis } from "../types"
+import { Cache } from "../cache";
+import { checkDenyListScript } from "./scripts";
+import { updateIpDenyList } from "./ip-deny-list";
 
 
 const denyListCache = new Cache(new Map());
@@ -46,21 +48,28 @@ export const checkDenyList = async (
   redis: Redis,
   prefix: string,
   members: string[]
-): Promise<DeniedValue> => {
-  const deniedMembers = await redis.smismember(
-    [prefix, "denyList", "all"].join(":"),
+): Promise<DenyListResponse> => {
+  const [ deniedValues, ipDenyListStatus ] = await redis.eval(
+    checkDenyListScript,
+    [
+      [prefix, DenyListExtension, "all"].join(":"),
+      [prefix, IpDenyListStatusKey].join(":"),
+    ],
     members
-  );
+  ) as [boolean[], number];
 
-  let deniedMember: DeniedValue = undefined;
-  deniedMembers.map((memberDenied, index) => {
+  let deniedValue: DeniedValue = undefined;
+  deniedValues.map((memberDenied, index) => {
     if (memberDenied) {
       blockMember(members[index])
-      deniedMember = members[index]
+      deniedValue = members[index]
     }
   })
 
-  return deniedMember;
+  return {
+    deniedValue,
+    invalidIpDenyList: ipDenyListStatus === -2
+  };
 };
 
 /**
@@ -71,15 +80,28 @@ export const checkDenyList = async (
  * @param denyListResponse 
  * @returns 
  */
-export const resolveResponses = (
-  [ratelimitResponse, denyListResponse]: LimitPayload
+export const resolveLimitPayload = (
+  redis: Redis,
+  prefix: string,
+  [ratelimitResponse, denyListResponse]: LimitPayload,
+  threshold: number
 ): RatelimitResponse => {
-  if (denyListResponse) {
+
+  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;
 };
 

Diff for: src/deny-list/index.ts

@@ -0,0 +1 @@
+export * from "./deny-list"