Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

DX-960: Update Deny List from SDK #112

Merged
merged 10 commits into from
Jun 13, 2024
2 changes: 1 addition & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
83 changes: 68 additions & 15 deletions src/deny-list.test.ts → src/deny-list/deny-list.test.ts
Original file line number Diff line number Diff line change
@@ -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", () => {
Expand All @@ -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,
Expand All @@ -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 () => {
Expand Down
52 changes: 37 additions & 15 deletions src/deny-list.ts → src/deny-list/deny-list.ts
Original file line number Diff line number Diff line change
@@ -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());
Expand Down Expand Up @@ -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
};
};

/**
Expand All @@ -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;
};

Expand Down
1 change: 1 addition & 0 deletions src/deny-list/index.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
export * from "./deny-list"
Loading
Loading