//! AAAA suppression for guests on hosts without an IPv6 route.
//!
//! When the host can't route v6 (`supermachine.host_ipv6=0`), a guest
//! that resolves a dual-stack hostname gets AAAA records it can never
//! connect to: the TSI connect fast-fails host-side. Clients with
//! happy-eyeballs fall back to v4; clients without it — alpine `apk`,
//! busybox `wget` — try the first (v6) address and give up. The
//! guest-side mitigation (init-oci writes `/etc/gai.conf` to bias
//! getaddrinfo toward A records) only works on glibc; musl ignores
//! gai.conf entirely, so Alpine images — the exact clients the fix
//! targets — were left broken.
//!
//! This module fixes it libc-agnostically at the single choke point
//! every resolver shares: the muxer's host-side UDP relay. For flows
//! to port 53, when the host has no v6 route, DNS *responses* whose
//! question is AAAA are rewritten to an empty NOERROR ("NODATA") —
//! the wire-true statement "this name has no AAAA records". Every
//! stub resolver (musl, glibc, c-ares, Go, Node) treats NODATA by
//! using the A results from the parallel/serial A query. Queries are
//! forwarded untouched; A/MX/TXT/etc. responses pass through
//! unmodified; hosts WITH v6 routes never engage the filter (full
//! AAAA + AF_TSI6 egress).
//!
//! The rewrite keeps header + question and zeroes ANCOUNT / NSCOUNT /
//! ARCOUNT (clearing TC). Surgically deleting only the AAAA RRs would
//! require rewriting DNS compression pointers in every later record —
//! all risk, no benefit, since the response exists only to say "use
//! v4".

/// QTYPE for AAAA records (RFC 3596).
const QTYPE_AAAA: u16 = 28;

/// If `msg` is a DNS response whose single question asks for AAAA
/// and which carries any answer/authority/additional records,
/// return the NODATA rewrite. `None` means "forward unmodified" —
/// non-DNS payloads, queries, non-AAAA questions, multi-question
/// packets (never produced by real stubs), already-empty responses,
/// and anything malformed all pass through untouched: this filter
/// must never turn working traffic into broken traffic.
pub fn rewrite_aaaa_response_to_nodata(msg: &[u8]) -> Option<Vec<u8>> {
    // Header: ID(2) FLAGS(2) QDCOUNT(2) ANCOUNT(2) NSCOUNT(2) ARCOUNT(2)
    if msg.len() < 12 {
        return None;
    }
    let flags = u16::from_be_bytes([msg[2], msg[3]]);
    if flags & 0x8000 == 0 {
        return None; // QR=0: a query, not a response.
    }
    let qdcount = u16::from_be_bytes([msg[4], msg[5]]);
    if qdcount != 1 {
        return None;
    }
    let ancount = u16::from_be_bytes([msg[6], msg[7]]);
    let nscount = u16::from_be_bytes([msg[8], msg[9]]);
    let arcount = u16::from_be_bytes([msg[10], msg[11]]);
    if ancount == 0 && nscount == 0 && arcount == 0 {
        return None; // Already NODATA-shaped; nothing to strip.
    }

    // Walk the QNAME labels. Compression pointers are illegal in
    // the question of a sanely-formed response (nothing earlier to
    // point at) — treat one as malformed and pass through.
    let mut off = 12usize;
    loop {
        let len = *msg.get(off)? as usize;
        if len == 0 {
            off += 1;
            break;
        }
        if len & 0xC0 != 0 {
            return None;
        }
        off += 1 + len;
        if off >= msg.len() {
            return None;
        }
    }
    // QTYPE + QCLASS
    if off + 4 > msg.len() {
        return None;
    }
    let qtype = u16::from_be_bytes([msg[off], msg[off + 1]]);
    if qtype != QTYPE_AAAA {
        return None;
    }
    let question_end = off + 4;

    let mut out = msg[..question_end].to_vec();
    out[6] = 0; // ANCOUNT
    out[7] = 0;
    out[8] = 0; // NSCOUNT
    out[9] = 0;
    out[10] = 0; // ARCOUNT (drops EDNS OPT — fine in a synthesized reply)
    out[11] = 0;
    // Clear TC so no resolver retries the truncated-looking reply
    // over TCP (which would bypass this filter).
    out[2] &= !0x02;
    Some(out)
}

/// DNS message builders shared between this module's tests and the
/// muxer's UDP-relay wiring tests.
#[cfg(test)]
pub(crate) mod testutil {
    use super::QTYPE_AAAA;

    /// Build a DNS message: header + one question for `example.com`,
    /// then `extra` (answer-section bytes etc., counted by an/ns/ar).
    pub(crate) fn dns_msg(
        qr_response: bool,
        qtype: u16,
        an: u16,
        ns: u16,
        ar: u16,
        extra: &[u8],
    ) -> Vec<u8> {
        let mut m = Vec::new();
        m.extend_from_slice(&0x1234u16.to_be_bytes()); // ID
        let flags: u16 = if qr_response { 0x8180 } else { 0x0100 };
        m.extend_from_slice(&flags.to_be_bytes());
        m.extend_from_slice(&1u16.to_be_bytes()); // QDCOUNT
        m.extend_from_slice(&an.to_be_bytes());
        m.extend_from_slice(&ns.to_be_bytes());
        m.extend_from_slice(&ar.to_be_bytes());
        for label in ["example", "com"] {
            m.push(label.len() as u8);
            m.extend_from_slice(label.as_bytes());
        }
        m.push(0);
        m.extend_from_slice(&qtype.to_be_bytes());
        m.extend_from_slice(&1u16.to_be_bytes()); // QCLASS=IN
        m.extend_from_slice(extra);
        m
    }

    /// One compressed answer RR (NAME = pointer to offset 12).
    pub(crate) fn aaaa_rr() -> Vec<u8> {
        let mut rr = vec![0xC0, 0x0C]; // pointer to question name
        rr.extend_from_slice(&QTYPE_AAAA.to_be_bytes());
        rr.extend_from_slice(&1u16.to_be_bytes()); // CLASS=IN
        rr.extend_from_slice(&60u32.to_be_bytes()); // TTL
        rr.extend_from_slice(&16u16.to_be_bytes()); // RDLENGTH
        rr.extend_from_slice(&[0u8; 16]); // ::
        rr
    }
}

#[cfg(test)]
mod tests {
    use super::testutil::{aaaa_rr, dns_msg};
    use super::*;

    #[test]
    fn aaaa_response_becomes_nodata() {
        let rr = aaaa_rr();
        let msg = dns_msg(true, QTYPE_AAAA, 1, 0, 0, &rr);
        let out = rewrite_aaaa_response_to_nodata(&msg).expect("must rewrite");
        // Header ID + question preserved, counts zeroed, RRs gone.
        assert_eq!(&out[..2], &msg[..2]);
        assert_eq!(u16::from_be_bytes([out[4], out[5]]), 1); // QDCOUNT intact
        assert_eq!(&out[6..12], &[0u8; 6]); // AN/NS/AR zeroed
        assert_eq!(out.len(), msg.len() - rr.len());
        assert_eq!(&out[12..], &msg[12..out.len()]); // question bytes intact
                                                     // RCODE preserved (NOERROR), QR still set.
        assert_eq!(out[2] & 0x80, 0x80);
        assert_eq!(out[3] & 0x0F, 0);
    }

    #[test]
    fn tc_bit_cleared_in_rewrite() {
        let mut msg = dns_msg(true, QTYPE_AAAA, 1, 0, 0, &aaaa_rr());
        msg[2] |= 0x02; // set TC
        let out = rewrite_aaaa_response_to_nodata(&msg).expect("must rewrite");
        assert_eq!(
            out[2] & 0x02,
            0,
            "TC must be cleared or stubs retry over TCP"
        );
    }

    #[test]
    fn cname_chain_in_aaaa_response_also_stripped() {
        // AAAA answers often arrive as CNAME + AAAA; both must go.
        let mut extra = vec![0xC0, 0x0C, 0, 5, 0, 1, 0, 0, 0, 60, 0, 2, 1, b'x'];
        extra.extend_from_slice(&aaaa_rr());
        let msg = dns_msg(true, QTYPE_AAAA, 2, 0, 0, &extra);
        let out = rewrite_aaaa_response_to_nodata(&msg).expect("must rewrite");
        assert_eq!(&out[6..8], &[0, 0]);
    }

    #[test]
    fn a_response_passes_through() {
        let msg = dns_msg(true, 1, 1, 0, 1, &[0u8; 20]);
        assert_eq!(rewrite_aaaa_response_to_nodata(&msg), None);
    }

    #[test]
    fn aaaa_query_passes_through() {
        let msg = dns_msg(false, QTYPE_AAAA, 0, 0, 0, &[]);
        assert_eq!(rewrite_aaaa_response_to_nodata(&msg), None);
    }

    #[test]
    fn empty_aaaa_response_passes_through() {
        // Upstream already said NODATA (possibly with SOA in NS —
        // but here fully empty): nothing to do.
        let msg = dns_msg(true, QTYPE_AAAA, 0, 0, 0, &[]);
        assert_eq!(rewrite_aaaa_response_to_nodata(&msg), None);
    }

    #[test]
    fn malformed_and_non_dns_pass_through() {
        assert_eq!(rewrite_aaaa_response_to_nodata(b""), None);
        assert_eq!(rewrite_aaaa_response_to_nodata(b"hi"), None);
        assert_eq!(rewrite_aaaa_response_to_nodata(&[0u8; 12]), None);
        // Truncated mid-QNAME: label length runs past the buffer.
        let mut msg = dns_msg(true, QTYPE_AAAA, 1, 0, 0, &aaaa_rr());
        msg.truncate(14);
        assert_eq!(rewrite_aaaa_response_to_nodata(&msg), None);
        // Compression pointer in the question: refuse to parse.
        let mut ptr_q = dns_msg(true, QTYPE_AAAA, 1, 0, 0, &[]);
        ptr_q[12] = 0xC0;
        assert_eq!(rewrite_aaaa_response_to_nodata(&ptr_q), None);
    }

    #[test]
    fn multi_question_passes_through() {
        let mut msg = dns_msg(true, QTYPE_AAAA, 1, 0, 0, &aaaa_rr());
        msg[5] = 2; // QDCOUNT=2
        assert_eq!(rewrite_aaaa_response_to_nodata(&msg), None);
    }

    #[test]
    fn servfail_aaaa_response_still_rewritten_preserving_rcode() {
        let mut msg = dns_msg(true, QTYPE_AAAA, 0, 1, 0, &aaaa_rr());
        msg[3] = (msg[3] & 0xF0) | 0x02; // RCODE=SERVFAIL
        let out = rewrite_aaaa_response_to_nodata(&msg).expect("ns-only payload stripped");
        assert_eq!(out[3] & 0x0F, 0x02, "RCODE must be preserved");
    }

    use proptest::prelude::*;

    proptest! {
        #![proptest_config(ProptestConfig::with_cases(4096))]

        /// The filter sits on the hostile side of the relay — any UDP
        /// payload a remote "DNS server" sends lands here. It must never
        /// panic, and whenever it DOES rewrite, the output must be a
        /// well-formed NODATA: header + question prefix of the input
        /// with counts zeroed and TC clear. (The AAAA bug itself was a
        /// "mechanism silently doesn't apply" failure — this guards the
        /// inverse class: the mechanism applying where it must not, or
        /// emitting garbage that breaks a resolver that WAS working.)
        #[test]
        fn arbitrary_bytes_never_panic_and_rewrites_are_wellformed(
            bytes in proptest::collection::vec(any::<u8>(), 0..600)
        ) {
            if let Some(out) = rewrite_aaaa_response_to_nodata(&bytes) {
                // Rewrites only ever shrink-or-keep and stay DNS-shaped.
                prop_assert!(out.len() >= 12 + 1 + 4); // header + root qname + qtype/qclass
                prop_assert!(out.len() <= bytes.len());
                prop_assert_eq!(&out[..2], &bytes[..2], "ID preserved");
                prop_assert_eq!(out[2] & 0x80, 0x80, "QR stays response");
                prop_assert_eq!(out[2] & 0x02, 0, "TC cleared");
                prop_assert_eq!(&out[4..6], &bytes[4..6], "QDCOUNT preserved");
                prop_assert_eq!(&out[6..12], &[0u8; 6][..], "AN/NS/AR zeroed");
                prop_assert_eq!(&out[12..], &bytes[12..out.len()], "question bytes verbatim");
                // And the input really was an AAAA response with payload —
                // the only thing allowed to trigger a rewrite.
                prop_assert_eq!(bytes[2] & 0x80, 0x80);
                let qt = u16::from_be_bytes([bytes[out.len() - 4], bytes[out.len() - 3]]);
                prop_assert_eq!(qt, QTYPE_AAAA);
            }
        }

        /// Valid responses for any non-AAAA qtype must pass through
        /// untouched no matter what the sections contain.
        #[test]
        fn non_aaaa_responses_never_rewritten(
            qtype in (0u16..=u16::MAX).prop_filter("non-AAAA", |q| *q != QTYPE_AAAA),
            an in 0u16..4,
            ns in 0u16..4,
            ar in 0u16..4,
            extra in proptest::collection::vec(any::<u8>(), 0..128),
        ) {
            let msg = dns_msg(true, qtype, an, ns, ar, &extra);
            prop_assert_eq!(rewrite_aaaa_response_to_nodata(&msg), None);
        }

        /// Truncating a rewritable AAAA response at ANY byte boundary
        /// must yield either a clean pass-through or a well-formed
        /// rewrite — never a panic, never an out-of-bounds slice.
        #[test]
        fn truncation_at_any_point_is_safe(cut in 0usize..64) {
            let full = dns_msg(true, QTYPE_AAAA, 1, 0, 0, &aaaa_rr());
            let cut = cut.min(full.len());
            let _ = rewrite_aaaa_response_to_nodata(&full[..cut]);
        }
    }
}
