From dc5b0010dd772207ec662062bfa6da5fe712f987 Mon Sep 17 00:00:00 2001 From: blaine-arcjet <146491715+blaine-arcjet@users.noreply.github.com> Date: Wed, 7 Feb 2024 13:09:13 -0700 Subject: [PATCH] feat: Allow user-defined characteristics on rate limit options (#203) Closes #202 Closes https://github.com/arcjet/arcjet/issues/597 This adds TypeScript support for custom characteristics. While it was supported via the protocol, there was no indication to the user that they needed to add the field to every request. By changing the types of our primitives, we can make the prop required on the request object (while filtering our well-known characteristics). Draft because it needs some tests. --- arcjet/index.ts | 145 +++++++++++++++++++++++++-------- arcjet/test/index.edge.test.ts | 29 +++++-- arcjet/test/index.node.test.ts | 128 +++++++++++++++++++++++++++++ 3 files changed, 262 insertions(+), 40 deletions(-) diff --git a/arcjet/index.ts b/arcjet/index.ts index 7aa8acafd..ebe5ef507 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,18 @@ 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< + Simplify< + UnionToIntersection< + { requested: number } | PropsForCharacteristic + > + > +> { const rules: ArcjetTokenBucketRateLimitRule<{ requested: number }>[] = []; if (typeof options === "undefined") { @@ -603,7 +654,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 +678,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< + Simplify>> +> { const rules: ArcjetFixedWindowRateLimitRule<{}>[] = []; if (typeof options === "undefined") { @@ -638,7 +695,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 +719,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< + Simplify>> +> { // 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< + Simplify>> +> { const rules: ArcjetSlidingWindowRateLimitRule<{}>[] = []; if (typeof options === "undefined") { @@ -682,7 +747,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 +934,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); diff --git a/arcjet/test/index.node.test.ts b/arcjet/test/index.node.test.ts index 1888057cf..19d26b66f 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", @@ -2152,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", @@ -2316,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",