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

/**
 * v0.8.2 P0-1: 适配器图片下载 SSRF 安全通道测试
 *
 * 锁定单元：
 * - sniffImageMediaType 魔数嗅探
 * - 协议白名单（仅 http/https）
 * - IP 直连私网/云元数据拒绝（resolvePublicAddresses 单一事实来源）
 * - DNS 解析到私网拒绝（rebinding 形态）
 * - 成功路径：下载 → base64 + 类型钳制（png/jpeg/gif/webp 白名单）
 * - 非图片 content-type 拒绝
 * - 字节上限（maxBytes + content-length 双闸）
 * - 逐跳重定向复检（每一跳重新进入完整校验）
 */

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

vi.mock('electron-log', () => ({
  default: { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() },
}));

const proxyMock = vi.hoisted(() => ({ isProxyActive: vi.fn(() => false) }));
vi.mock('../../../../utils/network-proxy', () => ({
  isProxyActive: proxyMock.isProxyActive,
}));

// undici mock：ssrfPinnedFetch 的 pinned 路径需要 Agent + fetch
const undiciMock = vi.hoisted(() => {
  class FakeAgent {
    async close(): Promise<void> {}
  }
  const fetch = vi.fn();
  return { FakeAgent, fetch };
});
vi.mock('undici', () => ({
  Agent: undiciMock.FakeAgent,
  fetch: undiciMock.fetch,
}));

import { fetchImageAsBase64, sniffImageMediaType, __imageFetcher } from '../ssrf-image-fetch';
import type { ImageFetcher } from '../ssrf-image-fetch';

const PNG_MAGIC = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0x00]);

function imageResponse(bytes: Uint8Array, contentType = 'image/png'): Response {
  return new Response(bytes, { status: 200, headers: { 'content-type': contentType } });
}

beforeEach(() => {
  undiciMock.fetch.mockReset();
  proxyMock.isProxyActive.mockReturnValue(false);
});

afterEach(() => {
  vi.restoreAllMocks();
});

describe('sniffImageMediaType', () => {
  it('识别 PNG / JPEG / GIF / WEBP 魔数', () => {
    const png = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]);
    expect(sniffImageMediaType(png)).toBe('image/png');
    const jpeg = Buffer.from([0xff, 0xd8, 0xff, 0xe0]);
    expect(sniffImageMediaType(jpeg)).toBe('image/jpeg');
    const gif = Buffer.from('GIF89a');
    expect(sniffImageMediaType(gif)).toBe('image/gif');
    const webp = Buffer.concat([
      Buffer.from('RIFF'),
      Buffer.from([0x00, 0x00, 0x00, 0x00]),
      Buffer.from('WEBP'),
    ]);
    expect(sniffImageMediaType(webp)).toBe('image/webp');
  });

  it('非图片内容返回 null', () => {
    expect(sniffImageMediaType(Buffer.from('hello world, plain text'))).toBeNull();
    expect(sniffImageMediaType(Buffer.alloc(0))).toBeNull();
  });
});

describe('fetchImageAsBase64 — 校验拒绝矩阵', () => {
  it('拒绝非 http/https 协议', async () => {
    await expect(fetchImageAsBase64('file:///etc/passwd')).rejects.toThrow(/protocol not allowed/);
    await expect(fetchImageAsBase64('ftp://example.com/a.png')).rejects.toThrow(
      /protocol not allowed/,
    );
  });

  it('拒绝 IP 直连私网/回环/云元数据', async () => {
    await expect(fetchImageAsBase64('http://127.0.0.1/img.png')).rejects.toThrow(/Blocked SSRF/);
    await expect(fetchImageAsBase64('http://10.1.2.3/img.png')).rejects.toThrow(/Blocked SSRF/);
    await expect(fetchImageAsBase64('http://192.168.1.1/img.png')).rejects.toThrow(/Blocked SSRF/);
    await expect(fetchImageAsBase64('http://172.16.0.9/img.png')).rejects.toThrow(/Blocked SSRF/);
    // 可回读通道的核心威胁：云元数据
    await expect(fetchImageAsBase64('http://169.254.169.254/latest/meta-data')).rejects.toThrow(
      /Blocked SSRF/,
    );
  });

  it('拒绝域名解析到私网（DNS rebinding 形态）', async () => {
    await expect(fetchImageAsBase64('http://nx.invalid.example/img.png')).rejects.toThrow(
      /Blocked SSRF/,
    );
  });
});

describe('fetchImageAsBase64 — 成功与钳制路径', () => {
  it('公网图片下载 → base64 + mediaType（IP 直连，无需 DNS）', async () => {
    const original = __imageFetcher.current;
    __imageFetcher.current = (async () => imageResponse(PNG_MAGIC)) as unknown as ImageFetcher;
    try {
      const r = await fetchImageAsBase64('http://93.184.216.34/a.png');
      expect(r.mediaType).toBe('image/png');
      expect(Buffer.from(r.base64, 'base64').equals(PNG_MAGIC)).toBe(true);
    } finally {
      __imageFetcher.current = original;
    }
  });

  it('非图片 content-type（text/html）拒绝', async () => {
    const original = __imageFetcher.current;
    __imageFetcher.current = (async () =>
      imageResponse(Buffer.from('<html></html>'), 'text/html')) as unknown as ImageFetcher;
    try {
      await expect(fetchImageAsBase64('http://93.184.216.34/a.png')).rejects.toThrow(
        /non-image content-type/,
      );
    } finally {
      __imageFetcher.current = original;
    }
  });

  it('白名单外图片类型（image/bmp）拒绝', async () => {
    const original = __imageFetcher.current;
    __imageFetcher.current = (async () =>
      imageResponse(Buffer.from('BMxx'), 'image/bmp')) as unknown as ImageFetcher;
    try {
      await expect(fetchImageAsBase64('http://93.184.216.34/a.bmp')).rejects.toThrow(
        /unsupported media type/,
      );
    } finally {
      __imageFetcher.current = original;
    }
  });

  it('字节上限：超过 maxBytes 拒绝（防大图内存峰值）', async () => {
    const original = __imageFetcher.current;
    __imageFetcher.current = (async () =>
      imageResponse(Buffer.alloc(64, 1))) as unknown as ImageFetcher;
    try {
      await expect(
        fetchImageAsBase64('http://93.184.216.34/a.png', { maxBytes: 8 }),
      ).rejects.toThrow(/size limit/);
    } finally {
      __imageFetcher.current = original;
    }
  });

  it('逐跳重定向复检：302 跳转后以下一跳 URL 再次下载，终图可用', async () => {
    const original = __imageFetcher.current;
    const seen: string[] = [];
    const fetcher: ImageFetcher = async (url: string) => {
      seen.push(url);
      if (seen.length === 1) {
        return new Response(null, {
          status: 302,
          headers: { location: 'http://93.184.216.34/final.png' },
        });
      }
      return imageResponse(PNG_MAGIC);
    };
    __imageFetcher.current = fetcher;
    try {
      const r = await fetchImageAsBase64('http://93.184.216.34/redirect.png');
      expect(r.mediaType).toBe('image/png');
      expect(seen).toEqual(['http://93.184.216.34/redirect.png', 'http://93.184.216.34/final.png']);
    } finally {
      __imageFetcher.current = original;
    }
  });
});