feat: Leverage local rate limiting if available in environment

analyze/index.ts

@@ -4,6 +4,11 @@ import initWasm, {
   is_valid_email,
   type EmailValidationConfig,
 } from "./wasm/arcjet_analyze_js_req.js";
+import * as rateLimit from "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.js";
+import type {
+  Root as RateLimitAPI,
+  ImportObject as RateLimitImports,
+} from "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.js";
 
 export { type EmailValidationConfig };
 
@@ -163,3 +168,213 @@ export async function detectBot(
     };
   }
 }
+
+function nowInSeconds(): number {
+  return Math.floor(Date.now() / 1000);
+}
+
+class Cache<T> {
+  expires: Map<string, number>;
+  data: Map<string, T>;
+
+  constructor() {
+    this.expires = new Map();
+    this.data = new Map();
+  }
+
+  get(key: string) {
+    const ttl = this.ttl(key);
+    if (ttl > 0) {
+      return this.data.get(key);
+    } else {
+      // Cleanup if expired
+      this.expires.delete(key);
+      this.data.delete(key);
+    }
+  }
+
+  set(key: string, value: T, expireAt: number) {
+    this.expires.set(key, expireAt);
+    this.data.set(key, value);
+  }
+
+  ttl(key: string): number {
+    const now = nowInSeconds();
+    const expiresAt = this.expires.get(key) ?? now;
+    return expiresAt - now;
+  }
+}
+
+const rateLimitCache = new Cache<string>();
+
+const wasmCache = new Map<string, WebAssembly.Module>();
+
+async function moduleFromPath(path: string): Promise<WebAssembly.Module> {
+  const cachedModule = wasmCache.get(path);
+  if (typeof cachedModule !== "undefined") {
+    return cachedModule;
+  }
+
+  if (process.env["NEXT_RUNTIME"] === "edge") {
+    if (path === "arcjet_analyze_bindings_rate_limit.component.core.wasm") {
+      const mod = await import(
+        // @ts-expect-error
+        "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core.wasm?module"
+      );
+      wasmCache.set(path, mod.default);
+      return mod.default;
+    }
+    if (path === "arcjet_analyze_bindings_rate_limit.component.core2.wasm") {
+      const mod = await import(
+        // @ts-expect-error
+        "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core2.wasm?module"
+      );
+      wasmCache.set(path, mod.default);
+      return mod.default;
+    }
+    if (path === "arcjet_analyze_bindings_rate_limit.component.core3.wasm") {
+      const mod = await import(
+        // @ts-expect-error
+        "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core3.wasm?module"
+      );
+      wasmCache.set(path, mod.default);
+      return mod.default;
+    }
+  } else {
+    if (path === "arcjet_analyze_bindings_rate_limit.component.core.wasm") {
+      const { wasm } = await import(
+        "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core.js"
+      );
+      const mod = await WebAssembly.compile(await wasm());
+      wasmCache.set(path, mod);
+      return mod;
+    }
+    if (path === "arcjet_analyze_bindings_rate_limit.component.core2.wasm") {
+      const { wasm } = await import(
+        "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core2.js"
+      );
+      const mod = await WebAssembly.compile(await wasm());
+      wasmCache.set(path, mod);
+      return mod;
+    }
+    if (path === "arcjet_analyze_bindings_rate_limit.component.core3.wasm") {
+      const { wasm } = await import(
+        "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core3.js"
+      );
+      const mod = await WebAssembly.compile(await wasm());
+      wasmCache.set(path, mod);
+      return mod;
+    }
+  }
+
+  throw new Error(`Unknown path: ${path}`);
+}
+
+const rateLimitImports: RateLimitImports = {
+  "arcjet:rate-limit/storage": {
+    get(key) {
+      return rateLimitCache.get(key);
+    },
+    set(key, value, expiresAt) {
+      rateLimitCache.set(key, value, expiresAt);
+    },
+  },
+  "arcjet:rate-limit/time": {
+    now() {
+      return Math.floor(Date.now() / 1000);
+    },
+  },
+};
+
+async function initRateLimit(): Promise<RateLimitAPI | undefined> {
+  try {
+    return rateLimit.instantiate(moduleFromPath, rateLimitImports);
+  } catch {
+    // TODO: Log unsupported wasm error
+  }
+}
+
+export async function tokenBucket(
+  config: {
+    key: string;
+    characteristics?: string[];
+    refill_rate: number;
+    interval: number;
+    capacity: number;
+  },
+  request: unknown,
+): Promise<
+  | {
+      allowed: boolean;
+      max: number;
+      remaining: number;
+      reset: number;
+    }
+  | undefined
+> {
+  const rl = await initRateLimit();
+
+  if (typeof rl !== "undefined") {
+    const configJson = JSON.stringify(config);
+    const requestJson = JSON.stringify(request);
+    const resultJson = rl.tokenBucket(configJson, requestJson);
+    // TODO: Try/catch and Validate
+    return JSON.parse(resultJson);
+  }
+}
+
+export async function fixedWindow(
+  config: {
+    key: string;
+    characteristics?: string[];
+    max: number;
+    window: number;
+  },
+  request: unknown,
+): Promise<
+  | {
+      allowed: boolean;
+      max: number;
+      remaining: number;
+      reset: number;
+    }
+  | undefined
+> {
+  const rl = await initRateLimit();
+
+  if (typeof rl !== "undefined") {
+    const configJson = JSON.stringify(config);
+    const requestJson = JSON.stringify(request);
+    const resultJson = rl.fixedWindow(configJson, requestJson);
+    // TODO: Try/catch and Validate
+    return JSON.parse(resultJson);
+  }
+}
+
+export async function slidingWindow(
+  config: {
+    key: string;
+    characteristics?: string[];
+    max: number;
+    interval: number;
+  },
+  request: unknown,
+): Promise<
+  | {
+      allowed: boolean;
+      max: number;
+      remaining: number;
+      reset: number;
+    }
+  | undefined
+> {
+  const rl = await initRateLimit();
+
+  if (typeof rl !== "undefined") {
+    const configJson = JSON.stringify(config);
+    const requestJson = JSON.stringify(request);
+    const resultJson = rl.slidingWindow(configJson, requestJson);
+    // TODO: Try/catch and Validate
+    return JSON.parse(resultJson);
+  }
+}

analyze/rollup.config.js

@@ -3,29 +3,35 @@ import { createConfig } from "@arcjet/rollup-config";
 export default createConfig(import.meta.url, {
   plugins: [
     {
-      name: "externalize-wasm",
-      resolveId(source) {
+      name: "externalize-wasm-js",
+      resolveId(id) {
         // We need to externalize this because otherwise some JS that edge can't
         // understand gets included in the bundle
-        if (source === "./wasm/arcjet_analyze_js_req.js") {
-          return {
-            id: "./wasm/arcjet_analyze_js_req.js",
-            external: true,
-          };
+        if (id === "./wasm/arcjet_analyze_js_req.js") {
+          return { id, external: true };
         }
-        if (source === "./wasm/arcjet_analyze_js_req_bg.wasm?module") {
-          return {
-            id: "./wasm/arcjet_analyze_js_req_bg.wasm?module",
-            external: true,
-          };
+        // TODO: Generation of the below files could be handled via rollup
+        // plugin so we wouldn't need to externalize here
+        if (id === "./wasm/arcjet.wasm.js") {
+          return { id, external: true };
         }
-        // TODO: Generation of this file can be handled via rollup plugin so we
-        // wouldn't need to externalize here
-        if (source === "./wasm/arcjet.wasm.js") {
-          return {
-            id: "./wasm/arcjet.wasm.js",
-            external: true,
-          };
+        if (
+          id ===
+          "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core.js"
+        ) {
+          return { id, external: true };
+        }
+        if (
+          id ===
+          "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core2.js"
+        ) {
+          return { id, external: true };
+        }
+        if (
+          id ===
+          "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core3.js"
+        ) {
+          return { id, external: true };
         }
         return null;
       },

analyze/wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core.d.ts

@@ -0,0 +1,31 @@
+// @generated by wasm2module - DO NOT EDIT
+/* tslint:disable */
+/* eslint-disable */
+
+/**
+ * This file contains an Arcjet Wasm binary inlined as a base64
+ * [Data URL](https://developer.mozilla.org/en-US/docs/Web/HTTP/Basics_of_HTTP/Data_URLs)
+ * with the application/wasm MIME type.
+ *
+ * This was chosen to save on storage space over inlining the file directly as
+ * a Uint8Array, which would take up ~3x the space of the Wasm file. See
+ * https://blobfolio.com/2019/better-binary-batter-mixing-base64-and-uint8array/
+ * for more details.
+ *
+ * It is then decoded into an ArrayBuffer to be used directly via WebAssembly's
+ * `compile()` function in our entry point file.
+ *
+ * This is all done to avoid trying to read or bundle the Wasm asset in various
+ * ways based on the platform or bundler a user is targeting. One example being
+ * that Next.js requires special `asyncWebAssembly` webpack config to load our
+ * Wasm file if we don't do this.
+ *
+ * In the future, we hope to do away with this workaround when all bundlers
+ * properly support consistent asset bundling techniques.
+ */
+
+/**
+ * Returns an ArrayBuffer for the Arcjet Wasm binary, decoded from a base64 Data
+ * URL.
+ */
+export function wasm(): Promise<ArrayBuffer>;

analyze/wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core2.d.ts

@@ -0,0 +1,31 @@
+// @generated by wasm2module - DO NOT EDIT
+/* tslint:disable */
+/* eslint-disable */
+
+/**
+ * This file contains an Arcjet Wasm binary inlined as a base64
+ * [Data URL](https://developer.mozilla.org/en-US/docs/Web/HTTP/Basics_of_HTTP/Data_URLs)
+ * with the application/wasm MIME type.
+ *
+ * This was chosen to save on storage space over inlining the file directly as
+ * a Uint8Array, which would take up ~3x the space of the Wasm file. See
+ * https://blobfolio.com/2019/better-binary-batter-mixing-base64-and-uint8array/
+ * for more details.
+ *
+ * It is then decoded into an ArrayBuffer to be used directly via WebAssembly's
+ * `compile()` function in our entry point file.
+ *
+ * This is all done to avoid trying to read or bundle the Wasm asset in various
+ * ways based on the platform or bundler a user is targeting. One example being
+ * that Next.js requires special `asyncWebAssembly` webpack config to load our
+ * Wasm file if we don't do this.
+ *
+ * In the future, we hope to do away with this workaround when all bundlers
+ * properly support consistent asset bundling techniques.
+ */
+
+/**
+ * Returns an ArrayBuffer for the Arcjet Wasm binary, decoded from a base64 Data
+ * URL.
+ */
+export function wasm(): Promise<ArrayBuffer>;

analyze/wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core2.js

@@ -0,0 +1,38 @@
+// @generated by wasm2module - DO NOT EDIT
+/* eslint-disable */
+// @ts-nocheck
+
+/**
+ * This file contains an Arcjet Wasm binary inlined as a base64
+ * [Data URL](https://developer.mozilla.org/en-US/docs/Web/HTTP/Basics_of_HTTP/Data_URLs)
+ * with the application/wasm MIME type.
+ *
+ * This was chosen to save on storage space over inlining the file directly as
+ * a Uint8Array, which would take up ~3x the space of the Wasm file. See
+ * https://blobfolio.com/2019/better-binary-batter-mixing-base64-and-uint8array/
+ * for more details.
+ *
+ * It is then decoded into an ArrayBuffer to be used directly via WebAssembly's
+ * `compile()` function in our entry point file.
+ *
+ * This is all done to avoid trying to read or bundle the Wasm asset in various
+ * ways based on the platform or bundler a user is targeting. One example being
+ * that Next.js requires special `asyncWebAssembly` webpack config to load our
+ * Wasm file if we don't do this.
+ *
+ * In the future, we hope to do away with this workaround when all bundlers
+ * properly support consistent asset bundling techniques.
+ */
+
+const wasmBase64 = "data:application/wasm;base64,AGFzbQEAAAABEAJgA39/fwBgBn9/f39/fwADAwIAAQQFAXABAgIHFAMBMAAAATEAAQgkaW1wb3J0cwEACiMCDQAgACABIAJBABEAAAsTACAAIAEgAiADIAQgBUEBEQEACwAuCXByb2R1Y2VycwEMcHJvY2Vzc2VkLWJ5AQ13aXQtY29tcG9uZW50BjAuMjAuMABtBG5hbWUAExJ3aXQtY29tcG9uZW50OnNoaW0BUQIAJmluZGlyZWN0LWFyY2pldDpyYXRlLWxpbWl0L3N0b3JhZ2UtZ2V0ASZpbmRpcmVjdC1hcmNqZXQ6cmF0ZS1saW1pdC9zdG9yYWdlLXNldA==";
+/**
+ * Returns an ArrayBuffer for the Arcjet Wasm binary, decoded from a base64 Data
+ * URL.
+ */
+// TODO: Switch back to top-level await when our platforms all support it
+export async function wasm() {
+  // This uses fetch to decode the wasm data url
+  const wasmDecode = await fetch(wasmBase64);
+  // And then we return it as an ArrayBuffer
+  return wasmDecode.arrayBuffer();
+}