import { describe, it, expect } from 'vitest';
import * as http from 'node:http';
import * as net from 'node:net';
import { resolveAndCheck, pinnedConnect, pinnedFetch, type LookupFn } from './ssrf-strict.js';

describe('net/ssrf-strict resolveAndCheck', () => {
  it('passes a public IPv4 literal with allowPrivate=false', async () => {
    const r = await resolveAndCheck({ host: '8.8.8.8', allowPrivate: false });
    expect(r).toEqual({ ok: true, ip: '8.8.8.8', family: 4 });
  });

  it('rejects a private IPv4 literal with allowPrivate=false', async () => {
    const r = await resolveAndCheck({ host: '10.0.0.5', allowPrivate: false });
    expect(r.ok).toBe(false);
    if (!r.ok) expect(r.reason).toMatch(/Private/);
  });

  it('passes a private IPv4 literal with allowPrivate=true', async () => {
    const r = await resolveAndCheck({ host: '10.0.0.5', allowPrivate: true });
    expect(r).toEqual({ ok: true, ip: '10.0.0.5', family: 4 });
  });

  it('rejects loopback IPv4 with allowPrivate=false', async () => {
    const r = await resolveAndCheck({ host: '127.0.0.1', allowPrivate: false });
    expect(r.ok).toBe(false);
  });

  it('rejects link-local IPv4 (IMDS range) with allowPrivate=false', async () => {
    const r = await resolveAndCheck({ host: '169.254.169.254', allowPrivate: false });
    expect(r.ok).toBe(false);
  });

  it('passes a public IPv6 literal', async () => {
    const r = await resolveAndCheck({ host: '2001:4860:4860::8888', allowPrivate: false });
    expect(r.ok).toBe(true);
    if (r.ok) expect(r.family).toBe(6);
  });

  it('rejects IPv6 loopback', async () => {
    const r = await resolveAndCheck({ host: '::1', allowPrivate: false });
    expect(r.ok).toBe(false);
  });

  it('rejects when DNS returns a private address (allowPrivate=false)', async () => {
    const lookup: LookupFn = async () => [{ address: '10.0.0.10', family: 4 }];
    const r = await resolveAndCheck({ host: 'evil.example.com', allowPrivate: false, lookup });
    expect(r.ok).toBe(false);
    if (!r.ok) expect(r.reason).toMatch(/Forbidden IP in DNS response/);
  });

  it('passes when DNS returns a private address with allowPrivate=true', async () => {
    const lookup: LookupFn = async () => [{ address: '10.0.0.10', family: 4 }];
    const r = await resolveAndCheck({ host: 'host.lan', allowPrivate: true, lookup });
    expect(r).toEqual({ ok: true, ip: '10.0.0.10', family: 4 });
  });

  it('rejects when DNS returns no addresses', async () => {
    const lookup: LookupFn = async () => [];
    const r = await resolveAndCheck({ host: 'nothing.example', allowPrivate: false, lookup });
    expect(r.ok).toBe(false);
    if (!r.ok) expect(r.reason).toMatch(/no addresses/);
  });

  it('rejects on DNS lookup error', async () => {
    const lookup: LookupFn = async () => {
      throw new Error('ENOTFOUND');
    };
    const r = await resolveAndCheck({ host: 'bad.example', allowPrivate: false, lookup });
    expect(r.ok).toBe(false);
    if (!r.ok) expect(r.reason).toMatch(/DNS lookup failed/);
  });

  it('rejects empty or invalid host', async () => {
    const r1 = await resolveAndCheck({ host: '', allowPrivate: true });
    expect(r1.ok).toBe(false);
    const r2 = await resolveAndCheck({ host: null as unknown as string, allowPrivate: true });
    expect(r2.ok).toBe(false);
  });

  it('returns the first resolved address when multiple are returned', async () => {
    const lookup: LookupFn = async () => [
      { address: '8.8.8.8', family: 4 },
      { address: '8.8.4.4', family: 4 },
    ];
    const r = await resolveAndCheck({ host: 'multi.example', allowPrivate: false, lookup });
    expect(r.ok).toBe(true);
    if (r.ok) expect(r.ip).toBe('8.8.8.8');
  });

  it('rejects if ANY DNS-returned address is private (mixed list)', async () => {
    // Defense-in-depth against DNS that returns a public + private pair.
    const lookup: LookupFn = async () => [
      { address: '8.8.8.8', family: 4 },
      { address: '10.0.0.1', family: 4 },
    ];
    const r = await resolveAndCheck({ host: 'mixed.example', allowPrivate: false, lookup });
    expect(r.ok).toBe(false);
  });
});

describe('net/ssrf-strict pinnedConnect', () => {
  it('rejects when ip is not a literal address', async () => {
    await expect(
      pinnedConnect({ ip: 'example.com', port: 22, timeoutMs: 100 }),
    ).rejects.toThrow(/not a literal address/);
  });

  it('connects to a local echo server using its literal IP', async () => {
    const server = net.createServer();
    await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', () => resolve()));
    const addr = server.address();
    if (typeof addr !== 'object' || !addr) throw new Error('no address');
    try {
      const sock = await pinnedConnect({ ip: '127.0.0.1', port: addr.port, timeoutMs: 1000 });
      expect(sock).toBeInstanceOf(net.Socket);
      sock.destroy();
    } finally {
      await new Promise<void>((resolve) => server.close(() => resolve()));
    }
  });

  it('rejects with connect_timeout when target is unreachable', async () => {
    // 192.0.2.0/24 is TEST-NET-1, guaranteed not routable.
    await expect(
      pinnedConnect({ ip: '192.0.2.1', port: 22, timeoutMs: 100 }),
    ).rejects.toThrow(/connect_timeout|ECONN|EHOSTUNREACH|ENETUNREACH/);
  });
});

describe('net/ssrf-strict pinnedFetch', () => {
  it('connects to the pinned IP and keeps the body readable after the Agent close is scheduled', async () => {
    const server = http.createServer((_req, res) => {
      res.writeHead(200, { 'content-type': 'text/plain' });
      res.end('pinned-ok');
    });
    await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', resolve));
    const addr = server.address() as net.AddressInfo;
    try {
      // The hostname is unresolvable; only the pinned lookup can reach the server.
      const res = await pinnedFetch(`http://pinned-fetch.invalid:${addr.port}/`, {
        pinnedIp: '127.0.0.1',
        family: 4,
      });
      expect(res.status).toBe(200);
      // pinnedFetch schedules agent.close() before returning; the graceful
      // close must not cut off an unconsumed body.
      expect(await res.text()).toBe('pinned-ok');
    } finally {
      await new Promise<void>((resolve) => server.close(() => resolve()));
    }
  });

  it('rejects (and tears the Agent down) when the connection fails', async () => {
    const server = http.createServer();
    await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', resolve));
    const addr = server.address() as net.AddressInfo;
    await new Promise<void>((resolve) => server.close(() => resolve()));
    // Port is now closed → ECONNREFUSED → undici fetch rejects.
    await expect(
      pinnedFetch(`http://pinned-fetch.invalid:${addr.port}/`, {
        pinnedIp: '127.0.0.1',
        family: 4,
      }),
    ).rejects.toThrow();
  });
});
