Skip to content

Commit

Permalink
feat: Leverage local rate limiting if available in environment
Browse files Browse the repository at this point in the history
  • Loading branch information
blaine-arcjet committed Feb 7, 2024
1 parent de36130 commit c354c30
Show file tree
Hide file tree
Showing 13 changed files with 538 additions and 17 deletions.
112 changes: 112 additions & 0 deletions analyze/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 };

Expand Down Expand Up @@ -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);
}
27 changes: 27 additions & 0 deletions analyze/rollup.config.js
Original file line number Diff line number Diff line change
Expand Up @@ -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") {
Expand Down
Binary file not shown.
Binary file not shown.
Binary file not shown.
Original file line number Diff line number Diff line change
@@ -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>;

Loading

0 comments on commit c354c30

Please sign in to comment.