feat: Leverage local rate limiting if available in environment

Author: blaine-arcjet

analyze/index.ts

@@ -4,6 +4,7 @@ import initWasm, {
   is_valid_email,
   type EmailValidationConfig,
 } from "./wasm/arcjet_analyze_js_req.js";
+import { instantiate } from "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.js";
 
 export { type EmailValidationConfig };
 
@@ -163,3 +164,114 @@ 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, ttl: number) {
+    this.expires.set(key, nowInSeconds() + ttl);
+    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>();
+
+export async function fixedWindow(
+  config: {
+    key: string;
+    characteristics?: string[];
+    max: number;
+    window: number;
+  },
+  request: unknown,
+): Promise<{
+  allowed: boolean;
+  max: number;
+  remaining: number;
+  reset: number;
+}> {
+  const configJson = JSON.stringify(config);
+  const requestJson = JSON.stringify(request);
+
+  const abc = await instantiate(
+    async function (path: string) {
+      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"
+          );
+          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"
+          );
+          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"
+          );
+          return mod.default;
+        }
+      } else {
+        // const { wasm } = await import("./wasm/arcjet.wasm.js");
+        // wasmModule = await WebAssembly.compile(await wasm());
+        return Promise.reject("TODO");
+      }
+    },
+    {
+      "arcjet:rate-limit/storage": {
+        get(key) {
+          return rateLimitCache.get(key);
+        },
+        set(key, value, ttl) {
+          rateLimitCache.set(key, value, ttl);
+        },
+      },
+      "arcjet:rate-limit/time": {
+        now() {
+          return nowInSeconds();
+        },
+      },
+    },
+  );
+
+  const resultJson = abc.fixedWindow(configJson, requestJson);
+  // TODO: Try/catch and Validate
+  return JSON.parse(resultJson);
+}

analyze/rollup.config.js

@@ -19,6 +19,33 @@ export default createConfig(import.meta.url, {
             external: true,
           };
         }
+        if (
+          source ===
+          "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core.wasm?module"
+        ) {
+          return {
+            id: "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core.wasm?module",
+            external: true,
+          };
+        }
+        if (
+          source ===
+          "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core2.wasm?module"
+        ) {
+          return {
+            id: "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core2.wasm?module",
+            external: true,
+          };
+        }
+        if (
+          source ===
+          "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core3.wasm?module"
+        ) {
+          return {
+            id: "./wasm/rate-limit/arcjet_analyze_bindings_rate_limit.component.core3.wasm?module",
+            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") {

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

@@ -0,0 +1,37 @@
+import { ArcjetRateLimitStorage } from './interfaces/arcjet-rate-limit-storage.js';
+import { ArcjetRateLimitTime } from './interfaces/arcjet-rate-limit-time.js';
+export interface ImportObject {
+  'arcjet:rate-limit/storage': typeof ArcjetRateLimitStorage,
+  'arcjet:rate-limit/time': typeof ArcjetRateLimitTime,
+}
+export interface Root {
+  tokenBucket(config: string, request: string): string,
+  fixedWindow(config: string, request: string): string,
+  slidingWindow(config: string, request: string): string,
+}
+
+/**
+* Instantiates this component with the provided imports and
+* returns a map of all the exports of the component.
+*
+* This function is intended to be similar to the
+* `WebAssembly.instantiate` function. The second `imports`
+* argument is the "import object" for wasm, except here it
+* uses component-model-layer types instead of core wasm
+* integers/numbers/etc.
+*
+* The first argument to this function, `getCoreModule`, is
+* used to compile core wasm modules within the component.
+* Components are composed of core wasm modules and this callback
+* will be invoked per core wasm module. The caller of this
+* function is responsible for reading the core wasm module
+* identified by `path` and returning its compiled
+* `WebAssembly.Module` object. This would use `compileStreaming`
+* on the web, for example.
+*/
+export function instantiate(
+getCoreModule: (path: string) => Promise<WebAssembly.Module>,
+imports: ImportObject,
+instantiateCore?: (module: WebAssembly.Module, imports: Record<string, any>) => Promise<WebAssembly.Instance>
+): Promise<Root>;
+