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<string, unknown>) - All 153 tests pass, typecheck clean
This commit is contained in:
+1
-1
@@ -252,7 +252,7 @@ async function handleRelay(req: Request): Promise<Response> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// ── Middleware: Rate limiting ────────────────────────────────────
|
// ── Middleware: Rate limiting ────────────────────────────────────
|
||||||
const rateCheck = rateLimiter.check(clientIP);
|
const rateCheck = await rateLimiter.checkAsync(clientIP);
|
||||||
if (!rateCheck.allowed) {
|
if (!rateCheck.allowed) {
|
||||||
logRelayEvent({
|
logRelayEvent({
|
||||||
method,
|
method,
|
||||||
|
|||||||
+1
-1
@@ -25,7 +25,7 @@ describe("handleHealth", () => {
|
|||||||
|
|
||||||
test("should include status, uptime, and version fields", async () => {
|
test("should include status, uptime, and version fields", async () => {
|
||||||
const response = handleHealth();
|
const response = handleHealth();
|
||||||
const body = await response.json();
|
const body = await response.json() as Record<string, unknown>;
|
||||||
expect(body.status).toBe("ok");
|
expect(body.status).toBe("ok");
|
||||||
expect(typeof body.uptime).toBe("number");
|
expect(typeof body.uptime).toBe("number");
|
||||||
expect(body.version).toBe("1.0.0");
|
expect(body.version).toBe("1.0.0");
|
||||||
|
|||||||
+23
-3
@@ -24,6 +24,7 @@ import {
|
|||||||
createErrorResponse,
|
createErrorResponse,
|
||||||
createCorsPreflightResponse,
|
createCorsPreflightResponse,
|
||||||
getCorsHeaders,
|
getCorsHeaders,
|
||||||
|
setSsrfDnsCheck,
|
||||||
} from "./lib/relay-utils";
|
} from "./lib/relay-utils";
|
||||||
|
|
||||||
import { checkBodySize } from "./middleware/body-limiter";
|
import { checkBodySize } from "./middleware/body-limiter";
|
||||||
@@ -86,6 +87,13 @@ proxyPool.tryLoad(
|
|||||||
const sessionPool = new SessionProxyPool(proxyPool);
|
const sessionPool = new SessionProxyPool(proxyPool);
|
||||||
sessionPool.setFailureThreshold(3);
|
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 -----------------------------------------------
|
// --- WebSocket relay data type -----------------------------------------------
|
||||||
|
|
||||||
interface WSRelayData {
|
interface WSRelayData {
|
||||||
@@ -302,7 +310,7 @@ async function handleRelay(
|
|||||||
}
|
}
|
||||||
|
|
||||||
// -- Middleware: Rate limiting ---------------------------------------------
|
// -- Middleware: Rate limiting ---------------------------------------------
|
||||||
const rateCheck = rateLimiter.check(clientIP);
|
const rateCheck = await rateLimiter.checkAsync(clientIP);
|
||||||
if (!rateCheck.allowed) {
|
if (!rateCheck.allowed) {
|
||||||
logRelayEvent({
|
logRelayEvent({
|
||||||
method,
|
method,
|
||||||
@@ -438,8 +446,20 @@ function handleWebSocketUpgrade(
|
|||||||
const relayPath = req.headers.get("x-relay-path") ?? "/";
|
const relayPath = req.headers.get("x-relay-path") ?? "/";
|
||||||
|
|
||||||
const normalized = normalizeTargetUrl(target, relayPath);
|
const normalized = normalizeTargetUrl(target, relayPath);
|
||||||
if (!normalized) return undefined;
|
if (!normalized) {
|
||||||
if (!isAllowedTarget(new URL(normalized.toString()))) return undefined;
|
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();
|
const targetUrl = normalized.toString();
|
||||||
|
|
||||||
|
|||||||
@@ -28,14 +28,12 @@ export function closeAllActiveReaders(): void {
|
|||||||
/**
|
/**
|
||||||
* Returns `true` when development features (HMR, verbose console) should
|
* Returns `true` when development features (HMR, verbose console) should
|
||||||
* be enabled. Controlled by the `NODE_ENV` / `BUN_ENV` env var — defaults
|
* be enabled. Controlled by the `NODE_ENV` / `BUN_ENV` env var — defaults
|
||||||
* to `true` for convenience during local development.
|
* to `false` (production-safe). Set `NODE_ENV=development` to enable.
|
||||||
*
|
|
||||||
* Set `NODE_ENV=production` or `BUN_ENV=production` to disable.
|
|
||||||
*/
|
*/
|
||||||
export function isDevMode(): boolean {
|
export function isDevMode(): boolean {
|
||||||
const env = (process.env.NODE_ENV ?? process.env.BUN_ENV ?? "").toLowerCase();
|
const env = (process.env.NODE_ENV ?? process.env.BUN_ENV ?? "").toLowerCase();
|
||||||
if (env === "production") return false;
|
if (env === "development" || env === "dev") return true;
|
||||||
return true;
|
return false;
|
||||||
}
|
}
|
||||||
|
|
||||||
// ─── SSE line buffer (fixes chunk-boundary corruption) ───────────────────
|
// ─── SSE line buffer (fixes chunk-boundary corruption) ───────────────────
|
||||||
@@ -51,6 +49,7 @@ export function isDevMode(): boolean {
|
|||||||
*/
|
*/
|
||||||
export class SSELineBuffer {
|
export class SSELineBuffer {
|
||||||
private buffer = "";
|
private buffer = "";
|
||||||
|
private readonly MAX_BUFFER_SIZE = 1024 * 1024; // 1MB limit to prevent OOM
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Feed a chunk of decoded text and return complete lines.
|
* Feed a chunk of decoded text and return complete lines.
|
||||||
@@ -58,6 +57,11 @@ export class SSELineBuffer {
|
|||||||
*/
|
*/
|
||||||
add(chunk: string): string[] {
|
add(chunk: string): string[] {
|
||||||
this.buffer += chunk;
|
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 [];
|
if (!this.buffer.includes("\n")) return [];
|
||||||
|
|
||||||
const parts = this.buffer.split("\n");
|
const parts = this.buffer.split("\n");
|
||||||
|
|||||||
@@ -29,8 +29,10 @@ const EXPIRY_BUFFER_MS = 300_000; // 5 minutes
|
|||||||
*
|
*
|
||||||
* Format: `sha256(hostname|platform|arch|cpu|username)`
|
* Format: `sha256(hostname|platform|arch|cpu|username)`
|
||||||
* This matches the 9Router reference implementation.
|
* 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<string> {
|
||||||
const hostname = os.hostname();
|
const hostname = os.hostname();
|
||||||
const platform = process.platform;
|
const platform = process.platform;
|
||||||
const arch = process.arch;
|
const arch = process.arch;
|
||||||
@@ -39,9 +41,11 @@ function generateDeviceFingerprint(): string {
|
|||||||
const username = process.env.USER ?? process.env.USERNAME ?? "unknown";
|
const username = process.env.USER ?? process.env.USERNAME ?? "unknown";
|
||||||
|
|
||||||
const raw = `${hostname}|${platform}|${arch}|${cpuModel}|${username}`;
|
const raw = `${hostname}|${platform}|${arch}|${cpuModel}|${username}`;
|
||||||
const hasher = new Bun.CryptoHasher("sha256");
|
const encoder = new TextEncoder();
|
||||||
hasher.update(raw);
|
const data = encoder.encode(raw);
|
||||||
return hasher.digest("hex") as string;
|
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 -----------------------------------------------------------
|
// --- JWT bootstrap -----------------------------------------------------------
|
||||||
@@ -55,7 +59,7 @@ function generateDeviceFingerprint(): string {
|
|||||||
* @throws If the bootstrap request fails or returns an unexpected response.
|
* @throws If the bootstrap request fails or returns an unexpected response.
|
||||||
*/
|
*/
|
||||||
async function bootstrapJwt(): Promise<string> {
|
async function bootstrapJwt(): Promise<string> {
|
||||||
const fingerprint = generateDeviceFingerprint();
|
const fingerprint = await generateDeviceFingerprint();
|
||||||
|
|
||||||
const resp = await fetch(MIMO_BOOTSTRAP_URL, {
|
const resp = await fetch(MIMO_BOOTSTRAP_URL, {
|
||||||
method: "POST",
|
method: "POST",
|
||||||
|
|||||||
+117
-27
@@ -59,6 +59,45 @@ export class ProxyPool {
|
|||||||
}
|
}
|
||||||
|
|
||||||
const text = readFileSync(filePath, "utf-8");
|
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<void> {
|
||||||
|
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
|
const lines = text
|
||||||
.split("\n")
|
.split("\n")
|
||||||
.map((l: string) => l.trim())
|
.map((l: string) => l.trim())
|
||||||
@@ -87,7 +126,7 @@ export class ProxyPool {
|
|||||||
this.currentIndex = 0;
|
this.currentIndex = 0;
|
||||||
this.failures.clear();
|
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;
|
return this.proxies.length > 0;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Async convenience -- load from a path or skip.
|
||||||
|
*/
|
||||||
|
async tryLoadAsync(filePath?: string): Promise<boolean> {
|
||||||
|
if (filePath) await this.loadAsync(filePath);
|
||||||
|
return this.proxies.length > 0;
|
||||||
|
}
|
||||||
|
|
||||||
// -- Access ------------------------------------------------------------------
|
// -- Access ------------------------------------------------------------------
|
||||||
|
|
||||||
/** Total number of proxies in the pool. */
|
/** Total number of proxies in the pool. */
|
||||||
@@ -137,6 +184,54 @@ export class ProxyPool {
|
|||||||
return `http://${auth}${entry.host}:${entry.port}`;
|
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 ----------------------------------------------------------------
|
// -- Rotation ----------------------------------------------------------------
|
||||||
|
|
||||||
/**
|
/**
|
||||||
@@ -342,7 +437,7 @@ export class SessionProxyPool {
|
|||||||
const existing = this.sessions.get(sessionId);
|
const existing = this.sessions.get(sessionId);
|
||||||
if (existing !== undefined) {
|
if (existing !== undefined) {
|
||||||
logPool(`acquire existing session=${sessionId.slice(0, 8)} proxyIndex=${existing.proxyIndex}`);
|
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);
|
const index = this.pickLeastUsedIndex(model);
|
||||||
@@ -357,18 +452,18 @@ export class SessionProxyPool {
|
|||||||
}
|
}
|
||||||
usedBy.add(sessionId);
|
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}`);
|
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. */
|
/** Return the current proxy URL for a session (no rotation), or null. */
|
||||||
getProxyUrl(sessionId: string): string | null {
|
getProxyUrl(sessionId: string): string | null {
|
||||||
const info = this.sessions.get(sessionId);
|
const info = this.sessions.get(sessionId);
|
||||||
if (!info) return null;
|
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}`);
|
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. */
|
/** Remove a session from all tracking. */
|
||||||
@@ -383,7 +478,7 @@ export class SessionProxyPool {
|
|||||||
}
|
}
|
||||||
|
|
||||||
this.sessions.delete(sessionId);
|
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}`);
|
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;
|
if (info.failures < this.failureThreshold) return false;
|
||||||
|
|
||||||
const oldIndex = info.proxyIndex;
|
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)
|
// Mark in underlying pool using public API
|
||||||
const savedIdx = (this.pool as any).currentIndex as number;
|
const savedIdx = this.pool.getCurrentIndex();
|
||||||
(this.pool as any).currentIndex = oldIndex;
|
this.pool.setCurrentIndex(oldIndex);
|
||||||
this.pool.markFailed(this.failureThreshold);
|
this.pool.markFailed(this.failureThreshold);
|
||||||
(this.pool as any).currentIndex = savedIdx;
|
this.pool.setCurrentIndex(savedIdx);
|
||||||
|
|
||||||
// Remove session from old proxy usage tracking
|
// Remove session from old proxy usage tracking
|
||||||
const usedBy = this.proxyUsage.get(oldIndex);
|
const usedBy = this.proxyUsage.get(oldIndex);
|
||||||
@@ -438,7 +533,7 @@ export class SessionProxyPool {
|
|||||||
}
|
}
|
||||||
newUsedBy.add(sessionId);
|
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}`);
|
logPool(`markFailed rotated session=${sessionId.slice(0, 8)} ${oldEntry?.host} -> ${newEntry?.host}`);
|
||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
@@ -454,7 +549,7 @@ export class SessionProxyPool {
|
|||||||
if (!info) return false;
|
if (!info) return false;
|
||||||
|
|
||||||
const oldIndex = info.proxyIndex;
|
const oldIndex = info.proxyIndex;
|
||||||
const oldEntry = this.poolEntryAtIndex(oldIndex);
|
const oldEntry = this.pool.getEntryAtIndex(oldIndex);
|
||||||
|
|
||||||
// Remove from old proxy usage first
|
// Remove from old proxy usage first
|
||||||
const usedBy = this.proxyUsage.get(oldIndex);
|
const usedBy = this.proxyUsage.get(oldIndex);
|
||||||
@@ -470,7 +565,7 @@ export class SessionProxyPool {
|
|||||||
let bestCount = Infinity;
|
let bestCount = Infinity;
|
||||||
for (let step = 1; step <= this.pool.size; step++) {
|
for (let step = 1; step <= this.pool.size; step++) {
|
||||||
const i = (oldIndex + step) % this.pool.size;
|
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;
|
const count = this.proxyUsage.get(i)?.size ?? 0;
|
||||||
if (count < bestCount) {
|
if (count < bestCount) {
|
||||||
bestCount = count;
|
bestCount = count;
|
||||||
@@ -493,7 +588,7 @@ export class SessionProxyPool {
|
|||||||
}
|
}
|
||||||
newUsedBy.add(sessionId);
|
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}`);
|
logPool(`rotateNow session=${sessionId.slice(0, 8)} ${oldEntry?.host} -> ${newEntry?.host}`);
|
||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
@@ -514,27 +609,22 @@ export class SessionProxyPool {
|
|||||||
markRateLimited(sessionId: string, model: string): void {
|
markRateLimited(sessionId: string, model: string): void {
|
||||||
const info = this.sessions.get(sessionId);
|
const info = this.sessions.get(sessionId);
|
||||||
if (!info) return;
|
if (!info) return;
|
||||||
const savedIdx = (this.pool as any).currentIndex as number;
|
const savedIdx = this.pool.getCurrentIndex();
|
||||||
(this.pool as any).currentIndex = info.proxyIndex;
|
this.pool.setCurrentIndex(info.proxyIndex);
|
||||||
this.pool.markRateLimited(model);
|
this.pool.markRateLimited(model);
|
||||||
(this.pool as any).currentIndex = savedIdx;
|
this.pool.setCurrentIndex(savedIdx);
|
||||||
}
|
}
|
||||||
|
|
||||||
// -- Internals ----------------------------------------------------------------
|
// -- Internals ----------------------------------------------------------------
|
||||||
|
|
||||||
/** Get the ProxyEntry at a given index. Forward reference to local type. */
|
/** Get the ProxyEntry at a given index. Forward reference to local type. */
|
||||||
private poolEntryAtIndex(index: number): ProxyEntry | null {
|
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. */
|
/** Build proxy URL string by index. */
|
||||||
private formatProxyUrlAtIndex(index: number): string | null {
|
private formatProxyUrlAtIndex(index: number): string | null {
|
||||||
const entry = this.poolEntryAtIndex(index);
|
return this.pool.getProxyUrlAtIndex(index);
|
||||||
if (!entry) return null;
|
|
||||||
const auth = entry.username
|
|
||||||
? `${encodeURIComponent(entry.username)}:${encodeURIComponent(entry.password)}@`
|
|
||||||
: "";
|
|
||||||
return `http://${auth}${entry.host}:${entry.port}`;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/** Return the index of the proxy with the fewest active sessions, or -1. */
|
/** Return the index of the proxy with the fewest active sessions, or -1. */
|
||||||
@@ -545,7 +635,7 @@ export class SessionProxyPool {
|
|||||||
let bestCount = Infinity;
|
let bestCount = Infinity;
|
||||||
|
|
||||||
for (let i = 0; i < this.pool.size; i++) {
|
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;
|
const count = this.proxyUsage.get(i)?.size ?? 0;
|
||||||
logPool(`pickLeastUsed proxy[${i}] count=${count}`);
|
logPool(`pickLeastUsed proxy[${i}] count=${count}`);
|
||||||
if (count < bestCount) {
|
if (count < bestCount) {
|
||||||
|
|||||||
@@ -692,7 +692,7 @@ describe("createErrorResponse", () => {
|
|||||||
});
|
});
|
||||||
expect(result.status).toBe(504);
|
expect(result.status).toBe(504);
|
||||||
expect(result.headers.get("Content-Type")).toBe("application/json");
|
expect(result.headers.get("Content-Type")).toBe("application/json");
|
||||||
const body = await result.json();
|
const body = await result.json() as Record<string, unknown>;
|
||||||
expect(body.error).toBe(true);
|
expect(body.error).toBe(true);
|
||||||
expect(body.code).toBe("TIMEOUT");
|
expect(body.code).toBe("TIMEOUT");
|
||||||
expect(body.message).toBe("Upstream timed out");
|
expect(body.message).toBe("Upstream timed out");
|
||||||
@@ -714,7 +714,7 @@ describe("createErrorResponse", () => {
|
|||||||
message: "DNS resolution failed",
|
message: "DNS resolution failed",
|
||||||
});
|
});
|
||||||
expect(result.status).toBe(502);
|
expect(result.status).toBe(502);
|
||||||
const body = await result.json();
|
const body = await result.json() as Record<string, unknown>;
|
||||||
expect(body.code).toBe("DNS_FAILURE");
|
expect(body.code).toBe("DNS_FAILURE");
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
@@ -849,7 +849,7 @@ describe("integration: full relay flow", () => {
|
|||||||
const classified = classifyFetchError(error);
|
const classified = classifyFetchError(error);
|
||||||
const response = createErrorResponse(classified);
|
const response = createErrorResponse(classified);
|
||||||
expect(response.status).toBe(504);
|
expect(response.status).toBe(504);
|
||||||
const body = await response.json();
|
const body = await response.json() as Record<string, unknown>;
|
||||||
expect(body.code).toBe("TIMEOUT");
|
expect(body.code).toBe("TIMEOUT");
|
||||||
expect(body.error).toBe(true);
|
expect(body.error).toBe(true);
|
||||||
});
|
});
|
||||||
|
|||||||
+27
-3
@@ -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 -------------------------------------------------------
|
// --- CORS configuration -------------------------------------------------------
|
||||||
|
|
||||||
/** Return CORS headers. Origin defaults to "*" but can be overridden via env. */
|
/** Return CORS headers. Origin defaults to "*" but can be overridden via env. */
|
||||||
@@ -39,20 +56,27 @@ export function getAllowedOrigin(): string {
|
|||||||
}
|
}
|
||||||
|
|
||||||
let cachedCorsHeaders: Record<string, string> | null = null;
|
let cachedCorsHeaders: Record<string, string> | null = null;
|
||||||
|
let cachedCorsOrigin: string | null = null;
|
||||||
|
|
||||||
/** Rebuild the CORS headers map (call after changing origin at runtime). */
|
/** Rebuild the CORS headers map (call after changing origin at runtime). */
|
||||||
export function rebuildCorsHeaders(): Record<string, string> {
|
export function rebuildCorsHeaders(): Record<string, string> {
|
||||||
|
const origin = getAllowedOrigin();
|
||||||
cachedCorsHeaders = {
|
cachedCorsHeaders = {
|
||||||
"Access-Control-Allow-Origin": getAllowedOrigin(),
|
"Access-Control-Allow-Origin": origin,
|
||||||
"Access-Control-Allow-Methods": "GET, POST, PUT, DELETE, PATCH, OPTIONS",
|
"Access-Control-Allow-Methods": "GET, POST, PUT, DELETE, PATCH, OPTIONS",
|
||||||
"Access-Control-Allow-Headers": "*",
|
"Access-Control-Allow-Headers": "*",
|
||||||
};
|
};
|
||||||
|
cachedCorsOrigin = origin;
|
||||||
return cachedCorsHeaders;
|
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<string, string> {
|
export function getCorsHeaders(): Record<string, string> {
|
||||||
if (!cachedCorsHeaders) rebuildCorsHeaders();
|
const currentOrigin = getAllowedOrigin();
|
||||||
|
// Invalidate cache if CORS_ORIGIN env changed
|
||||||
|
if (!cachedCorsHeaders || cachedCorsOrigin !== currentOrigin) {
|
||||||
|
rebuildCorsHeaders();
|
||||||
|
}
|
||||||
return cachedCorsHeaders!;
|
return cachedCorsHeaders!;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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<any>;
|
||||||
|
put(key: string, value: any, options?: { expirationTtl?: number }): Promise<void>;
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- Global singletons (survives warm starts) --------------------------------
|
||||||
|
|
||||||
|
let rateLimiter: ReturnType<typeof createRateLimiter> | 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 = `<!DOCTYPE html>
|
||||||
|
<html lang="en">
|
||||||
|
<head>
|
||||||
|
<meta charset="UTF-8">
|
||||||
|
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||||
|
<title>Edge Proxy Relay</title>
|
||||||
|
<style>
|
||||||
|
*, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }
|
||||||
|
body { font-family: system-ui, -apple-system, sans-serif; line-height: 1.6; color: #e1e4e8; background: #0d1117; display: flex; align-items: center; justify-content: center; min-height: 100vh; }
|
||||||
|
main { text-align: center; }
|
||||||
|
h1 { font-size: 2rem; color: #58a6ff; margin-bottom: 0.5rem; }
|
||||||
|
p { color: #8b949e; }
|
||||||
|
a { color: #58a6ff; }
|
||||||
|
.status { color: #3fb950; }
|
||||||
|
</style>
|
||||||
|
</head>
|
||||||
|
<body>
|
||||||
|
<main>
|
||||||
|
<h1>Edge Proxy Relay</h1>
|
||||||
|
<p class="status">Server is running</p>
|
||||||
|
<p><a href="/health">/health</a> · <a href="/docs">/docs</a></p>
|
||||||
|
</main>
|
||||||
|
</body>
|
||||||
|
</html>`;
|
||||||
|
|
||||||
|
return new Response(html, {
|
||||||
|
status: 200,
|
||||||
|
headers: {
|
||||||
|
"Content-Type": "text/html; charset=utf-8",
|
||||||
|
},
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
function handleDocs(isWebSocketSupported: boolean): Response {
|
||||||
|
const html = `<!DOCTYPE html>
|
||||||
|
<html lang="en">
|
||||||
|
<head>
|
||||||
|
<meta charset="UTF-8">
|
||||||
|
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||||
|
<title>Edge Proxy Relay — Docs</title>
|
||||||
|
<style>
|
||||||
|
*, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }
|
||||||
|
body { font-family: system-ui, -apple-system, sans-serif; line-height: 1.6; color: #e1e4e8; background: #0d1117; padding: 2rem; }
|
||||||
|
main { max-width: 800px; margin: 0 auto; }
|
||||||
|
h1 { font-size: 2rem; margin-bottom: 0.5rem; color: #58a6ff; }
|
||||||
|
h2 { font-size: 1.25rem; margin: 2rem 0 0.75rem; color: #c9d1d9; border-bottom: 1px solid #30363d; padding-bottom: 0.25rem; }
|
||||||
|
p, li { color: #8b949e; }
|
||||||
|
code { background: #161b22; padding: 0.2em 0.4em; border-radius: 4px; font-size: 0.9em; color: #f0f6fc; }
|
||||||
|
pre { background: #161b22; padding: 1rem; border-radius: 6px; overflow-x: auto; margin: 0.75rem 0; border: 1px solid #30363d; }
|
||||||
|
pre code { background: none; padding: 0; }
|
||||||
|
table { width: 100%; border-collapse: collapse; margin: 0.75rem 0; }
|
||||||
|
th, td { text-align: left; padding: 0.5rem 0.75rem; border: 1px solid #30363d; }
|
||||||
|
th { background: #161b22; color: #c9d1d9; }
|
||||||
|
ul { padding-left: 1.5rem; margin: 0.5rem 0; }
|
||||||
|
.endpoint { background: #161b22; border: 1px solid #30363d; border-radius: 6px; padding: 1rem; margin: 1rem 0; }
|
||||||
|
.endpoint h3 { color: #58a6ff; font-family: monospace; margin-bottom: 0.5rem; }
|
||||||
|
.status { color: #3fb950; }
|
||||||
|
a { color: #58a6ff; }
|
||||||
|
</style>
|
||||||
|
</head>
|
||||||
|
<body>
|
||||||
|
<main>
|
||||||
|
<h1>Edge Proxy Relay</h1>
|
||||||
|
<p>Forward HTTP${isWebSocketSupported ? ' and WebSocket' : ''} requests to any target server via the <code>x-relay-target</code> header.</p>
|
||||||
|
|
||||||
|
<h2>Endpoints</h2>
|
||||||
|
|
||||||
|
<div class="endpoint">
|
||||||
|
<h3>GET /health</h3>
|
||||||
|
<p>Health check. Returns <span class="status">200 OK</span> with server status, uptime, and version.</p>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div class="endpoint">
|
||||||
|
<h3>GET /docs</h3>
|
||||||
|
<p>This page.</p>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div class="endpoint">
|
||||||
|
<h3>Any Path (Catch-all Relay)</h3>
|
||||||
|
<p>Send a request with the <code>x-relay-target</code> header and this proxy forwards it.</p>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<h2>Usage — HTTP Relay</h2>
|
||||||
|
<pre><code>curl -s \\
|
||||||
|
-H "x-relay-target: https://httpbin.org" \\
|
||||||
|
-H "x-relay-path: /get" \\
|
||||||
|
"https://your-proxy.example/any/path"</code></pre>
|
||||||
|
<table>
|
||||||
|
<tr><th>Header</th><th>Required</th><th>Description</th></tr>
|
||||||
|
<tr><td><code>x-relay-target</code></td><td>Yes</td><td>Base URL of the upstream (http:// or https://)</td></tr>
|
||||||
|
<tr><td><code>x-relay-path</code></td><td>No</td><td>Path to append (default: <code>/</code>)</td></tr>
|
||||||
|
</table>
|
||||||
|
|
||||||
|
${isWebSocketSupported ? `
|
||||||
|
<h2>Usage — WebSocket Relay</h2>
|
||||||
|
<pre><code>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);</code></pre>
|
||||||
|
` : `<p><strong>Note:</strong> WebSocket relay is not available on this deployment.</p>`}
|
||||||
|
|
||||||
|
<h2>Status Codes</h2>
|
||||||
|
<table>
|
||||||
|
<tr><th>Code</th><th>Meaning</th></tr>
|
||||||
|
<tr><td>204</td><td>CORS preflight success (OPTIONS)</td></tr>
|
||||||
|
<tr><td>400</td><td>Missing <code>x-relay-target</code> header</td></tr>
|
||||||
|
<tr><td>403</td><td>Target blocked (SSRF protection / not allowed)</td></tr>
|
||||||
|
<tr><td>413</td><td>Request body exceeds size limit</td></tr>
|
||||||
|
<tr><td>429</td><td>Rate limit exceeded</td></tr>
|
||||||
|
<tr><td>502</td><td>Upstream network / DNS error</td></tr>
|
||||||
|
<tr><td>504</td><td>Upstream timeout</td></tr>
|
||||||
|
</table>
|
||||||
|
</main>
|
||||||
|
</body>
|
||||||
|
</html>`;
|
||||||
|
|
||||||
|
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<Response> {
|
||||||
|
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<string>;
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function handleRequest(
|
||||||
|
req: Request,
|
||||||
|
env: RouterEnv,
|
||||||
|
clientIP: string,
|
||||||
|
options: RouterOptions = {},
|
||||||
|
): Promise<Response | undefined> {
|
||||||
|
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;
|
||||||
|
}
|
||||||
@@ -80,7 +80,7 @@ describe("body-limiter", () => {
|
|||||||
headers: { "Content-Length": "200" },
|
headers: { "Content-Length": "200" },
|
||||||
});
|
});
|
||||||
const result = checkBodySize(request);
|
const result = checkBodySize(request);
|
||||||
const body = await result!.json();
|
const body = await result!.json() as Record<string, unknown>;
|
||||||
expect(body.error).toBe("Payload Too Large");
|
expect(body.error).toBe("Payload Too Large");
|
||||||
expect(body.maxSizeBytes).toBe(100);
|
expect(body.maxSizeBytes).toBe(100);
|
||||||
});
|
});
|
||||||
|
|||||||
@@ -10,76 +10,76 @@ import { createRateLimiter } from "./rate-limiter";
|
|||||||
|
|
||||||
describe("rate-limiter", () => {
|
describe("rate-limiter", () => {
|
||||||
describe("createRateLimiter", () => {
|
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 });
|
const limiter = createRateLimiter({ maxRequests: 5, windowMs: 60_000 });
|
||||||
for (let i = 0; i < 5; i++) {
|
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);
|
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 });
|
const limiter = createRateLimiter({ maxRequests: 3, windowMs: 60_000 });
|
||||||
for (let i = 0; i < 3; i++) {
|
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.allowed).toBe(false);
|
||||||
expect(blocked.retryAfterMs).toBeDefined();
|
expect(blocked.retryAfterMs).toBeDefined();
|
||||||
expect(typeof blocked.retryAfterMs).toBe("number");
|
expect(typeof blocked.retryAfterMs).toBe("number");
|
||||||
});
|
});
|
||||||
|
|
||||||
test("should return retryAfterMs when blocked", () => {
|
test("should return retryAfterMs when blocked", async () => {
|
||||||
const limiter = createRateLimiter({
|
const limiter = createRateLimiter({
|
||||||
maxRequests: 1,
|
maxRequests: 1,
|
||||||
windowMs: 60_000,
|
windowMs: 60_000,
|
||||||
});
|
});
|
||||||
limiter.check("retry-key");
|
await limiter.checkAsync("retry-key");
|
||||||
const blocked = limiter.check("retry-key");
|
const blocked = await limiter.checkAsync("retry-key");
|
||||||
expect(blocked.allowed).toBe(false);
|
expect(blocked.allowed).toBe(false);
|
||||||
expect(blocked.retryAfterMs).toBeGreaterThan(0);
|
expect(blocked.retryAfterMs).toBeGreaterThan(0);
|
||||||
expect(blocked.retryAfterMs).toBeLessThanOrEqual(60_000);
|
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 });
|
const limiter = createRateLimiter({ maxRequests: 2, windowMs: 60_000 });
|
||||||
limiter.check("reset-key");
|
await limiter.checkAsync("reset-key");
|
||||||
limiter.check("reset-key");
|
await limiter.checkAsync("reset-key");
|
||||||
// would be blocked, but...
|
// would be blocked, but...
|
||||||
limiter.reset("reset-key");
|
await limiter.resetAsync("reset-key");
|
||||||
// ...should be allowed again
|
// ...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 });
|
const limiter = createRateLimiter({ maxRequests: 2, windowMs: 60_000 });
|
||||||
expect(limiter.check("key-a").allowed).toBe(true);
|
expect((await limiter.checkAsync("key-a")).allowed).toBe(true);
|
||||||
expect(limiter.check("key-a").allowed).toBe(true);
|
expect((await limiter.checkAsync("key-a")).allowed).toBe(true);
|
||||||
expect(limiter.check("key-b").allowed).toBe(true); // key-b unaffected
|
expect((await limiter.checkAsync("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(false); // key-a blocked
|
||||||
});
|
});
|
||||||
|
|
||||||
test("should create with default options", () => {
|
test("should create with default options", () => {
|
||||||
const limiter = createRateLimiter();
|
const limiter = createRateLimiter();
|
||||||
expect(limiter.check).toBeDefined();
|
expect(limiter.checkAsync).toBeDefined();
|
||||||
expect(limiter.reset).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 });
|
const limiter = createRateLimiter({ maxRequests: 100, windowMs: 60_000 });
|
||||||
for (let i = 0; i < 100; i++) {
|
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 });
|
const limiter = createRateLimiter({ maxRequests: 1, windowMs: 60_000 });
|
||||||
limiter.check("after-reset-key");
|
await limiter.checkAsync("after-reset-key");
|
||||||
const blocked = limiter.check("after-reset-key");
|
const blocked = await limiter.checkAsync("after-reset-key");
|
||||||
expect(blocked.allowed).toBe(false);
|
expect(blocked.allowed).toBe(false);
|
||||||
limiter.reset("after-reset-key");
|
await limiter.resetAsync("after-reset-key");
|
||||||
expect(limiter.check("after-reset-key").allowed).toBe(true);
|
expect((await limiter.checkAsync("after-reset-key")).allowed).toBe(true);
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|||||||
@@ -11,11 +11,16 @@
|
|||||||
export interface RateLimiterOptions {
|
export interface RateLimiterOptions {
|
||||||
maxRequests?: number;
|
maxRequests?: number;
|
||||||
windowMs?: number;
|
windowMs?: number;
|
||||||
|
/** Optional async KV store for distributed rate limiting */
|
||||||
|
kv?: {
|
||||||
|
get(key: string): Promise<number[] | null>;
|
||||||
|
set(key: string, value: number[], expirationTtl?: number): Promise<void>;
|
||||||
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
export interface RateLimiter {
|
export interface RateLimiter {
|
||||||
check(key: string): { allowed: boolean; retryAfterMs?: number };
|
checkAsync(key: string): Promise<{ allowed: boolean; retryAfterMs?: number }>;
|
||||||
reset(key: string): void;
|
resetAsync(key: string): Promise<void>;
|
||||||
}
|
}
|
||||||
|
|
||||||
const DEFAULT_MAX_REQUESTS = 100;
|
const DEFAULT_MAX_REQUESTS = 100;
|
||||||
@@ -25,17 +30,17 @@ const CLEANUP_INTERVAL_DIVISOR = 10;
|
|||||||
export function createRateLimiter(options?: RateLimiterOptions): RateLimiter {
|
export function createRateLimiter(options?: RateLimiterOptions): RateLimiter {
|
||||||
const maxRequests = options?.maxRequests ?? DEFAULT_MAX_REQUESTS;
|
const maxRequests = options?.maxRequests ?? DEFAULT_MAX_REQUESTS;
|
||||||
const windowMs = options?.windowMs ?? DEFAULT_WINDOW_MS;
|
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<string, number[]>();
|
const store = new Map<string, number[]>();
|
||||||
|
|
||||||
|
// Maximum number of unique keys to track (prevents unbounded memory growth)
|
||||||
|
const MAX_STORE_KEYS = 10_000;
|
||||||
|
|
||||||
// ── helpers ──────────────────────────────────────────────────────
|
// ── helpers ──────────────────────────────────────────────────────
|
||||||
|
|
||||||
/** Remove timestamps outside the sliding window. Returns the pruned slice. */
|
function pruneTimestamps(timestamps: number[], now: number): number[] {
|
||||||
function prune(key: string, now: number): number[] {
|
|
||||||
const timestamps = store.get(key);
|
|
||||||
if (!timestamps) return [];
|
|
||||||
|
|
||||||
const cutoff = now - windowMs;
|
const cutoff = now - windowMs;
|
||||||
const result: number[] = [];
|
const result: number[] = [];
|
||||||
for (let i = 0; i < timestamps.length; i++) {
|
for (let i = 0; i < timestamps.length; i++) {
|
||||||
@@ -43,17 +48,12 @@ export function createRateLimiter(options?: RateLimiterOptions): RateLimiter {
|
|||||||
result.push(timestamps[i]);
|
result.push(timestamps[i]);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if (result.length === 0) {
|
|
||||||
store.delete(key);
|
|
||||||
} else {
|
|
||||||
store.set(key, result);
|
|
||||||
}
|
|
||||||
return result;
|
return result;
|
||||||
}
|
}
|
||||||
|
|
||||||
/** Periodically sweep the entire store to free memory. */
|
/** Periodically sweep the entire memory store to free memory. */
|
||||||
function periodicCleanup(): void {
|
function periodicCleanup(): void {
|
||||||
|
if (kv) return; // Cleanup handled by TTL in KV
|
||||||
const now = Date.now();
|
const now = Date.now();
|
||||||
const cutoff = now - windowMs;
|
const cutoff = now - windowMs;
|
||||||
|
|
||||||
@@ -70,30 +70,67 @@ export function createRateLimiter(options?: RateLimiterOptions): RateLimiter {
|
|||||||
store.set(key, pruned);
|
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)
|
if (!kv) {
|
||||||
const cleanupHandle = setInterval(
|
// Schedule periodic cleanup (every windowMs / 10)
|
||||||
periodicCleanup,
|
const cleanupHandle = setInterval(
|
||||||
windowMs / CLEANUP_INTERVAL_DIVISOR,
|
periodicCleanup,
|
||||||
);
|
windowMs / CLEANUP_INTERVAL_DIVISOR,
|
||||||
// Allow the process to exit even if the interval is still active
|
);
|
||||||
if (
|
// Allow the process to exit even if the interval is still active
|
||||||
cleanupHandle &&
|
if (
|
||||||
typeof cleanupHandle === "object" &&
|
cleanupHandle &&
|
||||||
"unref" in cleanupHandle
|
typeof cleanupHandle === "object" &&
|
||||||
) {
|
"unref" in cleanupHandle
|
||||||
(cleanupHandle as NodeJS.Timeout).unref();
|
) {
|
||||||
|
(cleanupHandle as NodeJS.Timeout).unref();
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── public API ───────────────────────────────────────────────────
|
// ── public API ───────────────────────────────────────────────────
|
||||||
|
|
||||||
return {
|
return {
|
||||||
check(key: string): { allowed: boolean; retryAfterMs?: number } {
|
async checkAsync(key: string): Promise<{ allowed: boolean; retryAfterMs?: number }> {
|
||||||
const now = Date.now();
|
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);
|
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) {
|
if (timestamps.length <= maxRequests) {
|
||||||
return { allowed: true };
|
return { allowed: true };
|
||||||
@@ -110,8 +147,12 @@ export function createRateLimiter(options?: RateLimiterOptions): RateLimiter {
|
|||||||
return { allowed: false, retryAfterMs };
|
return { allowed: false, retryAfterMs };
|
||||||
},
|
},
|
||||||
|
|
||||||
reset(key: string): void {
|
async resetAsync(key: string): Promise<void> {
|
||||||
store.delete(key);
|
if (kv) {
|
||||||
|
await kv.set(key, [], 1).catch(() => {});
|
||||||
|
} else {
|
||||||
|
store.delete(key);
|
||||||
|
}
|
||||||
},
|
},
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|||||||
+20
-12
@@ -84,8 +84,8 @@ function getClientIP(req: Request): string {
|
|||||||
|
|
||||||
// --- Auth Helper ---------------------------------------------------------------
|
// --- Auth Helper ---------------------------------------------------------------
|
||||||
|
|
||||||
function requireAuth(req: Request): Response | null {
|
function requireAuth(req: Request, env: Env): Response | null {
|
||||||
const apiKey = process.env.API_KEY ?? "sk-dummy-key";
|
const apiKey = env.API_KEY ?? "sk-dummy-key";
|
||||||
const header = req.headers.get("authorization") ?? req.headers.get("x-api-key") ?? "";
|
const header = req.headers.get("authorization") ?? req.headers.get("x-api-key") ?? "";
|
||||||
const key = header.replace(/^Bearer\s+/i, "").trim();
|
const key = header.replace(/^Bearer\s+/i, "").trim();
|
||||||
if (key === apiKey) return null;
|
if (key === apiKey) return null;
|
||||||
@@ -268,7 +268,7 @@ async function handleRelay(req: Request, env: Env): Promise<Response> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// -- Middleware: Rate limiting ---------------------------------------------
|
// -- Middleware: Rate limiting ---------------------------------------------
|
||||||
const rateCheck = limiter.check(clientIP);
|
const rateCheck = await limiter.checkAsync(clientIP);
|
||||||
if (!rateCheck.allowed) {
|
if (!rateCheck.allowed) {
|
||||||
logRelayEvent({
|
logRelayEvent({
|
||||||
method,
|
method,
|
||||||
@@ -387,12 +387,20 @@ export default {
|
|||||||
|
|
||||||
if (url.pathname === "/health") return handleHealth();
|
if (url.pathname === "/health") return handleHealth();
|
||||||
if (url.pathname === "/docs" || url.pathname === "/test") {
|
if (url.pathname === "/docs" || url.pathname === "/test") {
|
||||||
const file = Bun.file("public/test-api.html");
|
try {
|
||||||
const exists = await file.exists();
|
const file = Bun.file("public/test-api.html");
|
||||||
return new Response(exists ? file : "Not found", {
|
const exists = await file.exists();
|
||||||
status: exists ? 200 : 404,
|
return new Response(exists ? file : "Not found", {
|
||||||
headers: { "Content-Type": "text/html; charset=utf-8" },
|
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 (
|
if (
|
||||||
url.pathname === "/" &&
|
url.pathname === "/" &&
|
||||||
@@ -425,7 +433,7 @@ export default {
|
|||||||
if (url.pathname === "/v1/chat/completions") {
|
if (url.pathname === "/v1/chat/completions") {
|
||||||
if (req.method === "OPTIONS") return createCorsPreflightResponse();
|
if (req.method === "OPTIONS") return createCorsPreflightResponse();
|
||||||
if (req.method !== "POST") return new Response("Method Not Allowed", { status: 405 });
|
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;
|
if (authErr) return authErr;
|
||||||
try {
|
try {
|
||||||
const body = await req.json();
|
const body = await req.json();
|
||||||
@@ -441,7 +449,7 @@ export default {
|
|||||||
if (url.pathname === "/v1/messages") {
|
if (url.pathname === "/v1/messages") {
|
||||||
if (req.method === "OPTIONS") return createCorsPreflightResponse();
|
if (req.method === "OPTIONS") return createCorsPreflightResponse();
|
||||||
if (req.method !== "POST") return new Response("Method Not Allowed", { status: 405 });
|
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;
|
if (authErr) return authErr;
|
||||||
try {
|
try {
|
||||||
const body = await req.json();
|
const body = await req.json();
|
||||||
@@ -455,7 +463,7 @@ export default {
|
|||||||
}
|
}
|
||||||
|
|
||||||
if (url.pathname === "/v1/models" && req.method === "GET") {
|
if (url.pathname === "/v1/models" && req.method === "GET") {
|
||||||
const authErr = requireAuth(req);
|
const authErr = requireAuth(req, env);
|
||||||
if (authErr) return authErr;
|
if (authErr) return authErr;
|
||||||
const models = listModels().map((id) => ({
|
const models = listModels().map((id) => ({
|
||||||
id,
|
id,
|
||||||
|
|||||||
Reference in New Issue
Block a user