docfast/src/__tests__/browser-url-ssrf.test.ts

import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";

/**
 * TDD tests for renderUrlPdf SSRF protection (DNS pinning via hostResolverRules).
 * Covers branches in browser.ts lines ~300-331:
 * - HTTP request rewrite to pinned IP
 * - HTTPS request passthrough
 * - Blocking requests to non-target hosts
 * - No interception when hostResolverRules absent
 * - No interception when hostResolverRules doesn't match MAP regex
 * - Blocking cloud metadata endpoint requests
 */

vi.unmock("../services/browser.js");

vi.mock("../services/logger.js", () => ({
  default: { info: vi.fn(), warn: vi.fn(), error: vi.fn() },
}));

function createMockPage(overrides: Record<string, any> = {}) {
  return {
    setJavaScriptEnabled: vi.fn().mockResolvedValue(undefined),
    setContent: vi.fn().mockResolvedValue(undefined),
    addStyleTag: vi.fn().mockResolvedValue(undefined),
    pdf: vi.fn().mockResolvedValue(Buffer.from("%PDF-1.4 test")),
    goto: vi.fn().mockResolvedValue(undefined),
    close: vi.fn().mockResolvedValue(undefined),
    setRequestInterception: vi.fn().mockResolvedValue(undefined),
    removeAllListeners: vi.fn().mockReturnThis(),
    createCDPSession: vi.fn().mockResolvedValue({
      send: vi.fn().mockResolvedValue(undefined),
      detach: vi.fn().mockResolvedValue(undefined),
    }),
    cookies: vi.fn().mockResolvedValue([]),
    deleteCookie: vi.fn(),
    on: vi.fn(),
    ...overrides,
  } as any;
}

function createMockBrowser(pagesPerBrowser = 2) {
  const pages = Array.from({ length: pagesPerBrowser }, () => createMockPage());
  let pageIndex = 0;
  return {
    newPage: vi.fn().mockImplementation(() => Promise.resolve(pages[pageIndex++] || createMockPage())),
    close: vi.fn().mockResolvedValue(undefined),
    _pages: pages,
  } as any;
}

process.env.BROWSER_COUNT = "1";
process.env.PAGES_PER_BROWSER = "2";

describe("renderUrlPdf SSRF DNS pinning", () => {
  let browserModule: typeof import("../services/browser.js");
  let mockBrowsers: any[] = [];

  beforeEach(async () => {
    mockBrowsers = [];
    vi.resetModules();
    vi.doMock("puppeteer", () => ({
      default: {
        launch: vi.fn().mockImplementation(() => {
          const b = createMockBrowser(2);
          mockBrowsers.push(b);
          return Promise.resolve(b);
        }),
      },
    }));
    vi.doMock("../services/logger.js", () => ({
      default: { info: vi.fn(), warn: vi.fn(), error: vi.fn() },
    }));
    browserModule = await import("../services/browser.js");
    await browserModule.initBrowser();
  });

  afterEach(async () => {
    try { await browserModule.closeBrowser(); } catch {}
  });

  function getUsedPage() {
    return mockBrowsers
      .flatMap((b: any) => b._pages)
      .find((p: any) => p.on.mock.calls.some((c: any) => c[0] === "request"));
  }

  function getRequestHandler(page: any) {
    const call = page.on.mock.calls.find((c: any) => c[0] === "request");
    return call ? call[1] : null;
  }

  it("sets up request interception when valid MAP rule provided", async () => {
    await browserModule.renderUrlPdf("http://example.com", {
      hostResolverRules: "MAP example.com 93.184.216.34",
    });

    const page = getUsedPage();
    expect(page).toBeDefined();
    expect(page.setRequestInterception).toHaveBeenCalledWith(true);
  });

  it("rewrites HTTP requests to use pinned IP with Host header", async () => {
    await browserModule.renderUrlPdf("http://example.com/path", {
      hostResolverRules: "MAP example.com 93.184.216.34",
    });

    const handler = getRequestHandler(getUsedPage());
    expect(handler).not.toBeNull();

    const req = {
      url: () => "http://example.com/path?q=1",
      headers: () => ({ accept: "text/html" }),
      continue: vi.fn(),
      abort: vi.fn(),
    };
    handler(req);

    expect(req.continue).toHaveBeenCalledWith({
      url: "http://93.184.216.34/path?q=1",
      headers: { accept: "text/html", host: "example.com" },
    });
    expect(req.abort).not.toHaveBeenCalled();
  });

  it("continues HTTPS requests without URL rewrite (cert compatibility)", async () => {
    await browserModule.renderUrlPdf("https://example.com", {
      hostResolverRules: "MAP example.com 93.184.216.34",
    });

    const handler = getRequestHandler(getUsedPage());
    expect(handler).not.toBeNull();

    const req = {
      url: () => "https://example.com/page",
      headers: () => ({}),
      continue: vi.fn(),
      abort: vi.fn(),
    };
    handler(req);

    expect(req.continue).toHaveBeenCalledWith();
    expect(req.abort).not.toHaveBeenCalled();
  });

  it("aborts requests to non-target hosts (prevents redirect SSRF)", async () => {
    await browserModule.renderUrlPdf("http://example.com", {
      hostResolverRules: "MAP example.com 93.184.216.34",
    });

    const handler = getRequestHandler(getUsedPage());
    const req = {
      url: () => "http://evil.internal/admin",
      headers: () => ({}),
      continue: vi.fn(),
      abort: vi.fn(),
    };
    handler(req);

    expect(req.abort).toHaveBeenCalledWith("blockedbyclient");
    expect(req.continue).not.toHaveBeenCalled();
  });

  it("blocks cloud metadata endpoint (169.254.169.254)", async () => {
    await browserModule.renderUrlPdf("http://example.com", {
      hostResolverRules: "MAP example.com 93.184.216.34",
    });

    const handler = getRequestHandler(getUsedPage());
    const req = {
      url: () => "http://169.254.169.254/latest/meta-data/",
      headers: () => ({}),
      continue: vi.fn(),
      abort: vi.fn(),
    };
    handler(req);

    expect(req.abort).toHaveBeenCalledWith("blockedbyclient");
  });

  it("does NOT set up interception when hostResolverRules is absent", async () => {
    await browserModule.renderUrlPdf("http://example.com");

    const page = mockBrowsers.flatMap((b: any) => b._pages)
      .find((p: any) => p.goto.mock.calls.length > 0);
    expect(page).toBeDefined();

    // Should not have called setRequestInterception(true) — only false from recyclePage
    const trueCalls = page.setRequestInterception.mock.calls.filter(
      (c: any) => c[0] === true
    );
    expect(trueCalls.length).toBe(0);
  });

  it("skips DNS pinning when hostResolverRules doesn't match MAP regex", async () => {
    await browserModule.renderUrlPdf("http://example.com", {
      hostResolverRules: "BADFORMAT",
    });

    // CDP session was created (Network.enable), but no request interception
    const page = mockBrowsers.flatMap((b: any) => b._pages)
      .find((p: any) => p.goto.mock.calls.length > 0);
    const trueCalls = page.setRequestInterception.mock.calls.filter(
      (c: any) => c[0] === true
    );
    expect(trueCalls.length).toBe(0);
    // No request handler registered
    const requestCalls = page.on.mock.calls.filter((c: any) => c[0] === "request");
    expect(requestCalls.length).toBe(0);
  });
});