metona-ai-desktop/electron/harness/adapters/shared/ssrf-image-fetch.ts

/**
 * v0.8.2 P0-1: 适配器图片下载的 SSRF 安全通道（Anthropic / Ollama 共用）
 *
 * 背景：此前 anthropic.adapter.toImageBlock 与 ollama.adapter.resolveImageToBase64
 * 对消息里的 http(s) 图片 URL 直接 fetchWithTimeout 下载转 base64 —— 未接 SSRF
 * 校验、无字节上限、无 content-type 约束。与其他工具"SSRF 拦截即断"不同，
 * 该通道的下载结果会**以图片块进入模型上下文（数据可回读）**：模型可诱导用户
 * 发送 http://169.254.169.254/... 图片链接回读内网数据，属于可回读外泄通道；
 * 同时无上限的 arrayBuffer 会造成内存峰值。
 *
 * 收口方案（根治）：
 * - 校验与连接同源：resolvePublicAddresses（ssrf-guard 单一事实来源）+ DNS
 *   pinning（ssrfPinnedFetch），代理激活时按既有语义退化为"仅入口校验"；
 * - 逐跳重定向复检：redirect:'manual' + resolveRedirectTarget，每一跳都重新
 *   走完整校验（对齐 web_fetch 的逐跳语义），最多 3 跳；
 * - 字节上限：content-length 预检 + 流式增量累计双闸（10MB，防大图内存峰值）；
 * - 类型白名单：content-type 必须 image/*（或缺失/二进制时按魔数嗅探），
 *   并钳制到 Provider 实际支持集合（png/jpeg/gif/webp）。
 *
 * 失败语义由调用方决定：适配器保持"跳过该图、不阻断请求"（log.warn + 占位）。
 */

import { ssrfPinnedFetch, resolveRedirectTarget } from '../../tools/built-in/ssrf-dispatcher';

/** 实际下载器签名（与 ssrfPinnedFetch 对齐） */
export type ImageFetcher = (
  url: string,
  init: RequestInit,
  timeoutMs: number,
  signal?: AbortSignal,
) => Promise<Response>;

/**
 * v0.8.2 P0-1: 下载器注入点 —— 生产恒为 ssrfPinnedFetch（校验与连接同源）；
 * 单元测试通过替换 current 注入受控下载器（undici 客户端无法经 global fetch 打桩）。
 */
export const __imageFetcher: { current: ImageFetcher } = { current: ssrfPinnedFetch };

/** 单张图片下载字节上限（10MB，对齐 view_image 工具 5MB×2 的量级） */
export const MAX_IMAGE_BYTES = 10 * 1024 * 1024;

/** 允许的图片 media type（Anthropic Messages API 支持集合；Ollama 同样接受） */
const ALLOWED_MEDIA_TYPES = new Set(['image/png', 'image/jpeg', 'image/gif', 'image/webp']);

const MAX_REDIRECTS = 3;

export interface ImageFetchOptions {
  timeoutMs?: number;
  maxBytes?: number;
  signal?: AbortSignal;
}

export interface ImageFetchResult {
  base64: string;
  mediaType: string;
}

/** 魔数嗅探：content-type 缺失或为通用二进制类型时判定真实图片类型 */
export function sniffImageMediaType(buf: Buffer): string | null {
  if (buf.length >= 8 && buf.subarray(0, 4).equals(Buffer.from([0x89, 0x50, 0x4e, 0x47]))) {
    return 'image/png';
  }
  if (buf.length >= 3 && buf[0] === 0xff && buf[1] === 0xd8 && buf[2] === 0xff) {
    return 'image/jpeg';
  }
  if (buf.length >= 6 && buf.subarray(0, 3).toString('ascii') === 'GIF') {
    return 'image/gif';
  }
  if (
    buf.length >= 12 &&
    buf.subarray(0, 4).toString('ascii') === 'RIFF' &&
    buf.subarray(8, 12).toString('ascii') === 'WEBP'
  ) {
    return 'image/webp';
  }
  return null;
}

/** 流式增量读取并强制字节上限（content-length 可伪造，以实际累计为准） */
async function readBodyWithCap(res: Response, maxBytes: number): Promise<Buffer> {
  const declared = res.headers.get('content-length');
  const declaredBytes = declared ? Number(declared) : NaN;
  if (Number.isFinite(declaredBytes) && declaredBytes > maxBytes) {
    throw new Error(`Image exceeds size limit: ${declaredBytes} > ${maxBytes} bytes`);
  }
  const body = res.body;
  if (!body) {
    throw new Error('Image response has no body');
  }
  const reader = body.getReader();
  const chunks: Buffer[] = [];
  let total = 0;
  try {
    for (;;) {
      const { done, value } = await reader.read();
      if (done) break;
      if (value) {
        total += value.byteLength;
        if (total > maxBytes) {
          throw new Error(`Image exceeds size limit: > ${maxBytes} bytes`);
        }
        chunks.push(Buffer.from(value));
      }
    }
  } finally {
    reader.releaseLock();
  }
  return Buffer.concat(chunks);
}

/**
 * SSRF 安全的图片下载：校验 → pinned 连接 → 逐跳重定向复检 → 类型/尺寸钳制。
 *
 * @throws 任何校验失败/网络失败/超限/类型不符均抛 Error（消息含 Blocked SSRF /
 *         Image exceeds / non-image 等），由调用方按"跳过该图"处理。
 */
export async function fetchImageAsBase64(
  url: string,
  options: ImageFetchOptions = {},
): Promise<ImageFetchResult> {
  const { timeoutMs = 30_000, maxBytes = MAX_IMAGE_BYTES, signal } = options;

  if (!/^https?:/i.test(url)) {
    throw new Error(`Blocked image fetch: protocol not allowed (url=${url.slice(0, 120)})`);
  }

  let currentUrl = url;
  let res: Response | null = null;
  for (let hop = 0; hop <= MAX_REDIRECTS; hop++) {
    // __imageFetcher（生产 = ssrfPinnedFetch）内部先 resolvePublicAddresses 全量
    // 校验再 pinned 连接；每一跳都进入本调用 —— 重定向目标同样受完整 SSRF 校验。
    res = await __imageFetcher.current(currentUrl, { redirect: 'manual' }, timeoutMs, signal);
    const next = resolveRedirectTarget(
      { status: res.status, headers: { get: (n: string) => res!.headers.get(n) } },
      currentUrl,
    );
    if (!next) break;
    if (hop === MAX_REDIRECTS) {
      throw new Error(`Image fetch exceeded ${MAX_REDIRECTS} redirects`);
    }
    currentUrl = next;
  }
  if (!res) throw new Error('Image fetch failed: no response');

  if (!res.ok) {
    throw new Error(`Image fetch failed: HTTP ${res.status} (url=${currentUrl.slice(0, 120)})`);
  }

  const buf = await readBodyWithCap(res, maxBytes);
  if (buf.length === 0) {
    throw new Error('Image fetch failed: empty body');
  }

  // 类型钳制：content-type 声明优先（必须 image/*），缺失/通用二进制按魔数嗅探；
  // 明确的非图片类型（text/html、application/json 等）直接拒绝。
  const declaredType = (res.headers.get('content-type') ?? '').split(';')[0].trim().toLowerCase();
  let mediaType: string | null = null;
  if (
    declaredType === '' ||
    declaredType === 'application/octet-stream' ||
    declaredType === 'binary/octet-stream'
  ) {
    mediaType = sniffImageMediaType(buf);
  } else if (declaredType.startsWith('image/')) {
    mediaType = declaredType;
  } else {
    throw new Error(`Blocked image fetch: non-image content-type "${declaredType}"`);
  }
  if (!mediaType || !ALLOWED_MEDIA_TYPES.has(mediaType)) {
    throw new Error(
      `Blocked image fetch: unsupported media type "${mediaType ?? declaredType}" (allowed: png/jpeg/gif/webp)`,
    );
  }

  return { base64: buf.toString('base64'), mediaType };
}