feat(proxy): add challenge-aware response policy

2026-08-17 12:11:23 +02:00 · 2026-07-25 12:47:26 +02:00
parent ce028e201c
commit d7cca3ec44
4 changed files with 247 additions and 0 deletions
@@ -0,0 +1,93 @@
+import { describe, expect, test } from "bun:test"
+import { ChallengeCache } from "../challengeCache"
+
+describe("ChallengeCache", () => {
+  test("returns undefined for unknown host", () => {
+    const c = new ChallengeCache()
+    expect(c.get("never-seen.example.com")).toBeUndefined()
+  })
+
+  test("set + get round-trips modes", () => {
+    const c = new ChallengeCache()
+    c.set("a.example.com", "direct")
+    c.set("b.example.com", "cf")
+    c.set("c.example.com", "unknown")
+    expect(c.get("a.example.com")).toBe("direct")
+    expect(c.get("b.example.com")).toBe("cf")
+    expect(c.get("c.example.com")).toBe("unknown")
+  })
+
+  test("respects TTL — entries past TTL are not returned", () => {
+    const c = new ChallengeCache({ ttlMs: 10 })
+    c.set("host.example", "direct")
+    expect(c.get("host.example")).toBe("direct")
+    // Sleep past TTL
+    const deadline = Date.now() + 50
+    while (Date.now() < deadline) {
+      // tight spin — 10ms ttl + 50ms margin
+    }
+    expect(c.get("host.example")).toBeUndefined()
+  })
+
+  test("prune() evicts expired entries and returns count", () => {
+    const c = new ChallengeCache({ ttlMs: 5 })
+    c.set("a.example", "direct")
+    c.set("b.example", "cf")
+    c.set("c.example", "direct")
+    const deadline = Date.now() + 30
+    while (Date.now() < deadline) {
+      // wait
+    }
+    const removed = c.prune()
+    expect(removed).toBe(3)
+    expect(c.size()).toBe(0)
+  })
+
+  test("delete() removes a specific entry", () => {
+    const c = new ChallengeCache()
+    c.set("x.example", "direct")
+    c.set("y.example", "cf")
+    c.delete("x.example")
+    expect(c.get("x.example")).toBeUndefined()
+    expect(c.get("y.example")).toBe("cf")
+  })
+
+  test("clear() empties everything", () => {
+    const c = new ChallengeCache()
+    c.set("a", "direct")
+    c.set("b", "cf")
+    expect(c.size()).toBe(2)
+    c.clear()
+    expect(c.size()).toBe(0)
+    expect(c.get("a")).toBeUndefined()
+    expect(c.get("b")).toBeUndefined()
+  })
+
+  test("size() reflects current entries", () => {
+    const c = new ChallengeCache()
+    expect(c.size()).toBe(0)
+    c.set("a", "direct")
+    c.set("b", "direct")
+    c.set("c", "cf")
+    expect(c.size()).toBe(3)
+    c.delete("a")
+    expect(c.size()).toBe(2)
+  })
+
+  test("set() refreshes lastCheck (extends TTL window)", () => {
+    const c = new ChallengeCache({ ttlMs: 50 })
+    c.set("host", "direct")
+    // Wait partway through TTL
+    const midDeadline = Date.now() + 25
+    while (Date.now() < midDeadline) {
+      // tight spin
+    }
+    c.set("host", "direct") // refresh timestamp
+    // Wait past original TTL but within refreshed window
+    const endDeadline = Date.now() + 35
+    while (Date.now() < endDeadline) {
+      // tight spin
+    }
+    expect(c.get("host")).toBe("direct") // refreshed — still valid
+  })
+})
@@ -0,0 +1,58 @@
+import { describe, expect, test } from "bun:test"
+import type { ScrapeResult } from "@trawl/types"
+import { responseFromScrapeResult } from "../responsePolicy"
+
+function result(overrides: Partial<ScrapeResult>): ScrapeResult {
+  return {
+    url: "https://example.test/",
+    html: "",
+    cookies: [],
+    userAgent: "test",
+    statusCode: 200,
+    tier: 3,
+    sessionCached: false,
+    timings: [],
+    totalMs: 1,
+    ...overrides,
+  }
+}
+
+describe("responseFromScrapeResult", () => {
+  test("returns rendered HTML after a browser tier instead of the raw challenge response", () => {
+    const response = responseFromScrapeResult(
+      result({
+        html: "<html><title>Real page</title></html>",
+        body: Buffer.from("<html><title>Just a moment...</title></html>"),
+        contentType: "text/html; charset=utf-8",
+        responseHeaders: {
+          "content-type": "text/html; charset=utf-8",
+          "content-encoding": "br",
+          "content-length": "999",
+        },
+      }),
+    )
+
+    expect(response.body.toString()).toContain("Real page")
+    expect(response.body.toString()).not.toContain("Just a moment")
+    expect(response.headers["content-encoding"]).toBeUndefined()
+    expect(response.headers["content-length"]).toBeUndefined()
+  })
+
+  test("preserves raw bytes for binary responses", () => {
+    const bytes = Uint8Array.from([0, 255, 1, 2, 3])
+    const response = responseFromScrapeResult(
+      result({
+        html: "",
+        body: bytes,
+        contentType: "application/octet-stream",
+        responseHeaders: {
+          "content-type": "application/octet-stream",
+          "content-range": "bytes 0-4/100",
+        },
+      }),
+    )
+
+    expect([...response.body]).toEqual([...bytes])
+    expect(response.headers["content-range"]).toBe("bytes 0-4/100")
+  })
+})
@@ -0,0 +1,54 @@
+// Short-lived routing memory for hosts that recently served a challenge.
+
+export type ChallengeMode = "direct" | "cf" | "unknown"
+
+interface CacheEntry {
+  mode: ChallengeMode
+  lastCheck: number
+}
+
+export class ChallengeCache {
+  private readonly entries = new Map<string, CacheEntry>()
+  private readonly ttlMs: number
+
+  constructor(opts: { ttlMs?: number } = {}) {
+    this.ttlMs = opts.ttlMs ?? 5 * 60 * 1000 // 5 minutes
+  }
+
+  get(hostname: string): ChallengeMode | undefined {
+    const entry = this.entries.get(hostname)
+    if (!entry) return
+    if (Date.now() - entry.lastCheck > this.ttlMs) {
+      this.entries.delete(hostname)
+      return
+    }
+    return entry.mode
+  }
+
+  set(hostname: string, mode: ChallengeMode): void {
+    this.entries.set(hostname, { mode, lastCheck: Date.now() })
+  }
+
+  delete(hostname: string): void {
+    this.entries.delete(hostname)
+  }
+
+  clear(): void {
+    this.entries.clear()
+  }
+
+  size(): number {
+    return this.entries.size
+  }
+
+  prune(): number {
+    const now = Date.now()
+    let removed = 0
+    for (const [host, entry] of this.entries) {
+      if (now - entry.lastCheck <= this.ttlMs) continue
+      this.entries.delete(host)
+      removed++
+    }
+    return removed
+  }
+}
@@ -0,0 +1,42 @@
+import type { ScrapeResult } from "@trawl/types"
+
+export interface ProxyBufferedResponse {
+  body: Buffer
+  contentType: string
+  headers: Record<string, string>
+}
+
+const TRANSFORMED_BODY_HEADERS = new Set([
+  "content-encoding",
+  "content-length",
+  "content-md5",
+  "content-range",
+  "accept-ranges",
+  "etag",
+  "transfer-encoding",
+])
+
+function isHtml(contentType: string): boolean {
+  const base = contentType.split(";", 1)[0]?.trim().toLowerCase()
+  return base === "text/html" || base === "application/xhtml+xml"
+}
+
+export function responseFromScrapeResult(result: ScrapeResult): ProxyBufferedResponse {
+  const contentType = result.contentType ?? result.responseHeaders?.["content-type"] ?? "text/html; charset=utf-8"
+  const useRenderedHtml = isHtml(contentType) && result.html.length > 0
+  const body = useRenderedHtml
+    ? Buffer.from(result.html, "utf8")
+    : result.body
+      ? Buffer.from(result.body)
+      : Buffer.from(result.html, "utf8")
+
+  const headers: Record<string, string> = {}
+  for (const [name, value] of Object.entries(result.responseHeaders ?? {})) {
+    const lower = name.toLowerCase()
+    if (useRenderedHtml && TRANSFORMED_BODY_HEADERS.has(lower)) continue
+    headers[lower] = value
+  }
+  headers["content-type"] = contentType
+
+  return { body, contentType, headers }
+}