metona-ai-desktop/electron/utils/__tests__/pinned-proxy.test.ts

/**
 * v0.8.0 P1-4/P1-5: Pinned CONNECT 代理 + 配置 URL 深校验测试。
 *
 * - pinned-proxy：公网/私网目标的 CONNECT 校验分流（resolver 注入 mock，
 *   不触网）；端口分配与幂等启动。
 * - assertSafeConfigTargetDeep：域名解析到云元数据 IP → 拒绝；DNS 失败 →
 *   DeepCheckSoftFailure（软失败放行）；IP 字面量高危段 → 拒绝。
 */

import { describe, it, expect, afterAll } from 'vitest';
import * as net from 'net';
import {
  createPinnedProxy,
  ensurePinnedProxy,
  isPinnedProxyActive,
} from '../../utils/pinned-proxy';
import {
  assertSafeConfigTargetDeep,
  DeepCheckSoftFailure,
  __dnsLookup,
} from '../../harness/tools/built-in/ssrf-guard';

// ===== Pinned CONNECT 代理 =====

describe('pinned CONNECT proxy (P1-4)', () => {
  const keepAlive: Array<net.Server> = [];

  it('public target → CONNECT tunnel established (200), traffic pinned to resolved IP', async () => {
    // 本地 echo 服务器充当"上游"：mock resolver 把域名 pin 到 127.0.0.1，
    // 验证隧道建立后字节双向透传（不触公网）
    const echo = net.createServer((socket) => {
      socket.pipe(socket);
    });
    await new Promise<void>((r) => echo.listen(0, '127.0.0.1', r));
    const echoPort = (echo.address() as net.AddressInfo).port;
    keepAlive.push(echo);

    const proxy = await createPinnedProxy(async () => {
      // 校验通过 → 返回校验期锁定的 IP（pinning 语义：连接用同一批 IP）
      return ['127.0.0.1'];
    });
    keepAlive.push(proxy as unknown as net.Server);

    const target = net.connect({ host: '127.0.0.1', port: proxy.port });
    const result = await new Promise<string>((resolve, reject) => {
      const timer = setTimeout(() => reject(new Error('timeout')), 3_000);
      let buf = '';
      target.on('data', (d) => {
        buf += d.toString();
        if (buf.includes('\r\n\r\n')) {
          // 200 已到 → 隧道建立，写 payload 验证透传
          target.write(`ping-to-${echoPort}`);
        }
        if (buf.includes(`ping-to-${echoPort}`)) {
          clearTimeout(timer);
          resolve(buf.split('\r\n')[0]);
        }
      });
      target.on('error', (e) => {
        clearTimeout(timer);
        reject(e);
      });
      target.write(`CONNECT example.com:${echoPort} HTTP/1.1\r\nHost: example.com\r\n\r\n`);
    });
    expect(result).toContain('200');
    target.destroy();
    echo.close();
    await proxy.close();
  });

  it('private/metadata target → 403 rejected (fail-closed)', async () => {
    const proxy = await createPinnedProxy(async () => {
      throw new Error('Blocked SSRF: resolves to private IP');
    });
    const target = net.connect({ host: '127.0.0.1', port: proxy.port });
    const response = await new Promise<string>((resolve, reject) => {
      const timer = setTimeout(() => reject(new Error('timeout')), 2_000);
      let buf = '';
      target.on('data', (d) => {
        buf += d.toString();
        if (buf.includes('\r\n\r\n')) {
          clearTimeout(timer);
          resolve(buf.split('\r\n')[0]);
        }
      });
      target.on('error', (e) => {
        clearTimeout(timer);
        reject(e);
      });
      target.write('CONNECT 169.254.169.254:80 HTTP/1.1\r\nHost: 169.254.169.254\r\n\r\n');
    });
    expect(response).toContain('403');
    target.destroy();
    await proxy.close();
  });

  it('invalid CONNECT target → 400', async () => {
    const proxy = await createPinnedProxy(async () => ['1.2.3.4']);
    const target = net.connect({ host: '127.0.0.1', port: proxy.port });
    const response = await new Promise<string>((resolve, reject) => {
      const timer = setTimeout(() => reject(new Error('timeout')), 2_000);
      let buf = '';
      target.on('data', (d) => {
        buf += d.toString();
        if (buf.includes('\r\n\r\n')) {
          clearTimeout(timer);
          resolve(buf.split('\r\n')[0]);
        }
      });
      target.on('error', (e) => {
        clearTimeout(timer);
        reject(e);
      });
      target.write('CONNECT no-port-target HTTP/1.1\r\n\r\n');
    });
    expect(response).toContain('400');
    target.destroy();
    await proxy.close();
  });

  it('ensurePinnedProxy: idempotent start + active flag', async () => {
    const port1 = await ensurePinnedProxy();
    const port2 = await ensurePinnedProxy();
    expect(port1).not.toBeNull();
    expect(port2).toBe(port1);
    expect(isPinnedProxyActive()).toBe(true);
  }, 10_000);

  afterAll(() => {
    for (const s of keepAlive) {
      try {
        s.close();
      } catch {
        /* ignore */
      }
    }
  });
});

// ===== 配置 URL 深校验 =====

describe('assertSafeConfigTargetDeep (P1-5)', () => {
  it('hostname resolving to metadata IP → rejected', async () => {
    const original = __dnsLookup.current;
    __dnsLookup.current = (async () => [
      { address: '169.254.169.254', family: 4 },
    ]) as unknown as typeof __dnsLookup.current;
    try {
      await expect(assertSafeConfigTargetDeep('https://attacker.example/search')).rejects.toThrow(
        /high-risk/,
      );
    } finally {
      __dnsLookup.current = original;
    }
  });

  it('hostname resolving to public IP → allowed', async () => {
    const original = __dnsLookup.current;
    __dnsLookup.current = (async () => [
      { address: '93.184.216.34', family: 4 },
    ]) as unknown as typeof __dnsLookup.current;
    try {
      await expect(assertSafeConfigTargetDeep('https://example.com')).resolves.toBeUndefined();
    } finally {
      __dnsLookup.current = original;
    }
  });

  it('DNS failure → DeepCheckSoftFailure (soft-fail allow, caller logs)', async () => {
    const original = __dnsLookup.current;
    __dnsLookup.current = (async () => {
      throw new Error('ENOTFOUND');
    }) as unknown as typeof __dnsLookup.current;
    try {
      await expect(assertSafeConfigTargetDeep('https://offline.example')).rejects.toBeInstanceOf(
        DeepCheckSoftFailure,
      );
    } finally {
      __dnsLookup.current = original;
    }
  });

  it('IP-literal link-local target → rejected without DNS', async () => {
    await expect(assertSafeConfigTargetDeep('http://169.254.169.254/latest')).rejects.toThrow(
      /link-local|metadata/,
    );
  });

  it('loopback target → allowed (local SearXNG/MCP legit)', async () => {
    await expect(
      assertSafeConfigTargetDeep('http://127.0.0.1:8080/search'),
    ).resolves.toBeUndefined();
  });
});