// Wire-format tests for the bnsd DNS server.
//
// The bug we chase: strict resolvers (Node getaddrinfo, Firefox trr-off) fail
// against the shipped daemon while lenient ones (nslookup, Chrome async-dns)
// work. The suspect is EDNS0 OPT handling and how the daemon frames NODATA.
//
// Each test constructs a real DNS query on the wire, drives the server, and
// parses the response. When possible we also drive Node's `dns.Resolver`
// against the running server — it is UDP-level like getaddrinfo, so if it
// rejects the response the shipped daemon has the same fault under real
// clients.

import { test } from "node:test";
import assert from "node:assert/strict";
import dgram from "node:dgram";
import { Resolver } from "node:dns/promises";
import {
  parseQuestion,
  extractOpt,
  buildReply,
  handleQuery,
  startDnsServer,
} from "../src/lib/dns-server.js";

// --- helpers to build/parse DNS packets on the wire ---------------------------

function encodeName(name) {
  const parts = name.split(".").filter(Boolean);
  let len = 1;
  for (const p of parts) len += 1 + p.length;
  const out = Buffer.alloc(len);
  let i = 0;
  for (const p of parts) {
    out[i++] = p.length;
    out.write(p, i, "ascii");
    i += p.length;
  }
  out[i] = 0;
  return out;
}

// Build a minimal DNS query. `edns` adds an OPT RR in the additional section
// with the given udpPayloadSize (default 1232 — what modern glibc/Firefox use).
function buildQuery(name, qtype, { id = 0x1234, rd = true, edns = false, udpPayloadSize = 1232 } = {}) {
  const header = Buffer.alloc(12);
  header.writeUInt16BE(id, 0);
  header.writeUInt16BE(rd ? 0x0100 : 0x0000, 2);
  header.writeUInt16BE(1, 4);                    // QDCOUNT
  header.writeUInt16BE(0, 6);                    // ANCOUNT
  header.writeUInt16BE(0, 8);                    // NSCOUNT
  header.writeUInt16BE(edns ? 1 : 0, 10);        // ARCOUNT
  const qname = encodeName(name);
  const qtc = Buffer.alloc(4);
  qtc.writeUInt16BE(qtype, 0);
  qtc.writeUInt16BE(1, 2);                       // QCLASS = IN
  const parts = [header, qname, qtc];
  if (edns) {
    // OPT RR: name=root(0), type=41, class=udpPayloadSize, TTL=0 (extRcode|ver|flags), rdlength=0
    const opt = Buffer.alloc(11);
    opt[0] = 0;                                  // root name
    opt.writeUInt16BE(41, 1);                    // TYPE = OPT
    opt.writeUInt16BE(udpPayloadSize, 3);
    opt.writeUInt32BE(0, 5);                     // TTL / extended rcode etc
    opt.writeUInt16BE(0, 9);                     // RDLENGTH
    parts.push(opt);
  }
  return Buffer.concat(parts);
}

function parseHeader(msg) {
  return {
    id: msg.readUInt16BE(0),
    flags: msg.readUInt16BE(2),
    qr: (msg[2] >> 7) & 1,
    aa: (msg[2] >> 2) & 1,
    tc: (msg[2] >> 1) & 1,
    rd: msg[2] & 1,
    ra: (msg[3] >> 7) & 1,
    rcode: msg[3] & 0x0f,
    qdcount: msg.readUInt16BE(4),
    ancount: msg.readUInt16BE(6),
    nscount: msg.readUInt16BE(8),
    arcount: msg.readUInt16BE(10),
  };
}

// Send one DNS query to `port` on 127.0.0.1 and resolve with the raw response.
function askServer(port, query, { timeoutMs = 2000 } = {}) {
  return new Promise((resolve, reject) => {
    const sock = dgram.createSocket("udp4");
    const timer = setTimeout(() => { try { sock.close(); } catch {} reject(new Error("client timeout")); }, timeoutMs);
    sock.on("message", (resp) => { clearTimeout(timer); try { sock.close(); } catch {} resolve(resp); });
    sock.on("error", (e) => { clearTimeout(timer); reject(e); });
    sock.send(query, port, "127.0.0.1");
  });
}

// Minimal DNS server that ALWAYS forwards its answer to `respFor(qname, qtype)`
// — used as a mock UPSTREAM. Returns { port, close }.
function startMockUpstream(respFor) {
  const sock = dgram.createSocket("udp4");
  sock.on("message", (msg, rinfo) => {
    const q = parseQuestion(msg);
    if (!q) return;
    // Echo the query back with an A record answer + preserve OPT if present.
    const opt = extractOpt(msg, q.qend);
    const answers = respFor(q.name, q.qtype);
    sock.send(buildReply(msg, q.qend, { answers, opt }), rinfo.port, rinfo.address);
  });
  return new Promise((resolve) => {
    sock.bind(0, "127.0.0.1", () => resolve({ port: sock.address().port, close: () => sock.close() }));
  });
}

async function startServer({ lookup, upstreamPort = 53 } = {}) {
  return new Promise((resolve) => {
    const sock = startDnsServer({
      port: 0,
      bind: "127.0.0.1",
      upstream: `127.0.0.1:${upstreamPort}`,
      lookup,
      onListening: (addr) => resolve({ port: addr.port, close: () => sock.close() }),
    });
  });
}

// --- parseQuestion / buildReply unit tests -----------------------------------

test("parseQuestion parses a plain A query", () => {
  const q = buildQuery("hello.bch", 1);
  const parsed = parseQuestion(q);
  assert.deepEqual(parsed && { name: parsed.name, qtype: parsed.qtype, qclass: parsed.qclass }, {
    name: "hello.bch", qtype: 1, qclass: 1,
  });
});

test("parseQuestion tolerates an EDNS OPT record following the question", () => {
  const q = buildQuery("hello.bch", 1, { edns: true });
  const parsed = parseQuestion(q);
  assert.equal(parsed && parsed.name, "hello.bch");
  const opt = extractOpt(q, parsed.qend);
  assert.ok(opt, "OPT record should be found in the additional section");
  assert.equal(opt.readUInt16BE(1), 41, "OPT TYPE = 41");
});

test("parseQuestion rejects a query with QDCOUNT != 1", () => {
  const q = buildQuery("hello.bch", 1);
  q.writeUInt16BE(0, 4);                         // QDCOUNT = 0
  assert.equal(parseQuestion(q), null);
});

test("parseQuestion rejects a truncated label", () => {
  const q = Buffer.from([
    0x12, 0x34, 0x01, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x05, 0x68, 0x65, // says label length 5 but only "he" follows
  ]);
  assert.equal(parseQuestion(q), null);
});

test("buildReply echoes the OPT record when the query had one", () => {
  const q = buildQuery("hello.bch", 1, { edns: true });
  const { qend } = parseQuestion(q);
  const resp = buildReply(q, qend, { answers: ["1.2.3.4"] });
  const h = parseHeader(resp);
  assert.equal(h.qr, 1);
  assert.equal(h.aa, 1);
  assert.equal(h.rcode, 0);
  assert.equal(h.qdcount, 1);
  assert.equal(h.ancount, 1);
  assert.equal(h.arcount, 1, "OPT record must be echoed — a strict client sees FORMERR otherwise");
});

test("buildReply produces NODATA (empty answers) with echoed OPT for AAAA on a known-A name", () => {
  const q = buildQuery("hello.bch", 28, { edns: true });        // AAAA
  const { qend } = parseQuestion(q);
  const resp = buildReply(q, qend, { answers: [] });
  const h = parseHeader(resp);
  assert.equal(h.rcode, 0, "NODATA is NOERROR, not NXDOMAIN");
  assert.equal(h.ancount, 0);
  assert.equal(h.arcount, 1, "OPT must be echoed on NODATA too");
});

test("buildReply drops OPT when the query had none", () => {
  const q = buildQuery("hello.bch", 1, { edns: false });
  const { qend } = parseQuestion(q);
  const resp = buildReply(q, qend, { answers: ["1.2.3.4"] });
  const h = parseHeader(resp);
  assert.equal(h.arcount, 0, "must not fabricate OPT — RFC 6891 §7");
});

test("buildReply preserves the query's RD flag", () => {
  const rdOn = buildQuery("hello.bch", 1, { rd: true });
  const rdOff = buildQuery("hello.bch", 1, { rd: false });
  const { qend: e1 } = parseQuestion(rdOn);
  const { qend: e2 } = parseQuestion(rdOff);
  assert.equal(parseHeader(buildReply(rdOn, e1, { answers: ["1.2.3.4"] })).rd, 1);
  assert.equal(parseHeader(buildReply(rdOff, e2, { answers: ["1.2.3.4"] })).rd, 0);
});

// --- integration: dns.Resolver against a live in-process server --------------

test("dns.Resolver.resolve4 on a BCNR-known name returns the on-chain IP", async (t) => {
  const upstream = await startMockUpstream(() => []);
  const server = await startServer({
    lookup: (name) => name === "hello.bch" ? { exists: true, ip: "203.0.113.10" } : null,
    upstreamPort: upstream.port,
  });
  t.after(() => { server.close(); upstream.close(); });

  const r = new Resolver();
  r.setServers([`127.0.0.1:${server.port}`]);
  const addrs = await r.resolve4("hello.bch");
  assert.deepEqual(addrs, ["203.0.113.10"]);
});

test("dns.Resolver.resolve6 on a BCNR-known name returns NODATA (no throw)", async (t) => {
  const upstream = await startMockUpstream(() => []);
  const server = await startServer({
    lookup: (name) => name === "hello.bch" ? { exists: true, ip: "203.0.113.10" } : null,
    upstreamPort: upstream.port,
  });
  t.after(() => { server.close(); upstream.close(); });

  const r = new Resolver();
  r.setServers([`127.0.0.1:${server.port}`]);
  // c-ares maps NODATA to ENODATA. Assert on that rather than a resolve() value.
  await assert.rejects(r.resolve6("hello.bch"), (e) => e.code === "ENODATA");
});

test("dns.Resolver.resolve4 on a non-BCNR name is forwarded to the mock upstream", async (t) => {
  const upstream = await startMockUpstream((name, qtype) => name === "example.com" && qtype === 1 ? ["198.51.100.5"] : []);
  const server = await startServer({
    lookup: () => null,                          // nothing on BCNR
    upstreamPort: upstream.port,
  });
  t.after(() => { server.close(); upstream.close(); });

  const r = new Resolver();
  r.setServers([`127.0.0.1:${server.port}`]);
  const addrs = await r.resolve4("example.com");
  assert.deepEqual(addrs, ["198.51.100.5"]);
});

// --- direct byte-level check: send EDNS query, confirm OPT survives ----------

test("EDNS0 OPT survives the round trip for a BCNR-known name", async (t) => {
  const upstream = await startMockUpstream(() => []);
  const server = await startServer({
    lookup: (name) => name === "hello.bch" ? { exists: true, ip: "1.2.3.4" } : null,
    upstreamPort: upstream.port,
  });
  t.after(() => { server.close(); upstream.close(); });

  const q = buildQuery("hello.bch", 1, { edns: true, udpPayloadSize: 4096 });
  const resp = await askServer(server.port, q);
  const h = parseHeader(resp);
  assert.equal(h.qr, 1);
  assert.equal(h.rcode, 0);
  assert.equal(h.ancount, 1);
  assert.equal(h.arcount, 1, "server dropped the OPT record — a strict resolver would fail");
});

// --- regression tests for the OLD inline implementation in bnsd.js -----------
//
// We copy the SHIPPED functions verbatim (as of git bf15d86) and prove they
// exhibit the failure the memory note describes. If these two tests pass, the
// bug is confirmed. If they fail (because someone already patched bnsd), the
// fix is already in place and the shipped daemon behaves like the new module.

const TTL = 60;
function oldParseQuestion(msg) {
  let i = 12;
  const labels = [];
  while (i < msg.length) {
    const len = msg[i];
    if (len === 0) { i += 1; break; }
    if (len > 63) return null;
    labels.push(msg.subarray(i + 1, i + 1 + len).toString("ascii"));
    i += 1 + len;
  }
  if (i + 4 > msg.length) return null;
  const qtype = msg.readUInt16BE(i);
  return { name: labels.join(".").toLowerCase(), qtype, qend: i + 4 };
}
function oldReply(query, qend, { rcode = 0, answers = [] }) {
  const question = query.subarray(12, qend);
  const header = Buffer.alloc(12);
  query.copy(header, 0, 0, 2);
  header.writeUInt16BE(0x8580 | rcode, 2);
  header.writeUInt16BE(1, 4);
  header.writeUInt16BE(answers.length, 6);
  const parts = [header, question];
  for (const ip of answers) {
    const rr = Buffer.alloc(16);
    rr.writeUInt16BE(0xc00c, 0);
    rr.writeUInt16BE(1, 2);
    rr.writeUInt16BE(1, 4);
    rr.writeUInt32BE(TTL, 6);
    rr.writeUInt16BE(4, 10);
    const o = ip.split(".").map(Number);
    for (let j = 0; j < 4; j++) rr[12 + j] = o[j];
    parts.push(rr);
  }
  return Buffer.concat(parts);
}

test("OLD reply() drops the OPT record when the query had one — this is the shipped bug", () => {
  const q = buildQuery("hello.bch", 1, { edns: true });
  const { qend } = oldParseQuestion(q);
  const resp = oldReply(q, qend, { answers: ["1.2.3.4"] });
  const h = parseHeader(resp);
  assert.equal(h.arcount, 0, "OLD reply drops the OPT — strict clients see FORMERR");
});

test("OLD reply() for NODATA (AAAA on known A) also drops OPT", () => {
  const q = buildQuery("hello.bch", 28, { edns: true });
  const { qend } = oldParseQuestion(q);
  const resp = oldReply(q, qend, { answers: [] });
  const h = parseHeader(resp);
  assert.equal(h.arcount, 0);
});