From cb6191902e6b528d12a1160de59683c16eef58ea Mon Sep 17 00:00:00 2001 From: MythEclipse Date: Fri, 19 Jun 2026 18:10:49 +0700 Subject: [PATCH] fix: critical bugs, serverless stability, and security hardening Bug Fixes: - Fix rate limiter API mismatch: check() -> checkAsync() in index.ts, worker.ts, api/relay.ts - Fix WebSocket SSRF silent drop: return error Response instead of undefined - Fix isDevMode() default: changed from true to false (production-safe) - Fix process.env -> env bindings in worker.ts requireAuth for Cloudflare Workers - Fix Bun.file() crash in Workers: add try/catch with fallback - Fix Bun.CryptoHasher -> Web Crypto API in mimo-auth.ts for Workers compat Architecture: - Add public methods to ProxyPool (getEntryAtIndex, getProxyUrlAtIndex, getCurrentIndex, setCurrentIndex) to remove all 'as any' casts in SessionProxyPool - Add addProxy() method for manual proxy management - Add loadAsync(), tryLoadAsync(), loadFromString() to ProxyPool Serverless Stability: - Add optional DNS rebinding protection via SSRF_DNS_CHECK env flag - CORS cache now auto-invalidates when CORS_ORIGIN env changes - Rate limiter max-size eviction (10k keys) prevents unbounded memory growth Tests: - Fix type assertions in test files (body as Record) - All 153 tests pass, typecheck clean --- api/relay.ts | 2 +- src/index.test.ts | 2 +- src/index.ts | 26 +- src/lib/fetch-utils.ts | 14 +- src/lib/mimo-auth.ts | 14 +- src/lib/proxy-pool.ts | 144 ++++++-- src/lib/relay-utils.test.ts | 6 +- src/lib/relay-utils.ts | 30 +- src/lib/router.ts | 553 ++++++++++++++++++++++++++++ src/middleware/body-limiter.test.ts | 2 +- src/middleware/rate-limiter.test.ts | 56 +-- src/middleware/rate-limiter.ts | 105 ++++-- src/worker.ts | 32 +- 13 files changed, 865 insertions(+), 121 deletions(-) create mode 100644 src/lib/router.ts diff --git a/api/relay.ts b/api/relay.ts index 979d4a0..553b513 100644 --- a/api/relay.ts +++ b/api/relay.ts @@ -252,7 +252,7 @@ async function handleRelay(req: Request): Promise { } // ── Middleware: Rate limiting ──────────────────────────────────── - const rateCheck = rateLimiter.check(clientIP); + const rateCheck = await rateLimiter.checkAsync(clientIP); if (!rateCheck.allowed) { logRelayEvent({ method, diff --git a/src/index.test.ts b/src/index.test.ts index 9165eba..16b27ea 100644 --- a/src/index.test.ts +++ b/src/index.test.ts @@ -25,7 +25,7 @@ describe("handleHealth", () => { test("should include status, uptime, and version fields", async () => { const response = handleHealth(); - const body = await response.json(); + const body = await response.json() as Record; expect(body.status).toBe("ok"); expect(typeof body.uptime).toBe("number"); expect(body.version).toBe("1.0.0"); diff --git a/src/index.ts b/src/index.ts index 9d7b930..c163ade 100644 --- a/src/index.ts +++ b/src/index.ts @@ -24,6 +24,7 @@ import { createErrorResponse, createCorsPreflightResponse, getCorsHeaders, + setSsrfDnsCheck, } from "./lib/relay-utils"; import { checkBodySize } from "./middleware/body-limiter"; @@ -86,6 +87,13 @@ proxyPool.tryLoad( const sessionPool = new SessionProxyPool(proxyPool); sessionPool.setFailureThreshold(3); +// --- SSRF DNS rebinding protection -------------------------------------------- + +if (process.env.SSRF_DNS_CHECK === "true") { + setSsrfDnsCheck(true); + console.log("[relay] SSRF DNS rebinding protection enabled"); +} + // --- WebSocket relay data type ----------------------------------------------- interface WSRelayData { @@ -302,7 +310,7 @@ async function handleRelay( } // -- Middleware: Rate limiting --------------------------------------------- - const rateCheck = rateLimiter.check(clientIP); + const rateCheck = await rateLimiter.checkAsync(clientIP); if (!rateCheck.allowed) { logRelayEvent({ method, @@ -438,8 +446,20 @@ function handleWebSocketUpgrade( const relayPath = req.headers.get("x-relay-path") ?? "/"; const normalized = normalizeTargetUrl(target, relayPath); - if (!normalized) return undefined; - if (!isAllowedTarget(new URL(normalized.toString()))) return undefined; + if (!normalized) { + return createErrorResponse({ + code: "INVALID_TARGET", + status: 400, + message: "Missing or invalid x-relay-target header", + }); + } + if (!isAllowedTarget(new URL(normalized.toString()))) { + return createErrorResponse({ + code: "SSRF_BLOCKED", + status: 403, + message: "Target domain not allowed", + }); + } const targetUrl = normalized.toString(); diff --git a/src/lib/fetch-utils.ts b/src/lib/fetch-utils.ts index 13f5f8c..1680170 100644 --- a/src/lib/fetch-utils.ts +++ b/src/lib/fetch-utils.ts @@ -28,14 +28,12 @@ export function closeAllActiveReaders(): void { /** * Returns `true` when development features (HMR, verbose console) should * be enabled. Controlled by the `NODE_ENV` / `BUN_ENV` env var — defaults - * to `true` for convenience during local development. - * - * Set `NODE_ENV=production` or `BUN_ENV=production` to disable. + * to `false` (production-safe). Set `NODE_ENV=development` to enable. */ export function isDevMode(): boolean { const env = (process.env.NODE_ENV ?? process.env.BUN_ENV ?? "").toLowerCase(); - if (env === "production") return false; - return true; + if (env === "development" || env === "dev") return true; + return false; } // ─── SSE line buffer (fixes chunk-boundary corruption) ─────────────────── @@ -51,6 +49,7 @@ export function isDevMode(): boolean { */ export class SSELineBuffer { private buffer = ""; + private readonly MAX_BUFFER_SIZE = 1024 * 1024; // 1MB limit to prevent OOM /** * Feed a chunk of decoded text and return complete lines. @@ -58,6 +57,11 @@ export class SSELineBuffer { */ add(chunk: string): string[] { this.buffer += chunk; + + if (this.buffer.length > this.MAX_BUFFER_SIZE) { + throw new Error(`SSELineBuffer exceeded maximum size of ${this.MAX_BUFFER_SIZE} bytes. Stream may be malicious or corrupted.`); + } + if (!this.buffer.includes("\n")) return []; const parts = this.buffer.split("\n"); diff --git a/src/lib/mimo-auth.ts b/src/lib/mimo-auth.ts index 0290911..3f884bf 100644 --- a/src/lib/mimo-auth.ts +++ b/src/lib/mimo-auth.ts @@ -29,8 +29,10 @@ const EXPIRY_BUFFER_MS = 300_000; // 5 minutes * * Format: `sha256(hostname|platform|arch|cpu|username)` * This matches the 9Router reference implementation. + * + * Uses Web Crypto API (available in Bun, Workers, and Node.js 20+). */ -function generateDeviceFingerprint(): string { +async function generateDeviceFingerprint(): Promise { const hostname = os.hostname(); const platform = process.platform; const arch = process.arch; @@ -39,9 +41,11 @@ function generateDeviceFingerprint(): string { const username = process.env.USER ?? process.env.USERNAME ?? "unknown"; const raw = `${hostname}|${platform}|${arch}|${cpuModel}|${username}`; - const hasher = new Bun.CryptoHasher("sha256"); - hasher.update(raw); - return hasher.digest("hex") as string; + const encoder = new TextEncoder(); + const data = encoder.encode(raw); + const hashBuffer = await crypto.subtle.digest("SHA-256", data); + const hashArray = new Uint8Array(hashBuffer); + return Array.from(hashArray).map((b) => b.toString(16).padStart(2, "0")).join(""); } // --- JWT bootstrap ----------------------------------------------------------- @@ -55,7 +59,7 @@ function generateDeviceFingerprint(): string { * @throws If the bootstrap request fails or returns an unexpected response. */ async function bootstrapJwt(): Promise { - const fingerprint = generateDeviceFingerprint(); + const fingerprint = await generateDeviceFingerprint(); const resp = await fetch(MIMO_BOOTSTRAP_URL, { method: "POST", diff --git a/src/lib/proxy-pool.ts b/src/lib/proxy-pool.ts index 42168ee..68f8e8b 100644 --- a/src/lib/proxy-pool.ts +++ b/src/lib/proxy-pool.ts @@ -59,6 +59,45 @@ export class ProxyPool { } const text = readFileSync(filePath, "utf-8"); + this.parseProxies(text, filePath); + } + + /** + * Async load proxies from a file at `filePath`. + * Uses Bun.file() when available, falls back to sync read. + */ + async loadAsync(filePath: string): Promise { + try { + // Try Bun.file first (Bun runtime) + if (typeof globalThis.Bun !== "undefined") { + const file = (globalThis as any).Bun.file(filePath); + const exists = await file.exists(); + if (!exists) { + console.warn(`[proxy-pool] File not found: ${filePath}`); + return; + } + const text = await file.text(); + this.parseProxies(text, filePath); + return; + } + } catch { + // Fall through to sync + } + // Fallback to sync read + this.load(filePath); + } + + /** + * Load proxies from a comma-separated string (for serverless envs). + * Format: "host1:port1:user1:pass1,host2:port2:user2:pass2" + */ + loadFromString(proxyList: string): void { + if (!proxyList) return; + const lines = proxyList.split(",").map((l) => l.trim()).filter(Boolean); + this.parseProxies(lines.join("\n"), "env:PROXY_LIST"); + } + + private parseProxies(text: string, source: string): void { const lines = text .split("\n") .map((l: string) => l.trim()) @@ -87,7 +126,7 @@ export class ProxyPool { this.currentIndex = 0; this.failures.clear(); - logPool(`loaded ${this.proxies.length} proxies from ${filePath}`); + logPool(`loaded ${this.proxies.length} proxies from ${source}`); } /** @@ -99,6 +138,14 @@ export class ProxyPool { return this.proxies.length > 0; } + /** + * Async convenience -- load from a path or skip. + */ + async tryLoadAsync(filePath?: string): Promise { + if (filePath) await this.loadAsync(filePath); + return this.proxies.length > 0; + } + // -- Access ------------------------------------------------------------------ /** Total number of proxies in the pool. */ @@ -137,6 +184,54 @@ export class ProxyPool { return `http://${auth}${entry.host}:${entry.port}`; } + // -- Manual proxy management ------------------------------------------------ + + /** + * Add a single proxy to the pool. + * Format: "host:port:username:password" (username:password optional) + */ + addProxy(proxyStr: string): void { + const parts = proxyStr.trim().split(":"); + if (parts.length < 2) return; + + const host = parts[0]!; + const port = Number.parseInt(parts[1]!, 10); + if (!Number.isFinite(port)) return; + + this.proxies.push({ + host, + port, + username: parts[2] ?? "", + password: parts.slice(3).join(":") ?? "", + }); + + logPool(`added proxy ${host}:${port} (total: ${this.proxies.length})`); + } + + // -- Public accessors (for SessionProxyPool) -------------------------------- + + /** Get the ProxyEntry at a given index. Returns null if out of bounds. */ + getEntryAtIndex(index: number): ProxyEntry | null { + return this.proxies[index] ?? null; + } + + /** Build proxy URL string by index. Returns null if out of bounds. */ + getProxyUrlAtIndex(index: number): string | null { + const entry = this.getEntryAtIndex(index); + if (!entry) return null; + return this.formatProxyUrl(entry); + } + + /** Get the current rotation index. */ + getCurrentIndex(): number { + return this.currentIndex; + } + + /** Set the rotation index (for session pool's direct manipulation). */ + setCurrentIndex(index: number): void { + this.currentIndex = index; + } + // -- Rotation ---------------------------------------------------------------- /** @@ -342,7 +437,7 @@ export class SessionProxyPool { const existing = this.sessions.get(sessionId); if (existing !== undefined) { logPool(`acquire existing session=${sessionId.slice(0, 8)} proxyIndex=${existing.proxyIndex}`); - return this.formatProxyUrlAtIndex(existing.proxyIndex); + return this.pool.getProxyUrlAtIndex(existing.proxyIndex); } const index = this.pickLeastUsedIndex(model); @@ -357,18 +452,18 @@ export class SessionProxyPool { } usedBy.add(sessionId); - const entry = this.poolEntryAtIndex(index); + const entry = this.pool.getEntryAtIndex(index); logPool(`acquire session=${sessionId.slice(0, 8)} -> proxyIndex=${index} host=${entry?.host} activeSessions=${this.activeSessions}`); - return this.formatProxyUrlAtIndex(index); + return this.pool.getProxyUrlAtIndex(index); } /** Return the current proxy URL for a session (no rotation), or null. */ getProxyUrl(sessionId: string): string | null { const info = this.sessions.get(sessionId); if (!info) return null; - const entry = this.poolEntryAtIndex(info.proxyIndex); + const entry = this.pool.getEntryAtIndex(info.proxyIndex); logPool(`getProxyUrl session=${sessionId.slice(0, 8)} proxyIndex=${info.proxyIndex} host=${entry?.host}`); - return this.formatProxyUrlAtIndex(info.proxyIndex); + return this.pool.getProxyUrlAtIndex(info.proxyIndex); } /** Remove a session from all tracking. */ @@ -383,7 +478,7 @@ export class SessionProxyPool { } this.sessions.delete(sessionId); - const entry = this.poolEntryAtIndex(info.proxyIndex); + const entry = this.pool.getEntryAtIndex(info.proxyIndex); logPool(`release session=${sessionId.slice(0, 8)} proxyIndex=${info.proxyIndex} host=${entry?.host} activeSessions=${this.activeSessions}`); } @@ -405,13 +500,13 @@ export class SessionProxyPool { if (info.failures < this.failureThreshold) return false; const oldIndex = info.proxyIndex; - const oldEntry = this.poolEntryAtIndex(oldIndex); + const oldEntry = this.pool.getEntryAtIndex(oldIndex); - // Mark in underlying pool (public API only works on currentIndex) - const savedIdx = (this.pool as any).currentIndex as number; - (this.pool as any).currentIndex = oldIndex; + // Mark in underlying pool using public API + const savedIdx = this.pool.getCurrentIndex(); + this.pool.setCurrentIndex(oldIndex); this.pool.markFailed(this.failureThreshold); - (this.pool as any).currentIndex = savedIdx; + this.pool.setCurrentIndex(savedIdx); // Remove session from old proxy usage tracking const usedBy = this.proxyUsage.get(oldIndex); @@ -438,7 +533,7 @@ export class SessionProxyPool { } newUsedBy.add(sessionId); - const newEntry = this.poolEntryAtIndex(newIndex); + const newEntry = this.pool.getEntryAtIndex(newIndex); logPool(`markFailed rotated session=${sessionId.slice(0, 8)} ${oldEntry?.host} -> ${newEntry?.host}`); return true; } @@ -454,7 +549,7 @@ export class SessionProxyPool { if (!info) return false; const oldIndex = info.proxyIndex; - const oldEntry = this.poolEntryAtIndex(oldIndex); + const oldEntry = this.pool.getEntryAtIndex(oldIndex); // Remove from old proxy usage first const usedBy = this.proxyUsage.get(oldIndex); @@ -470,7 +565,7 @@ export class SessionProxyPool { let bestCount = Infinity; for (let step = 1; step <= this.pool.size; step++) { const i = (oldIndex + step) % this.pool.size; - if (model && (this.pool as any).isIndexInCooldown(i, model)) continue; + if (model && this.pool.isIndexInCooldown(i, model)) continue; const count = this.proxyUsage.get(i)?.size ?? 0; if (count < bestCount) { bestCount = count; @@ -493,7 +588,7 @@ export class SessionProxyPool { } newUsedBy.add(sessionId); - const newEntry = this.poolEntryAtIndex(bestIndex); + const newEntry = this.pool.getEntryAtIndex(bestIndex); logPool(`rotateNow session=${sessionId.slice(0, 8)} ${oldEntry?.host} -> ${newEntry?.host}`); return true; } @@ -514,27 +609,22 @@ export class SessionProxyPool { markRateLimited(sessionId: string, model: string): void { const info = this.sessions.get(sessionId); if (!info) return; - const savedIdx = (this.pool as any).currentIndex as number; - (this.pool as any).currentIndex = info.proxyIndex; + const savedIdx = this.pool.getCurrentIndex(); + this.pool.setCurrentIndex(info.proxyIndex); this.pool.markRateLimited(model); - (this.pool as any).currentIndex = savedIdx; + this.pool.setCurrentIndex(savedIdx); } // -- Internals ---------------------------------------------------------------- /** Get the ProxyEntry at a given index. Forward reference to local type. */ private poolEntryAtIndex(index: number): ProxyEntry | null { - return (this.pool as any).proxies[index] ?? null; + return this.pool.getEntryAtIndex(index); } /** Build proxy URL string by index. */ private formatProxyUrlAtIndex(index: number): string | null { - const entry = this.poolEntryAtIndex(index); - if (!entry) return null; - const auth = entry.username - ? `${encodeURIComponent(entry.username)}:${encodeURIComponent(entry.password)}@` - : ""; - return `http://${auth}${entry.host}:${entry.port}`; + return this.pool.getProxyUrlAtIndex(index); } /** Return the index of the proxy with the fewest active sessions, or -1. */ @@ -545,7 +635,7 @@ export class SessionProxyPool { let bestCount = Infinity; for (let i = 0; i < this.pool.size; i++) { - if (model && (this.pool as any).isIndexInCooldown(i, model)) continue; + if (model && this.pool.isIndexInCooldown(i, model)) continue; const count = this.proxyUsage.get(i)?.size ?? 0; logPool(`pickLeastUsed proxy[${i}] count=${count}`); if (count < bestCount) { diff --git a/src/lib/relay-utils.test.ts b/src/lib/relay-utils.test.ts index 5877775..c66e463 100644 --- a/src/lib/relay-utils.test.ts +++ b/src/lib/relay-utils.test.ts @@ -692,7 +692,7 @@ describe("createErrorResponse", () => { }); expect(result.status).toBe(504); expect(result.headers.get("Content-Type")).toBe("application/json"); - const body = await result.json(); + const body = await result.json() as Record; expect(body.error).toBe(true); expect(body.code).toBe("TIMEOUT"); expect(body.message).toBe("Upstream timed out"); @@ -714,7 +714,7 @@ describe("createErrorResponse", () => { message: "DNS resolution failed", }); expect(result.status).toBe(502); - const body = await result.json(); + const body = await result.json() as Record; expect(body.code).toBe("DNS_FAILURE"); }); }); @@ -849,7 +849,7 @@ describe("integration: full relay flow", () => { const classified = classifyFetchError(error); const response = createErrorResponse(classified); expect(response.status).toBe(504); - const body = await response.json(); + const body = await response.json() as Record; expect(body.code).toBe("TIMEOUT"); expect(body.error).toBe(true); }); diff --git a/src/lib/relay-utils.ts b/src/lib/relay-utils.ts index c16f542..c4ecf81 100644 --- a/src/lib/relay-utils.ts +++ b/src/lib/relay-utils.ts @@ -29,6 +29,23 @@ export class RelayError extends Error { } } +// --- SSRF DNS check configuration --------------------------------------------- + +/** + * When true, the relay path will resolve DNS and verify no resolved IP is + * private/link-local. This protects against DNS rebinding attacks but adds + * latency (DNS lookup per request). Enable via SSRF_DNS_CHECK=true env var. + */ +let ssrfDnsCheckEnabled = false; + +export function setSsrfDnsCheck(enabled: boolean): void { + ssrfDnsCheckEnabled = enabled; +} + +export function isSsrfDnsCheckEnabled(): boolean { + return ssrfDnsCheckEnabled; +} + // --- CORS configuration ------------------------------------------------------- /** Return CORS headers. Origin defaults to "*" but can be overridden via env. */ @@ -39,20 +56,27 @@ export function getAllowedOrigin(): string { } let cachedCorsHeaders: Record | null = null; +let cachedCorsOrigin: string | null = null; /** Rebuild the CORS headers map (call after changing origin at runtime). */ export function rebuildCorsHeaders(): Record { + const origin = getAllowedOrigin(); cachedCorsHeaders = { - "Access-Control-Allow-Origin": getAllowedOrigin(), + "Access-Control-Allow-Origin": origin, "Access-Control-Allow-Methods": "GET, POST, PUT, DELETE, PATCH, OPTIONS", "Access-Control-Allow-Headers": "*", }; + cachedCorsOrigin = origin; return cachedCorsHeaders; } -/** Return CORS headers with the configured origin (cached after first call). */ +/** Return CORS headers with the configured origin (cached, auto-invalidates on env change). */ export function getCorsHeaders(): Record { - if (!cachedCorsHeaders) rebuildCorsHeaders(); + const currentOrigin = getAllowedOrigin(); + // Invalidate cache if CORS_ORIGIN env changed + if (!cachedCorsHeaders || cachedCorsOrigin !== currentOrigin) { + rebuildCorsHeaders(); + } return cachedCorsHeaders!; } diff --git a/src/lib/router.ts b/src/lib/router.ts new file mode 100644 index 0000000..ac01aa5 --- /dev/null +++ b/src/lib/router.ts @@ -0,0 +1,553 @@ +import { + normalizeTargetUrl, + isAllowedTarget, + isAllowedTargetAsync, + isSsrfDnsCheckEnabled, + filterRequestHeaders, + buildRelayRequest, + createRelayResponse, + createErrorResponse, + createCorsPreflightResponse, + getCorsHeaders, + classifyFetchError, +} from "./relay-utils"; + +import { checkBodySize } from "../middleware/body-limiter"; +import { createRateLimiter } from "../middleware/rate-limiter"; +import { logRelayEvent } from "../middleware/logger"; +import { ProxyPool, SessionProxyPool } from "./proxy-pool"; +import { handleChatCompletion, listModels } from "./ai-proxy"; +import { handleAnthropicMessages } from "./anthropic-proxy"; +import { fetchWithRetry } from "./fetch-utils"; + +// --- Types ------------------------------------------------------------------- + +export interface RouterEnv { + PORT?: string; + RELAY_TIMEOUT_MS?: string; + BODY_MAX_BYTES?: string; + RATE_LIMIT_MAX?: string; + RATE_LIMIT_WINDOW_MS?: string; + CORS_ORIGIN?: string; + NODE_ENV?: string; + API_KEY?: string; + PROXY_LIST?: string; // Comma-separated list of proxies for serverless + // Optional KV binding for rate limiter + KV?: { + get(key: string): Promise; + put(key: string, value: any, options?: { expirationTtl?: number }): Promise; + }; +} + +// --- Global singletons (survives warm starts) -------------------------------- + +let rateLimiter: ReturnType | null = null; +let proxyPool: ProxyPool | null = null; +let sessionPool: SessionProxyPool | null = null; + +const SERVER_START_TIME = Date.now(); +const RELAY_VERSION = "1.0.0"; + +// --- Helpers ----------------------------------------------------------------- + +function getNumericEnv(env: RouterEnv, key: keyof RouterEnv, fallback: number): number { + const raw = env[key]; + const val = typeof raw === "string" ? raw : (typeof process !== "undefined" ? process.env[key as string] : undefined); + return Number.parseInt(val ?? String(fallback), 10); +} + +function getEnv(env: RouterEnv, key: keyof RouterEnv, fallback: string): string { + const raw = env[key]; + const fromEnv = typeof raw === "string" ? raw : (typeof process !== "undefined" ? process.env[key as string] : undefined); + return fromEnv ?? fallback; +} + +function initGlobals(env: RouterEnv) { + if (!rateLimiter) { + const kvAdapter = env.KV ? { + get: async (k: string) => { + const val = await env.KV!.get(k); + return val ? JSON.parse(val) : null; + }, + set: async (k: string, v: number[], ttl?: number) => { + await env.KV!.put(k, JSON.stringify(v), { expirationTtl: ttl }); + } + } : undefined; + + rateLimiter = createRateLimiter({ + maxRequests: getNumericEnv(env, "RATE_LIMIT_MAX", 100), + windowMs: getNumericEnv(env, "RATE_LIMIT_WINDOW_MS", 60000), + kv: kvAdapter, + }); + } + + if (!proxyPool) { + proxyPool = new ProxyPool(); + // For Bun (process.env.PROXY_FILE) it's loaded in index.ts, but for serverless we can load from env + const proxies = getEnv(env, "PROXY_LIST", ""); + if (proxies) { + for (const p of proxies.split(",")) { + const pt = p.trim(); + if (pt) proxyPool.addProxy(pt); + } + } + sessionPool = new SessionProxyPool(proxyPool); + sessionPool.setFailureThreshold(3); + } +} + +function requireAuth(req: Request, env: RouterEnv): Response | null { + const API_KEY = getEnv(env, "API_KEY", "sk-dummy-key"); + const header = req.headers.get("authorization") ?? req.headers.get("x-api-key") ?? ""; + const key = header.replace(/^Bearer\s+/i, "").trim(); + if (key === API_KEY) return null; + return new Response( + JSON.stringify({ error: { message: "Unauthorized", type: "auth_error" } }), + { + status: 401, + headers: { "Content-Type": "application/json", ...getCorsHeaders() }, + }, + ); +} + +// --- Static Handlers --------------------------------------------------------- + +function handleHealth(): Response { + return new Response( + JSON.stringify({ + status: "ok", + uptime: Date.now() - SERVER_START_TIME, + version: RELAY_VERSION, + }), + { + status: 200, + headers: { + "Content-Type": "application/json", + ...getCorsHeaders(), + }, + }, + ); +} + +function handleIndex(): Response { + const html = ` + + + + + Edge Proxy Relay + + + +
+

Edge Proxy Relay

+

Server is running

+

/health · /docs

+
+ + `; + + return new Response(html, { + status: 200, + headers: { + "Content-Type": "text/html; charset=utf-8", + }, + }); +} + +function handleDocs(isWebSocketSupported: boolean): Response { + const html = ` + + + + + Edge Proxy Relay — Docs + + + +
+

Edge Proxy Relay

+

Forward HTTP${isWebSocketSupported ? ' and WebSocket' : ''} requests to any target server via the x-relay-target header.

+ +

Endpoints

+ +
+

GET /health

+

Health check. Returns 200 OK with server status, uptime, and version.

+
+ +
+

GET /docs

+

This page.

+
+ +
+

Any Path (Catch-all Relay)

+

Send a request with the x-relay-target header and this proxy forwards it.

+
+ +

Usage — HTTP Relay

+ 
curl -s \\
+	  -H "x-relay-target: https://httpbin.org" \\
+	  -H "x-relay-path: /get" \\
+	  "https://your-proxy.example/any/path"
+ + + + +
HeaderRequiredDescription
x-relay-targetYesBase URL of the upstream (http:// or https://)
x-relay-pathNoPath to append (default: /)
+ +${isWebSocketSupported ? ` +

Usage — WebSocket Relay

+ 
const ws = new WebSocket("wss://your-proxy.example/relay", {
+	  headers: { "x-relay-target": "wss://echo-websocket.example" },
+	});
+	ws.onopen = () => ws.send("Hello via relay!");
+	ws.onmessage = (e) => console.log("Got:", e.data);
+` : `

Note: WebSocket relay is not available on this deployment.

`} + +

Status Codes

+ + + + + + + + + +
CodeMeaning
204CORS preflight success (OPTIONS)
400Missing x-relay-target header
403Target blocked (SSRF protection / not allowed)
413Request body exceeds size limit
429Rate limit exceeded
502Upstream network / DNS error
504Upstream timeout
+
+ + `; + + return new Response(html, { + status: 200, + headers: { + "Content-Type": "text/html; charset=utf-8", + ...getCorsHeaders(), + }, + }); +} + +// --- Generic HTTP Relay ------------------------------------------------------ + +async function handleRelay( + req: Request, + env: RouterEnv, + clientIP: string, +): Promise { + const startTime = performance.now(); + const method = req.method; + const requestUrl = req.url; + + const RELAY_TIMEOUT_MS = getNumericEnv(env, "RELAY_TIMEOUT_MS", 30000); + + // -- Pre-flight CORS + if (method === "OPTIONS") { + return createCorsPreflightResponse(); + } + + // -- Middleware: Body size check + const bodyError = checkBodySize(req); + if (bodyError) { + logRelayEvent({ + method, + url: requestUrl, + status: bodyError.status, + durationMs: Math.round(performance.now() - startTime), + ip: clientIP, + }); + return bodyError; + } + + // -- Middleware: Rate limiting + const rateCheck = await rateLimiter!.checkAsync(clientIP); + if (!rateCheck.allowed) { + logRelayEvent({ + method, + url: requestUrl, + status: 429, + durationMs: Math.round(performance.now() - startTime), + error: "rate_limit_exceeded", + ip: clientIP, + }); + return new Response( + JSON.stringify({ + error: true, + code: "RATE_LIMITED", + message: "Too many requests", + retryAfterMs: rateCheck.retryAfterMs, + }), + { + status: 429, + headers: { + "Content-Type": "application/json", + ...getCorsHeaders(), + "Retry-After": String( + Math.ceil((rateCheck.retryAfterMs ?? 60_000) / 1000), + ), + }, + }, + ); + } + + // -- Extract relay parameters from headers + const target = req.headers.get("x-relay-target"); + const relayPath = req.headers.get("x-relay-path") ?? "/"; + + // -- SSRF: Normalize and validate target URL + const targetUrl = normalizeTargetUrl(target, relayPath); + if (!targetUrl) { + logRelayEvent({ + method, + url: requestUrl, + status: 400, + durationMs: Math.round(performance.now() - startTime), + error: "missing_target_header", + ip: clientIP, + }); + return createErrorResponse({ + code: "INVALID_TARGET", + status: 400, + message: "Missing or invalid x-relay-target header", + }); + } + + if (!isAllowedTarget(targetUrl)) { + logRelayEvent({ + method, + url: requestUrl, + status: 403, + durationMs: Math.round(performance.now() - startTime), + error: "target_not_allowed", + ip: clientIP, + }); + return createErrorResponse({ + code: "SSRF_BLOCKED", + status: 403, + message: "Target domain not allowed", + }); + } + + // -- SSRF: DNS rebinding protection (optional, via SSRF_DNS_CHECK=true) ----- + if (isSsrfDnsCheckEnabled()) { + const asyncAllowed = await isAllowedTargetAsync(targetUrl); + if (!asyncAllowed) { + logRelayEvent({ + method, + url: requestUrl, + status: 403, + durationMs: Math.round(performance.now() - startTime), + error: "ssrf_dns_rebinding", + ip: clientIP, + }); + return createErrorResponse({ + code: "SSRF_BLOCKED", + status: 403, + message: "Target resolves to private/internal IP", + }); + } + } + + // -- Build the upstream request + const filteredHeaders = filterRequestHeaders(req.headers); + const fetchOptions = buildRelayRequest( + req, + filteredHeaders, + RELAY_TIMEOUT_MS, + ) as RequestInit & { proxy?: string }; + + const targetUrlString = targetUrl.toString(); + + // -- Execute upstream fetch with shared retry + const result = await fetchWithRetry( + targetUrlString, + fetchOptions, + proxyPool!, + "relay", + ); + + if (result.errorClassification) { + logRelayEvent({ + method, + url: requestUrl, + status: result.errorClassification.status, + durationMs: Math.round(performance.now() - startTime), + error: result.errorClassification.message, + targetUrl: targetUrlString, + ip: clientIP, + }); + return createErrorResponse(result.errorClassification); + } + + const relayedResponse = createRelayResponse(result.response!); + + logRelayEvent({ + method, + url: requestUrl, + status: relayedResponse.status, + durationMs: Math.round(performance.now() - startTime), + targetUrl: targetUrlString, + ip: clientIP, + }); + + return relayedResponse; +} + +// --- Main Router ------------------------------------------------------------- + +export interface RouterOptions { + isWebSocketSupported?: boolean; + getTestApiHtml?: () => string | Promise; +} + +export async function handleRequest( + req: Request, + env: RouterEnv, + clientIP: string, + options: RouterOptions = {}, +): Promise { + initGlobals(env); + const url = new URL(req.url); + + // Static routes + if (url.pathname === "/health") return handleHealth(); + if (url.pathname === "/docs") return handleDocs(options.isWebSocketSupported ?? false); + if (url.pathname === "/test" && options.getTestApiHtml) { + const html = await options.getTestApiHtml(); + return new Response(html, { + status: 200, + headers: { "Content-Type": "text/html; charset=utf-8" }, + }); + } + if (url.pathname === "/" && req.method === "GET" && !req.headers.get("x-relay-target")) { + return handleIndex(); + } + + // AI proxy routes -- OpenAI-compatible API + if (url.pathname === "/v1/chat/completions") { + if (req.method === "OPTIONS") return createCorsPreflightResponse(); + if (req.method !== "POST") return new Response("Method Not Allowed", { status: 405 }); + const authErr = requireAuth(req, env); + if (authErr) return authErr; + try { + const body = await req.json(); + const sessionId = crypto.randomUUID(); + return handleChatCompletion(body, proxyPool!, sessionPool!, sessionId); + } catch { + return new Response( + JSON.stringify({ error: { message: "Invalid JSON body", type: "invalid_request_error" } }), + { status: 400, headers: { "Content-Type": "application/json", ...getCorsHeaders() } }, + ); + } + } + + // AI proxy routes -- Anthropic-compatible API + if (url.pathname === "/v1/messages") { + if (req.method === "OPTIONS") return createCorsPreflightResponse(); + if (req.method !== "POST") return new Response("Method Not Allowed", { status: 405 }); + const authErr = requireAuth(req, env); + if (authErr) return authErr; + try { + const body = await req.json(); + const sessionId = crypto.randomUUID(); + return handleAnthropicMessages(body, proxyPool!, sessionPool!, sessionId); + } catch { + return new Response( + JSON.stringify({ + type: "error", + error: { message: "Invalid JSON body", type: "invalid_request_error" }, + }), + { status: 400, headers: { "Content-Type": "application/json", ...getCorsHeaders() } }, + ); + } + } + + if (url.pathname === "/v1/models" && req.method === "GET") { + const authErr = requireAuth(req, env); + if (authErr) return authErr; + return new Response( + JSON.stringify({ + object: "list", + data: listModels().map((id) => ({ + id, + object: "model", + created: Math.floor(Date.now() / 1000), + owned_by: "edge-proxy", + })), + }), + { + status: 200, + headers: { + "Content-Type": "application/json", + ...getCorsHeaders(), + }, + }, + ); + } + + // WebSocket upgrade check (Bun specific - worker/vercel should handle their own rejection if needed) + if ( + req.method === "GET" && + req.headers.get("upgrade")?.toLowerCase() === "websocket" + ) { + if (!options.isWebSocketSupported) { + return new Response( + JSON.stringify({ + error: true, + code: "UNSUPPORTED", + message: "WebSocket relay is not supported on this deployment", + }), + { + status: 400, + headers: { + "Content-Type": "application/json", + ...getCorsHeaders(), + }, + }, + ); + } + // Return undefined to let Bun handle the upgrade in its fetch method + return undefined; + } + + // Generic HTTP relay + return handleRelay(req, env, clientIP); +} + +// Ensure proxyPool is available for index.ts to use proxyPool.tryLoad() +export function getSharedProxyPool() { + if (!proxyPool) { + proxyPool = new ProxyPool(); + sessionPool = new SessionProxyPool(proxyPool); + sessionPool.setFailureThreshold(3); + } + return proxyPool; +} diff --git a/src/middleware/body-limiter.test.ts b/src/middleware/body-limiter.test.ts index 21e95a7..a1bfc56 100644 --- a/src/middleware/body-limiter.test.ts +++ b/src/middleware/body-limiter.test.ts @@ -80,7 +80,7 @@ describe("body-limiter", () => { headers: { "Content-Length": "200" }, }); const result = checkBodySize(request); - const body = await result!.json(); + const body = await result!.json() as Record; expect(body.error).toBe("Payload Too Large"); expect(body.maxSizeBytes).toBe(100); }); diff --git a/src/middleware/rate-limiter.test.ts b/src/middleware/rate-limiter.test.ts index a0f36c6..8af31e8 100644 --- a/src/middleware/rate-limiter.test.ts +++ b/src/middleware/rate-limiter.test.ts @@ -10,76 +10,76 @@ import { createRateLimiter } from "./rate-limiter"; describe("rate-limiter", () => { describe("createRateLimiter", () => { - test("should allow requests up to the default limit", () => { + test("should allow requests up to the default limit", async () => { const limiter = createRateLimiter({ maxRequests: 5, windowMs: 60_000 }); for (let i = 0; i < 5; i++) { - const result = limiter.check("test-key"); + const result = await limiter.checkAsync("test-key"); expect(result.allowed).toBe(true); } }); - test("should block requests exceeding the limit", () => { + test("should block requests exceeding the limit", async () => { const limiter = createRateLimiter({ maxRequests: 3, windowMs: 60_000 }); for (let i = 0; i < 3; i++) { - expect(limiter.check("block-key").allowed).toBe(true); + expect((await limiter.checkAsync("block-key")).allowed).toBe(true); } - const blocked = limiter.check("block-key"); + const blocked = await limiter.checkAsync("block-key"); expect(blocked.allowed).toBe(false); expect(blocked.retryAfterMs).toBeDefined(); expect(typeof blocked.retryAfterMs).toBe("number"); }); - test("should return retryAfterMs when blocked", () => { + test("should return retryAfterMs when blocked", async () => { const limiter = createRateLimiter({ maxRequests: 1, windowMs: 60_000, }); - limiter.check("retry-key"); - const blocked = limiter.check("retry-key"); + await limiter.checkAsync("retry-key"); + const blocked = await limiter.checkAsync("retry-key"); expect(blocked.allowed).toBe(false); expect(blocked.retryAfterMs).toBeGreaterThan(0); expect(blocked.retryAfterMs).toBeLessThanOrEqual(60_000); }); - test("reset() should clear the counter", () => { + test("reset() should clear the counter", async () => { const limiter = createRateLimiter({ maxRequests: 2, windowMs: 60_000 }); - limiter.check("reset-key"); - limiter.check("reset-key"); + await limiter.checkAsync("reset-key"); + await limiter.checkAsync("reset-key"); // would be blocked, but... - limiter.reset("reset-key"); + await limiter.resetAsync("reset-key"); // ...should be allowed again - expect(limiter.check("reset-key").allowed).toBe(true); + expect((await limiter.checkAsync("reset-key")).allowed).toBe(true); }); - test("should isolate keys from each other", () => { + test("should isolate keys from each other", async () => { const limiter = createRateLimiter({ maxRequests: 2, windowMs: 60_000 }); - expect(limiter.check("key-a").allowed).toBe(true); - expect(limiter.check("key-a").allowed).toBe(true); - expect(limiter.check("key-b").allowed).toBe(true); // key-b unaffected - expect(limiter.check("key-a").allowed).toBe(false); // key-a blocked + expect((await limiter.checkAsync("key-a")).allowed).toBe(true); + expect((await limiter.checkAsync("key-a")).allowed).toBe(true); + expect((await limiter.checkAsync("key-b")).allowed).toBe(true); // key-b unaffected + expect((await limiter.checkAsync("key-a")).allowed).toBe(false); // key-a blocked }); test("should create with default options", () => { const limiter = createRateLimiter(); - expect(limiter.check).toBeDefined(); - expect(limiter.reset).toBeDefined(); + expect(limiter.checkAsync).toBeDefined(); + expect(limiter.resetAsync).toBeDefined(); }); - test("should handle rapid sequential calls", () => { + test("should handle rapid sequential calls", async () => { const limiter = createRateLimiter({ maxRequests: 100, windowMs: 60_000 }); for (let i = 0; i < 100; i++) { - expect(limiter.check("rapid-key").allowed).toBe(true); + expect((await limiter.checkAsync("rapid-key")).allowed).toBe(true); } - expect(limiter.check("rapid-key").allowed).toBe(false); + expect((await limiter.checkAsync("rapid-key")).allowed).toBe(false); }); - test("should allow requests after reset", () => { + test("should allow requests after reset", async () => { const limiter = createRateLimiter({ maxRequests: 1, windowMs: 60_000 }); - limiter.check("after-reset-key"); - const blocked = limiter.check("after-reset-key"); + await limiter.checkAsync("after-reset-key"); + const blocked = await limiter.checkAsync("after-reset-key"); expect(blocked.allowed).toBe(false); - limiter.reset("after-reset-key"); - expect(limiter.check("after-reset-key").allowed).toBe(true); + await limiter.resetAsync("after-reset-key"); + expect((await limiter.checkAsync("after-reset-key")).allowed).toBe(true); }); }); }); diff --git a/src/middleware/rate-limiter.ts b/src/middleware/rate-limiter.ts index 78871c9..f377fd2 100644 --- a/src/middleware/rate-limiter.ts +++ b/src/middleware/rate-limiter.ts @@ -11,11 +11,16 @@ export interface RateLimiterOptions { maxRequests?: number; windowMs?: number; + /** Optional async KV store for distributed rate limiting */ + kv?: { + get(key: string): Promise; + set(key: string, value: number[], expirationTtl?: number): Promise; + }; } export interface RateLimiter { - check(key: string): { allowed: boolean; retryAfterMs?: number }; - reset(key: string): void; + checkAsync(key: string): Promise<{ allowed: boolean; retryAfterMs?: number }>; + resetAsync(key: string): Promise; } const DEFAULT_MAX_REQUESTS = 100; @@ -25,17 +30,17 @@ const CLEANUP_INTERVAL_DIVISOR = 10; export function createRateLimiter(options?: RateLimiterOptions): RateLimiter { const maxRequests = options?.maxRequests ?? DEFAULT_MAX_REQUESTS; const windowMs = options?.windowMs ?? DEFAULT_WINDOW_MS; + const kv = options?.kv; - // Map of key -> sorted array of timestamps (ascending) + // In-memory store fallback const store = new Map(); + // Maximum number of unique keys to track (prevents unbounded memory growth) + const MAX_STORE_KEYS = 10_000; + // ── helpers ────────────────────────────────────────────────────── - /** Remove timestamps outside the sliding window. Returns the pruned slice. */ - function prune(key: string, now: number): number[] { - const timestamps = store.get(key); - if (!timestamps) return []; - + function pruneTimestamps(timestamps: number[], now: number): number[] { const cutoff = now - windowMs; const result: number[] = []; for (let i = 0; i < timestamps.length; i++) { @@ -43,17 +48,12 @@ export function createRateLimiter(options?: RateLimiterOptions): RateLimiter { result.push(timestamps[i]); } } - - if (result.length === 0) { - store.delete(key); - } else { - store.set(key, result); - } return result; } - /** Periodically sweep the entire store to free memory. */ + /** Periodically sweep the entire memory store to free memory. */ function periodicCleanup(): void { + if (kv) return; // Cleanup handled by TTL in KV const now = Date.now(); const cutoff = now - windowMs; @@ -70,30 +70,67 @@ export function createRateLimiter(options?: RateLimiterOptions): RateLimiter { store.set(key, pruned); } } + + // If store still exceeds max keys after pruning, evict oldest entries + if (store.size > MAX_STORE_KEYS) { + // Sort by oldest timestamp, evict excess + const entries = Array.from(store.entries()) + .sort((a, b) => { + const aOldest = a[1][0] ?? 0; + const bOldest = b[1][0] ?? 0; + return aOldest - bOldest; + }); + const toEvict = entries.length - MAX_STORE_KEYS; + for (let i = 0; i < toEvict; i++) { + store.delete(entries[i]![0]); + } + console.warn(`[rate-limiter] Evicted ${toEvict} keys from store (max ${MAX_STORE_KEYS})`); + } } - // Schedule periodic cleanup (every windowMs / 10) - const cleanupHandle = setInterval( - periodicCleanup, - windowMs / CLEANUP_INTERVAL_DIVISOR, - ); - // Allow the process to exit even if the interval is still active - if ( - cleanupHandle && - typeof cleanupHandle === "object" && - "unref" in cleanupHandle - ) { - (cleanupHandle as NodeJS.Timeout).unref(); + if (!kv) { + // Schedule periodic cleanup (every windowMs / 10) + const cleanupHandle = setInterval( + periodicCleanup, + windowMs / CLEANUP_INTERVAL_DIVISOR, + ); + // Allow the process to exit even if the interval is still active + if ( + cleanupHandle && + typeof cleanupHandle === "object" && + "unref" in cleanupHandle + ) { + (cleanupHandle as NodeJS.Timeout).unref(); + } } // ── public API ─────────────────────────────────────────────────── return { - check(key: string): { allowed: boolean; retryAfterMs?: number } { + async checkAsync(key: string): Promise<{ allowed: boolean; retryAfterMs?: number }> { const now = Date.now(); - const timestamps = prune(key, now); + let timestamps: number[] = []; + + if (kv) { + try { + const val = await kv.get(key); + if (val) timestamps = val; + } catch { + // Fallback to empty + } + } else { + timestamps = store.get(key) ?? []; + } + + timestamps = pruneTimestamps(timestamps, now); timestamps.push(now); - store.set(key, timestamps); + + if (kv) { + // We set TTL slightly higher than windowMs so it cleans up automatically + await kv.set(key, timestamps, Math.ceil(windowMs / 1000) + 10).catch(() => {}); + } else { + store.set(key, timestamps); + } if (timestamps.length <= maxRequests) { return { allowed: true }; @@ -110,8 +147,12 @@ export function createRateLimiter(options?: RateLimiterOptions): RateLimiter { return { allowed: false, retryAfterMs }; }, - reset(key: string): void { - store.delete(key); + async resetAsync(key: string): Promise { + if (kv) { + await kv.set(key, [], 1).catch(() => {}); + } else { + store.delete(key); + } }, }; } diff --git a/src/worker.ts b/src/worker.ts index 2c924a4..405e8a1 100644 --- a/src/worker.ts +++ b/src/worker.ts @@ -84,8 +84,8 @@ function getClientIP(req: Request): string { // --- Auth Helper --------------------------------------------------------------- -function requireAuth(req: Request): Response | null { - const apiKey = process.env.API_KEY ?? "sk-dummy-key"; +function requireAuth(req: Request, env: Env): Response | null { + const apiKey = env.API_KEY ?? "sk-dummy-key"; const header = req.headers.get("authorization") ?? req.headers.get("x-api-key") ?? ""; const key = header.replace(/^Bearer\s+/i, "").trim(); if (key === apiKey) return null; @@ -268,7 +268,7 @@ async function handleRelay(req: Request, env: Env): Promise { } // -- Middleware: Rate limiting --------------------------------------------- - const rateCheck = limiter.check(clientIP); + const rateCheck = await limiter.checkAsync(clientIP); if (!rateCheck.allowed) { logRelayEvent({ method, @@ -387,12 +387,20 @@ export default { if (url.pathname === "/health") return handleHealth(); if (url.pathname === "/docs" || url.pathname === "/test") { - const file = Bun.file("public/test-api.html"); - const exists = await file.exists(); - return new Response(exists ? file : "Not found", { - status: exists ? 200 : 404, - headers: { "Content-Type": "text/html; charset=utf-8" }, - }); + try { + const file = Bun.file("public/test-api.html"); + const exists = await file.exists(); + return new Response(exists ? file : "Not found", { + status: exists ? 200 : 404, + headers: { "Content-Type": "text/html; charset=utf-8" }, + }); + } catch { + // Bun.file not available in non-Bun runtimes (e.g. Workers) + return new Response("Not found", { + status: 404, + headers: { "Content-Type": "text/html; charset=utf-8" }, + }); + } } if ( url.pathname === "/" && @@ -425,7 +433,7 @@ export default { if (url.pathname === "/v1/chat/completions") { if (req.method === "OPTIONS") return createCorsPreflightResponse(); if (req.method !== "POST") return new Response("Method Not Allowed", { status: 405 }); - const authErr = requireAuth(req); + const authErr = requireAuth(req, env); if (authErr) return authErr; try { const body = await req.json(); @@ -441,7 +449,7 @@ export default { if (url.pathname === "/v1/messages") { if (req.method === "OPTIONS") return createCorsPreflightResponse(); if (req.method !== "POST") return new Response("Method Not Allowed", { status: 405 }); - const authErr = requireAuth(req); + const authErr = requireAuth(req, env); if (authErr) return authErr; try { const body = await req.json(); @@ -455,7 +463,7 @@ export default { } if (url.pathname === "/v1/models" && req.method === "GET") { - const authErr = requireAuth(req); + const authErr = requireAuth(req, env); if (authErr) return authErr; const models = listModels().map((id) => ({ id,