OmniRoute/tests/unit/nanobanana-image-handler.test.ts

import test from "node:test";
import assert from "node:assert/strict";
import dns from "node:dns";

import { handleImageGeneration } from "../../open-sse/handlers/imageGeneration.ts";

// Stub DNS for fetchRemoteImage's GHSA-cmhj-wh2f-9cgx DNS-rebinding guard
// (assertHostnameResolvesPublic in src/shared/network/remoteImageFetch.ts).
// The b64_json test mocks globalThis.fetch with an example.com URL that
// doesn't resolve in CI; the handler invokes fetchRemoteImage without
// exposing a `lookup` injection point, so we monkey-patch dns.promises.lookup
// to always return a public IP so the rebinding guard passes and the test
// exercises the mocked fetch behaviour as intended.
const originalDnsLookup = dns.promises.lookup;
(dns.promises as { lookup: unknown }).lookup = (async (
  _hostname: string,
  options?: { all?: boolean }
) => {
  const record = { address: "203.0.113.1", family: 4 };
  return options && options.all ? [record] : record;
}) as typeof dns.promises.lookup;
process.on("exit", () => {
  (dns.promises as { lookup: unknown }).lookup = originalDnsLookup;
});

test("handleImageGeneration(nanobanana): async submit+poll returns URL payload", async () => {
  const originalFetch = globalThis.fetch;
  let pollCount = 0;

  globalThis.fetch = async (url, options = {}) => {
    const u = String(url);

    if (u.includes("/generate-pro")) {
      const body = JSON.parse(options.body);
      assert.equal(body.prompt, "galaxy test");
      assert.equal(body.aspectRatio, "1:1");
      assert.equal(body.resolution, "2K");

      return new Response(
        JSON.stringify({ code: 200, msg: "success", data: { taskId: "task-handler-1" } }),
        { status: 200, headers: { "content-type": "application/json" } }
      );
    }

    if (u.includes("/record-info")) {
      pollCount += 1;
      if (pollCount < 2) {
        return new Response(
          JSON.stringify({ code: 200, msg: "success", data: { successFlag: 0 } }),
          {
            status: 200,
            headers: { "content-type": "application/json" },
          }
        );
      }

      return new Response(
        JSON.stringify({
          code: 200,
          msg: "success",
          data: {
            successFlag: 1,
            response: { resultImageUrl: "https://cdn.example.com/handler-result.jpg" },
          },
        }),
        { status: 200, headers: { "content-type": "application/json" } }
      );
    }

    throw new Error(`Unexpected URL: ${u}`);
  };

  try {
    const result = await handleImageGeneration({
      body: {
        model: "nanobanana/nanobanana-pro",
        prompt: "galaxy test",
        size: "1024x1280",
        poll_interval_ms: 1,
      },
      credentials: { apiKey: "test-key" },
      log: null,
    });

    assert.equal(result.success, true);
    assert.equal(result.data.data.length, 1);
    assert.equal(result.data.data[0].url, "https://cdn.example.com/handler-result.jpg");
  } finally {
    globalThis.fetch = originalFetch;
  }
});

test("handleImageGeneration(nanobanana): response_format=b64_json converts URL to b64", async () => {
  const originalFetch = globalThis.fetch;

  globalThis.fetch = async (url) => {
    const u = String(url);

    if (u.includes("/generate")) {
      return new Response(
        JSON.stringify({ code: 200, msg: "success", data: { taskId: "task-handler-2" } }),
        { status: 200, headers: { "content-type": "application/json" } }
      );
    }

    if (u.includes("/record-info")) {
      return new Response(
        JSON.stringify({
          code: 200,
          msg: "success",
          data: {
            successFlag: 1,
            response: { resultImageUrl: "https://cdn.example.com/handler-result-2.jpg" },
          },
        }),
        { status: 200, headers: { "content-type": "application/json" } }
      );
    }

    if (u.includes("handler-result-2.jpg")) {
      return new Response(new Uint8Array([0x89, 0x50, 0x4e, 0x47]), { status: 200 });
    }

    throw new Error(`Unexpected URL: ${u}`);
  };

  try {
    const result = await handleImageGeneration({
      body: {
        model: "nanobanana/nanobanana-flash",
        prompt: "galaxy test",
        response_format: "b64_json",
      },
      credentials: { apiKey: "test-key" },
      log: null,
    });

    assert.equal(result.success, true);
    assert.equal(result.data.data.length, 1);
    assert.equal(result.data.data[0].b64_json, "iVBORw==");
  } finally {
    globalThis.fetch = originalFetch;
  }
});

// GHSA-34rg-3pqj-35g9 — the `response_format=b64_json` path re-fetches the result URL the
// upstream task reports. That URL is upstream-supplied (lower risk than a request-body
// URL), but it went through `fetchRemoteImage()` with no explicit `guard`, i.e. under the
// OPERATOR outbound policy (`block-metadata` on a default install: LAN/loopback allowed,
// DNS check skipped). Pin it to `public-only`, mirroring the AI Horde result download.
for (const privateUrl of ["http://127.0.0.1:1/x.png", "http://192.168.1.50/x.png"]) {
  test(`handleImageGeneration(nanobanana): b64_json never downloads a private result URL (${privateUrl}) (GHSA-34rg-3pqj-35g9)`, async () => {
    const originalFetch = globalThis.fetch;
    const fetchedUrls: string[] = [];

    globalThis.fetch = async (url) => {
      const u = String(url);
      fetchedUrls.push(u);

      if (u.includes("/generate")) {
        return new Response(
          JSON.stringify({ code: 200, msg: "success", data: { taskId: "task-ssrf-1" } }),
          { status: 200, headers: { "content-type": "application/json" } }
        );
      }

      if (u.includes("/record-info")) {
        return new Response(
          JSON.stringify({
            code: 200,
            msg: "success",
            data: { successFlag: 1, response: { resultImageUrl: privateUrl } },
          }),
          { status: 200, headers: { "content-type": "application/json" } }
        );
      }

      if (u === privateUrl) {
        // Canary: on the vulnerable code these bytes come back to the caller as b64_json.
        return new Response(new Uint8Array([0x89, 0x50, 0x4e, 0x47]), { status: 200 });
      }

      throw new Error(`Unexpected URL: ${u}`);
    };

    try {
      const result = await handleImageGeneration({
        body: {
          model: "nanobanana/nanobanana-flash",
          prompt: "galaxy test",
          response_format: "b64_json",
        },
        credentials: { apiKey: "test-key" },
        log: null,
      });

      assert.equal(result.success, false);
      assert.match(String(result.error), /blocked/i);
      assert.ok(
        !fetchedUrls.includes(privateUrl),
        `the private result URL must never be fetched (fetched: ${fetchedUrls.join(", ")})`
      );
    } finally {
      globalThis.fetch = originalFetch;
    }
  });
}