From 97bfdf74026e396de85d441ff2651e1cc75ded7f Mon Sep 17 00:00:00 2001 From: Blaine Bublitz Date: Wed, 7 Feb 2024 10:37:10 -0700 Subject: [PATCH 1/3] feat: Allow user-defined characteristics on rate limit options --- arcjet/index.ts | 143 +++++++++++++++++++++++++-------- arcjet/test/index.edge.test.ts | 29 +++++-- 2 files changed, 132 insertions(+), 40 deletions(-) diff --git a/arcjet/index.ts b/arcjet/index.ts index 7aa8acafd..70d28b37d 100644 --- a/arcjet/index.ts +++ b/arcjet/index.ts @@ -111,6 +111,10 @@ function errorMessage(err: unknown): string { // https://github.com/sindresorhus/type-fest/blob/964466c9d59c711da57a5297ad954c13132a0001/source/simplify.d.ts // UnionToIntersection: // https://github.com/sindresorhus/type-fest/blob/017bf38ebb52df37c297324d97bcc693ec22e920/source/union-to-intersection.d.ts +// IsNever: +// https://github.com/sindresorhus/type-fest/blob/e02f228f6391bb2b26c32a55dfe1e3aa2386d515/source/primitive.d.ts +// LiteralCheck & IsStringLiteral: +// https://github.com/sindresorhus/type-fest/blob/e02f228f6391bb2b26c32a55dfe1e3aa2386d515/source/is-literal.d.ts // // Licensed: MIT License Copyright (c) Sindre Sorhus // (https://sindresorhus.com) @@ -149,6 +153,25 @@ type UnionToIntersection = ? // The `& Union` is to allow indexing by the resulting type Intersection & Union : never; +type IsNever = [T] extends [never] ? true : false; +type LiteralCheck< + T, + LiteralType extends + | null + | undefined + | string + | number + | boolean + | symbol + | bigint, +> = IsNever extends false // Must be wider than `never` + ? [T] extends [LiteralType] // Must be narrower than `LiteralType` + ? [LiteralType] extends [T] // Cannot be wider than `LiteralType` + ? false + : true + : false + : false; +type IsStringLiteral = LiteralCheck; export interface RemoteClient { decide( @@ -417,30 +440,31 @@ function runtime(): Runtime { } } -type TokenBucketRateLimitOptions = { +type TokenBucketRateLimitOptions = { mode?: ArcjetMode; match?: string; - characteristics?: string[]; + characteristics?: Characteristics; refillRate: number; interval: string | number; capacity: number; }; -type FixedWindowRateLimitOptions = { +type FixedWindowRateLimitOptions = { mode?: ArcjetMode; match?: string; - characteristics?: string[]; + characteristics?: Characteristics; window: string | number; max: number; }; -type SlidingWindowRateLimitOptions = { - mode?: ArcjetMode; - match?: string; - characteristics?: string[]; - interval: string | number; - max: number; -}; +type SlidingWindowRateLimitOptions = + { + mode?: ArcjetMode; + match?: string; + characteristics?: Characteristics; + interval: string | number; + max: number; + }; /** * Bot detection is disabled by default. The `bots` configuration block allows @@ -550,6 +574,25 @@ type PlainObject = { [key: string]: unknown }; export type Primitive = ArcjetRule[]; export type Product = ArcjetRule[]; +// User-defined characteristics alter the required props of an ArcjetRequest +// Note: If a user doesn't provide the object literal to our primitives +// directly, we fallback to no required props. They can opt-in by adding the +// `as const` suffix to the characteristics array. +type PropsForCharacteristic = IsStringLiteral extends true + ? T extends + | "ip.src" + | "http.host" + | "http.method" + | "http.request.uri.path" + | `http.request.headers["${string}"]` + | `http.request.cookie["${string}"]` + | `http.request.uri.args["${string}"]` + ? {} + : T extends string + ? Record + : never + : {}; +// Rules can specify they require specific props on an ArcjetRequest type PropsForRule = R extends ArcjetRule ? Props : {}; // We theoretically support an arbitrary amount of rule flattening, // but one level seems to be easiest; however, this puts a constraint of @@ -590,10 +633,16 @@ function isLocalRule( ); } -export function tokenBucket( - options?: TokenBucketRateLimitOptions, - ...additionalOptions: TokenBucketRateLimitOptions[] -): Primitive<{ requested: number }> { +export function tokenBucket< + const Characteristics extends readonly string[] = [], +>( + options?: TokenBucketRateLimitOptions, + ...additionalOptions: TokenBucketRateLimitOptions[] +): Primitive< + UnionToIntersection< + { requested: number } | PropsForCharacteristic + > +> { const rules: ArcjetTokenBucketRateLimitRule<{ requested: number }>[] = []; if (typeof options === "undefined") { @@ -603,7 +652,9 @@ export function tokenBucket( for (const opt of [options, ...additionalOptions]) { const mode = opt.mode === "LIVE" ? "LIVE" : "DRY_RUN"; const match = opt.match; - const characteristics = opt.characteristics; + const characteristics = Array.isArray(opt.characteristics) + ? opt.characteristics + : undefined; const refillRate = opt.refillRate; const interval = duration.parse(opt.interval); @@ -625,10 +676,14 @@ export function tokenBucket( return rules; } -export function fixedWindow( - options?: FixedWindowRateLimitOptions, - ...additionalOptions: FixedWindowRateLimitOptions[] -): Primitive { +export function fixedWindow< + const Characteristics extends readonly string[] = [], +>( + options?: FixedWindowRateLimitOptions, + ...additionalOptions: FixedWindowRateLimitOptions[] +): Primitive< + UnionToIntersection> +> { const rules: ArcjetFixedWindowRateLimitRule<{}>[] = []; if (typeof options === "undefined") { @@ -638,7 +693,9 @@ export function fixedWindow( for (const opt of [options, ...additionalOptions]) { const mode = opt.mode === "LIVE" ? "LIVE" : "DRY_RUN"; const match = opt.match; - const characteristics = opt.characteristics; + const characteristics = Array.isArray(opt.characteristics) + ? opt.characteristics + : undefined; const max = opt.max; const window = duration.parse(opt.window); @@ -660,19 +717,25 @@ export function fixedWindow( // This is currently kept for backwards compatibility but should be removed in // favor of the fixedWindow primitive. -export function rateLimit( - options?: FixedWindowRateLimitOptions, - ...additionalOptions: FixedWindowRateLimitOptions[] -): Primitive { +export function rateLimit( + options?: FixedWindowRateLimitOptions, + ...additionalOptions: FixedWindowRateLimitOptions[] +): Primitive< + UnionToIntersection> +> { // TODO(#195): We should also have a local rate limit using an in-memory data // structure if the environment supports it return fixedWindow(options, ...additionalOptions); } -export function slidingWindow( - options?: SlidingWindowRateLimitOptions, - ...additionalOptions: SlidingWindowRateLimitOptions[] -): Primitive { +export function slidingWindow< + const Characteristics extends readonly string[] = [], +>( + options?: SlidingWindowRateLimitOptions, + ...additionalOptions: SlidingWindowRateLimitOptions[] +): Primitive< + UnionToIntersection> +> { const rules: ArcjetSlidingWindowRateLimitRule<{}>[] = []; if (typeof options === "undefined") { @@ -682,7 +745,9 @@ export function slidingWindow( for (const opt of [options, ...additionalOptions]) { const mode = opt.mode === "LIVE" ? "LIVE" : "DRY_RUN"; const match = opt.match; - const characteristics = opt.characteristics; + const characteristics = Array.isArray(opt.characteristics) + ? opt.characteristics + : undefined; const max = opt.max; const interval = duration.parse(opt.interval); @@ -867,15 +932,23 @@ export function detectBot( return rules; } -export type ProtectSignupOptions = { - rateLimit?: SlidingWindowRateLimitOptions | SlidingWindowRateLimitOptions[]; +export type ProtectSignupOptions = { + rateLimit?: + | SlidingWindowRateLimitOptions + | SlidingWindowRateLimitOptions[]; bots?: BotOptions | BotOptions[]; email?: EmailOptions | EmailOptions[]; }; -export function protectSignup( - options?: ProtectSignupOptions, -): Product<{ email: string }> { +export function protectSignup( + options?: ProtectSignupOptions, +): Product< + Simplify< + UnionToIntersection< + { email: string } | PropsForCharacteristic + > + > +> { let rateLimitRules: Primitive<{}> = []; if (Array.isArray(options?.rateLimit)) { rateLimitRules = slidingWindow(...options.rateLimit); diff --git a/arcjet/test/index.edge.test.ts b/arcjet/test/index.edge.test.ts index 406cadff9..46df15fa4 100644 --- a/arcjet/test/index.edge.test.ts +++ b/arcjet/test/index.edge.test.ts @@ -36,11 +36,28 @@ describe("Arcjet: Env = Edge runtime", () => { rules: [ // Test rules foobarbaz(), - tokenBucket({ - refillRate: 1, - interval: 1, - capacity: 1, - }), + tokenBucket( + { + characteristics: [ + "ip.src", + "http.host", + "http.method", + "http.request.uri.path", + `http.request.headers["abc"]`, + `http.request.cookie["xyz"]`, + `http.request.uri.args["foobar"]`, + ], + refillRate: 1, + interval: 1, + capacity: 1, + }, + { + characteristics: ["userId"], + refillRate: 1, + interval: 1, + capacity: 1, + }, + ), rateLimit({ max: 1, window: "60s", @@ -61,6 +78,8 @@ describe("Arcjet: Env = Edge runtime", () => { path: "", headers: new Headers(), extra: {}, + userId: "user123", + foobar: 123, }); expect(decision.isErrored()).toBe(false); From 53d220f38551ec1816c6d15cc1a9cddff920baa3 Mon Sep 17 00:00:00 2001 From: Blaine Bublitz Date: Wed, 7 Feb 2024 11:56:43 -0700 Subject: [PATCH 2/3] type tests for required props --- arcjet/index.ts | 12 +++--- arcjet/test/index.node.test.ts | 72 ++++++++++++++++++++++++++++++++++ 2 files changed, 79 insertions(+), 5 deletions(-) diff --git a/arcjet/index.ts b/arcjet/index.ts index 70d28b37d..ebe5ef507 100644 --- a/arcjet/index.ts +++ b/arcjet/index.ts @@ -639,8 +639,10 @@ export function tokenBucket< options?: TokenBucketRateLimitOptions, ...additionalOptions: TokenBucketRateLimitOptions[] ): Primitive< - UnionToIntersection< - { requested: number } | PropsForCharacteristic + Simplify< + UnionToIntersection< + { requested: number } | PropsForCharacteristic + > > > { const rules: ArcjetTokenBucketRateLimitRule<{ requested: number }>[] = []; @@ -682,7 +684,7 @@ export function fixedWindow< options?: FixedWindowRateLimitOptions, ...additionalOptions: FixedWindowRateLimitOptions[] ): Primitive< - UnionToIntersection> + Simplify>> > { const rules: ArcjetFixedWindowRateLimitRule<{}>[] = []; @@ -721,7 +723,7 @@ export function rateLimit( options?: FixedWindowRateLimitOptions, ...additionalOptions: FixedWindowRateLimitOptions[] ): Primitive< - UnionToIntersection> + Simplify>> > { // TODO(#195): We should also have a local rate limit using an in-memory data // structure if the environment supports it @@ -734,7 +736,7 @@ export function slidingWindow< options?: SlidingWindowRateLimitOptions, ...additionalOptions: SlidingWindowRateLimitOptions[] ): Primitive< - UnionToIntersection> + Simplify>> > { const rules: ArcjetSlidingWindowRateLimitRule<{}>[] = []; diff --git a/arcjet/test/index.node.test.ts b/arcjet/test/index.node.test.ts index 1888057cf..7f0c93627 100644 --- a/arcjet/test/index.node.test.ts +++ b/arcjet/test/index.node.test.ts @@ -55,8 +55,47 @@ import arcjet, { fixedWindow, tokenBucket, slidingWindow, + Primitive, } from "../index"; +// Type helpers from https://github.com/sindresorhus/type-fest but adjusted for +// our use. +// +// IsEqual: +// https://github.com/sindresorhus/type-fest/blob/e02f228f6391bb2b26c32a55dfe1e3aa2386d515/source/is-equal.d.ts +// +// Licensed: MIT License Copyright (c) Sindre Sorhus +// (https://sindresorhus.com) +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: The above copyright +// notice and this permission notice shall be included in all copies or +// substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +type IsEqual = (() => G extends A ? 1 : 2) extends () => G extends B + ? 1 + : 2 + ? true + : false; + +// Type testing utilities +type Assert = T; +type Props

= P extends Primitive + ? Props + : never; +type RequiredProps

= IsEqual, E>; + // Instances of Headers contain symbols that may be different depending // on if they have been iterated or not, so we need this equality tester // to only match the items inside the Headers instance. @@ -1978,6 +2017,39 @@ describe("Primitive > tokenBucket", () => { expect(rules[0]).toHaveProperty("capacity", 120); }); + test("can specify user-defined characteristics which are reflected in required props", async () => { + const rules = tokenBucket({ + characteristics: ["userId"], + refillRate: 60, + interval: 60, + capacity: 120, + }); + type Test = Assert< + RequiredProps< + typeof rules, + { requested: number; userId: string | number | boolean } + > + >; + }); + + test("well-known characteristics don't affect the required props", async () => { + const rules = tokenBucket({ + characteristics: [ + "ip.src", + "http.host", + "http.method", + "http.request.uri.path", + `http.request.headers["abc"]`, + `http.request.cookie["xyz"]`, + `http.request.uri.args["foobar"]`, + ], + refillRate: 60, + interval: 60, + capacity: 120, + }); + type Test = Assert>; + }); + test("produces a rules based on single `limit` specified", async () => { const options = { match: "/test", From 1955b3c0b24c0f0ee024157870da85e06a9f6126 Mon Sep 17 00:00:00 2001 From: Blaine Bublitz Date: Wed, 7 Feb 2024 12:47:14 -0700 Subject: [PATCH 3/3] Add type tests to fixedWindow and slidingWindow --- arcjet/test/index.node.test.ts | 56 ++++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/arcjet/test/index.node.test.ts b/arcjet/test/index.node.test.ts index 7f0c93627..19d26b66f 100644 --- a/arcjet/test/index.node.test.ts +++ b/arcjet/test/index.node.test.ts @@ -2224,6 +2224,34 @@ describe("Primitive > fixedWindow", () => { expect(rules[0]).toHaveProperty("max", 1); }); + test("can specify user-defined characteristics which are reflected in required props", async () => { + const rules = fixedWindow({ + characteristics: ["userId"], + window: "1h", + max: 1, + }); + type Test = Assert< + RequiredProps + >; + }); + + test("well-known characteristics don't affect the required props", async () => { + const rules = fixedWindow({ + characteristics: [ + "ip.src", + "http.host", + "http.method", + "http.request.uri.path", + `http.request.headers["abc"]`, + `http.request.cookie["xyz"]`, + `http.request.uri.args["foobar"]`, + ], + window: "1h", + max: 1, + }); + type Test = Assert>; + }); + test("produces a rules based on single `limit` specified", async () => { const options = { match: "/test", @@ -2388,6 +2416,34 @@ describe("Primitive > slidingWindow", () => { expect(rules[0]).toHaveProperty("max", 1); }); + test("can specify user-defined characteristics which are reflected in required props", async () => { + const rules = slidingWindow({ + characteristics: ["userId"], + interval: "1h", + max: 1, + }); + type Test = Assert< + RequiredProps + >; + }); + + test("well-known characteristics don't affect the required props", async () => { + const rules = slidingWindow({ + characteristics: [ + "ip.src", + "http.host", + "http.method", + "http.request.uri.path", + `http.request.headers["abc"]`, + `http.request.cookie["xyz"]`, + `http.request.uri.args["foobar"]`, + ], + interval: "1h", + max: 1, + }); + type Test = Assert>; + }); + test("produces a rules based on single `limit` specified", async () => { const options = { match: "/test",