Harden the security layer (DNS pinning, rate-limit key, body cap, scan budget) and add a test suite

.github/workflows/ci.yml

@@ -27,6 +27,9 @@ jobs:
       - name: Typecheck
         run: npm run typecheck
 
+      - name: Test
+        run: npm run test
+
       - name: Lint
         run: npm run lint
 

components/ToolCard.tsx

@@ -3,18 +3,13 @@ import type { Tool } from "@/lib/tools/registry";
 
 export function ToolCard({ tool }: { tool: Tool }) {
   const isLive = tool.status === "live";
-  const Wrapper: any = isLive ? Link : "div";
-  const wrapperProps = isLive ? { href: tool.href } : {};
-
-  return (
-    <Wrapper
-      {...wrapperProps}
-      className={`group relative block rounded-2xl border p-5 transition ${
-        isLive
-          ? "border-ink-700 bg-ink-900/60 hover:border-accent-500/60 hover:bg-ink-800/80"
-          : "border-ink-700/60 bg-ink-900/30 opacity-60 cursor-not-allowed"
-      }`}
-    >
+  const className = `group relative block rounded-2xl border p-5 transition ${
+    isLive
+      ? "border-ink-700 bg-ink-900/60 hover:border-accent-500/60 hover:bg-ink-800/80"
+      : "border-ink-700/60 bg-ink-900/30 opacity-60 cursor-not-allowed"
+  }`;
+  const content = (
+    <>
       <div className="flex items-start justify-between gap-3">
         <div>
           <h3 className="text-lg font-semibold tracking-tight text-slate-100 group-hover:text-accent-400 transition">
@@ -45,6 +40,14 @@ export function ToolCard({ tool }: { tool: Tool }) {
           ))}
         </div>
       )}
-    </Wrapper>
+    </>
+  );
+
+  return isLive ? (
+    <Link href={tool.href} className={className}>
+      {content}
+    </Link>
+  ) : (
+    <div className={className}>{content}</div>
   );
 }

lib/cors/scan.ts

@@ -11,7 +11,7 @@ type ProbeSpec = {
   method: ProbeMethod;
 };
 
-function buildProbes(targetHost: string, isHttps: boolean): ProbeSpec[] {
+export function buildProbes(targetHost: string, isHttps: boolean): ProbeSpec[] {
   const scheme = isHttps ? "https" : "http";
   const probes: ProbeSpec[] = [
     { id: "baseline", label: "Baseline (no Origin)", origin: null, method: "GET" },
@@ -24,13 +24,13 @@ function buildProbes(targetHost: string, isHttps: boolean): ProbeSpec[] {
     { id: "null-origin", label: "Origin: null", origin: "null", method: "GET" },
     {
       id: "suffix-bypass",
-      label: "Suffix bypass — attacker domain ends in target",
-      origin: `https://attacker${targetHost.replace(/\./g, "")}.example`,
+      label: "Suffix bypass (origin ends with target host)",
+      origin: `https://not${targetHost}`,
       method: "GET",
     },
     {
       id: "prefix-bypass",
-      label: "Prefix bypass — target as a subdomain of attacker",
+      label: "Prefix bypass (origin starts with target host)",
       origin: `https://${targetHost}.${ATTACKER_HOST}`,
       method: "GET",
     },
@@ -58,15 +58,19 @@ function buildProbes(targetHost: string, isHttps: boolean): ProbeSpec[] {
   return probes;
 }
 
-async function runProbe(url: string, spec: ProbeSpec): Promise<ProbeResult> {
+async function runProbe(url: string, spec: ProbeSpec, deadlineMs?: number): Promise<ProbeResult> {
   const headers: Record<string, string> = {};
   if (spec.origin !== null) headers["Origin"] = spec.origin;
   if (spec.method === "OPTIONS") {
     headers["Access-Control-Request-Method"] = "PUT";
     headers["Access-Control-Request-Headers"] = "authorization,content-type";
   }
 
-  const res = await safeFetch(url, { method: spec.method, headers });
+  const res = await safeFetch(url, {
+    method: spec.method,
+    headers,
+    ...(deadlineMs === undefined ? {} : { deadlineMs }),
+  });
   if (!res.ok) {
     return {
       id: spec.id,
@@ -101,7 +105,7 @@ function isCredentialed(p: ProbeResult): boolean {
   return p.acac?.toLowerCase() === "true";
 }
 
-function analyze(probes: ProbeResult[]): FindingGroup[] {
+export function analyze(probes: ProbeResult[]): FindingGroup[] {
   const findings: Finding[] = [];
   const byId = new Map(probes.map((p) => [p.id, p]));
 
@@ -197,7 +201,7 @@ function analyze(probes: ProbeResult[]): FindingGroup[] {
 
   // Vary: Origin hygiene
   for (const p of probes) {
-    if (p.acao && p.acao !== "*" && p.id === "arbitrary-origin") {
+    if (p.acao && p.acao !== "*" && p.reflectsOrigin) {
       const varyHasOrigin = !!p.vary?.split(",").some((v) => v.trim().toLowerCase() === "origin");
       if (!varyHasOrigin) {
         findings.push({
@@ -209,7 +213,6 @@ function analyze(probes: ProbeResult[]): FindingGroup[] {
           value: `Vary=${p.vary ?? "(none)"}`,
           recommendation: "Add 'Vary: Origin' whenever ACAO is computed from the request.",
         });
-        break;
       }
     }
   }
@@ -273,14 +276,15 @@ function summarise(groups: FindingGroup[]): CorsReport["summary"] {
 export async function runCorsScan(
   input: string,
 ): Promise<{ ok: true; report: CorsReport } | { ok: false; reason: string }> {
-  const initial = await safeFetch(input, { method: "GET" });
+  const deadlineMs = Date.now() + 18_000;
+  const initial = await safeFetch(input, { method: "GET", deadlineMs });
   if (!initial.ok) return { ok: false, reason: initial.reason };
   const finalUrl = initial.data.finalUrl;
   const targetUrl = new URL(finalUrl);
   const isHttps = targetUrl.protocol === "https:";
 
   const specs = buildProbes(targetUrl.hostname, isHttps);
-  const probes = await Promise.all(specs.map((s) => runProbe(finalUrl, s)));
+  const probes = await Promise.all(specs.map((s) => runProbe(finalUrl, s, deadlineMs)));
   const groups = analyze(probes);
 
   return {

lib/misconfig/scan.ts

@@ -40,7 +40,8 @@ export async function runScan(input: string): Promise<
   | { ok: true; report: ScanReport }
   | { ok: false; reason: string }
 > {
-  const initial = await safeFetch(input, { method: "GET" });
+  const deadlineMs = Date.now() + 13_000;
+  const initial = await safeFetch(input, { method: "GET", deadlineMs });
   if (!initial.ok) return { ok: false, reason: initial.reason };
 
   const { finalUrl, status, headers, redirects, responseTimeMs } = initial.data;
@@ -65,7 +66,7 @@ export async function runScan(input: string): Promise<
     {
       id: "exposures",
       title: "File Exposure Probes",
-      findings: await runProbes(finalUrl),
+      findings: await runProbes(finalUrl, deadlineMs),
     },
   ];
 

lib/security/rate-limit.ts

@@ -4,9 +4,16 @@ const WINDOW_MS = 60_000;
 const MAX_PER_WINDOW = 12;
 
 const buckets = new Map<string, Bucket>();
+let lastSweepAt = 0;
 
 export function rateLimit(key: string): { ok: true } | { ok: false; retryAfterSec: number } {
   const now = Date.now();
+  if (now - lastSweepAt >= WINDOW_MS || buckets.size > 1000) {
+    for (const [bucketKey, value] of buckets) {
+      if (value.resetAt <= now) buckets.delete(bucketKey);
+    }
+    lastSweepAt = now;
+  }
   const bucket = buckets.get(key);
   if (!bucket || bucket.resetAt <= now) {
     buckets.set(key, { count: 1, resetAt: now + WINDOW_MS });
@@ -20,7 +27,14 @@ export function rateLimit(key: string): { ok: true } | { ok: false; retryAfterSe
 }
 
 export function clientKeyFromHeaders(headers: Headers): string {
+  const vercel = headers.get("x-vercel-forwarded-for")?.trim();
+  if (vercel) return vercel;
+  const realIp = headers.get("x-real-ip")?.trim();
+  if (realIp) return realIp;
   const xff = headers.get("x-forwarded-for");
-  if (xff) return xff.split(",")[0].trim();
-  return headers.get("x-real-ip") ?? "unknown";
+  if (xff) {
+    const entries = xff.split(",").map((value) => value.trim()).filter(Boolean);
+    if (entries.length > 0) return entries[entries.length - 1];
+  }
+  return "unknown";
 }

lib/security/safe-fetch.ts

@@ -1,4 +1,13 @@
-import { guardUrl } from "./ssrf";
+import http from "node:http";
+import https from "node:https";
+import type { IncomingMessage } from "node:http";
+import { guardUrl, pinnedLookup } from "./ssrf";
+
+export type SafeFetchInit = {
+  method?: string;
+  headers?: Record<string, string>;
+  deadlineMs?: number;
+};
 
 export type SafeFetchResult = {
   finalUrl: string;
@@ -14,79 +23,114 @@ const MAX_REDIRECTS = 3;
 const MAX_BODY_BYTES = 256 * 1024;
 const TIMEOUT_MS = 6000;
 
-async function readLimitedText(res: Response, max = MAX_BODY_BYTES): Promise<string> {
-  const reader = res.body?.getReader();
-  if (!reader) return "";
-  const decoder = new TextDecoder();
+async function readLimitedText(res: IncomingMessage, max = MAX_BODY_BYTES): Promise<string> {
+  const chunks: Buffer[] = [];
   let received = 0;
-  let out = "";
-  while (true) {
-    const { done, value } = await reader.read();
-    if (done) break;
-    received += value.byteLength;
-    out += decoder.decode(value, { stream: true });
-    if (received >= max) {
-      try {
-        await reader.cancel();
-      } catch {
-        // ignore
+  for await (const chunk of res) {
+    if (received < max) {
+      const bytes = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk);
+      const allowed = Math.min(bytes.length, max - received);
+      if (allowed > 0) {
+        chunks.push(bytes.subarray(0, allowed));
+        received += allowed;
+      }
+      if (received >= max) {
+        res.resume();
+        break;
       }
-      break;
     }
   }
-  out += decoder.decode();
-  return out;
+  return Buffer.concat(chunks).toString("utf8");
 }
 
 export async function safeFetch(
   inputUrl: string,
-  init: RequestInit = {},
+  init: SafeFetchInit = {},
 ): Promise<{ ok: true; data: SafeFetchResult } | { ok: false; reason: string }> {
   let current = inputUrl;
   const redirects: string[] = [];
   const start = Date.now();
 
   for (let i = 0; i <= MAX_REDIRECTS; i++) {
+    if (init.deadlineMs !== undefined && init.deadlineMs <= Date.now()) {
+      return { ok: false, reason: "Scan time budget exhausted." };
+    }
     const guard = await guardUrl(current);
     if (!guard.ok) return { ok: false, reason: guard.reason };
+    const remaining = init.deadlineMs === undefined ? TIMEOUT_MS : init.deadlineMs - Date.now();
+    if (remaining <= 0) {
+      return { ok: false, reason: "Scan time budget exhausted." };
+    }
 
-    let res: Response;
-    try {
-      res = await fetch(guard.url.toString(), {
-        ...init,
-        headers: {
-          "User-Agent": USER_AGENT,
-          Accept: "*/*",
-          ...(init.headers ?? {}),
+    const timeoutMs = Math.min(TIMEOUT_MS, remaining);
+    const headers = {
+      "User-Agent": USER_AGENT,
+      Accept: "*/*",
+      ...(init.headers ?? {}),
+    };
+    const requestFn = guard.url.protocol === "https:" ? https.request : http.request;
+    const response = await new Promise<
+      | { ok: true; status: number; headers: Headers; body: string }
+      | { ok: false; reason: string }
+    >((resolve) => {
+      const req = requestFn(
+        guard.url,
+        {
+          method: init.method ?? "GET",
+          headers,
+          lookup: pinnedLookup(guard.addresses),
+          ...(guard.url.protocol === "https:" ? { servername: guard.url.hostname } : {}),
+        },
+        (res) => {
+          const responseHeaders = new Headers();
+          for (const [name, value] of Object.entries(res.headers)) {
+            if (name === "set-cookie" && Array.isArray(value)) {
+              for (const cookie of value) responseHeaders.append(name, cookie);
+            } else if (typeof value === "string") {
+              responseHeaders.append(name, value);
+            } else if (Array.isArray(value)) {
+              responseHeaders.append(name, value.join(", "));
+            }
+          }
+          const body = readLimitedText(res);
+          body.then((text) => {
+            clearTimeout(timer);
+            resolve({
+              ok: true,
+              status: res.statusCode ?? 0,
+              headers: responseHeaders,
+              body: text,
+            });
+          }).catch((error: unknown) => {
+            clearTimeout(timer);
+            const msg = error instanceof Error ? error.message : "response read failed";
+            resolve({ ok: false, reason: `Network error: ${msg}` });
+          });
         },
-        redirect: "manual",
-        signal: AbortSignal.timeout(TIMEOUT_MS),
+      );
+      const timer = setTimeout(() => req.destroy(new Error("Request timed out.")), timeoutMs);
+      req.on("error", (error) => {
+        clearTimeout(timer);
+        resolve({ ok: false, reason: `Network error: ${error.message}` });
       });
-    } catch (e: unknown) {
-      const msg = e instanceof Error ? e.message : "fetch failed";
-      return { ok: false, reason: `Network error: ${msg}` };
-    }
+      req.end();
+    });
+    if (!response.ok) return response;
 
-    if (res.status >= 300 && res.status < 400 && res.headers.get("location")) {
-      const next = new URL(res.headers.get("location")!, guard.url).toString();
+    if (response.status >= 300 && response.status < 400 && response.headers.get("location")) {
+      const next = new URL(response.headers.get("location")!, guard.url).toString();
       redirects.push(next);
       current = next;
-      try {
-        await res.body?.cancel();
-      } catch {
-        // ignore
-      }
       continue;
     }
 
-    const body = await readLimitedText(res);
     return {
       ok: true,
       data: {
         finalUrl: guard.url.toString(),
-        status: res.status,
-        headers: res.headers,
-        body,
+        status: response.status,
+        headers: response.headers,
+        body: response.body,
         redirects,
         responseTimeMs: Date.now() - start,
       },