shadowvpn 0.6.0

//! Minimal DNS wire parsing — just enough for policy routing and Magic DNS.
//!
//! The proxy needs the queried name (to decide which upstream to use), the
//! IPv4 addresses in an answer (to add to the ipset), and — for Magic DNS —
//! the ability to synthesize a short `A`/`AAAA` or NXDOMAIN reply. Queries
//! and upstream responses are otherwise relayed verbatim.
//!
//! Helpers are total: any malformed input yields `None` / an empty vector
//! rather than panicking.

use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};

/// Fixed DNS header length in bytes.
const HEADER_LEN: usize = 12;
/// RR TYPE for an IPv4 address record (`A`).
pub const TYPE_A: u16 = 1;
/// RR TYPE for an IPv6 address record (`AAAA`).
pub const TYPE_AAAA: u16 = 28;
/// RR CLASS for the Internet (`IN`).
pub const CLASS_IN: u16 = 1;

/// Build a standard recursive `A`/`IN` query for `name` with transaction id
/// `id`. Used to pre-warm the cache; labels longer than 63 bytes are skipped.
pub fn build_query(id: u16, name: &str) -> Vec<u8> {
    let mut m = Vec::with_capacity(name.len() + 18);
    m.extend_from_slice(&id.to_be_bytes());
    m.extend_from_slice(&[0x01, 0x00]); // flags: RD (recursion desired)
    m.extend_from_slice(&[0x00, 0x01]); // QDCOUNT = 1
    m.extend_from_slice(&[0, 0, 0, 0, 0, 0]); // AN/NS/AR = 0
    for label in name.split('.') {
        if label.is_empty() || label.len() > 63 {
            continue;
        }
        m.push(label.len() as u8);
        m.extend_from_slice(label.as_bytes());
    }
    m.push(0); // root label
    m.extend_from_slice(&TYPE_A.to_be_bytes());
    m.extend_from_slice(&CLASS_IN.to_be_bytes());
    m
}

/// Extract the (lower-cased, dot-joined) name from the first question of a DNS
/// message, or `None` if there is no question or the message is malformed.
///
/// Question names do not use compression, so this reads plain labels.
pub fn question_name(msg: &[u8]) -> Option<String> {
    if msg.len() < HEADER_LEN {
        return None;
    }
    let qdcount = u16::from_be_bytes([msg[4], msg[5]]);
    if qdcount == 0 {
        return None;
    }
    let (name, _) = read_name(msg, HEADER_LEN)?;
    Some(name)
}

/// Extract the first question as `(name, qtype, qclass)` — the natural cache
/// key for a query. Returns `None` if there is no question or it is malformed.
pub fn question(msg: &[u8]) -> Option<(String, u16, u16)> {
    if msg.len() < HEADER_LEN {
        return None;
    }
    let qdcount = u16::from_be_bytes([msg[4], msg[5]]);
    if qdcount == 0 {
        return None;
    }
    let (name, pos) = read_name(msg, HEADER_LEN)?;
    if pos + 4 > msg.len() {
        return None;
    }
    let qtype = u16::from_be_bytes([msg[pos], msg[pos + 1]]);
    let qclass = u16::from_be_bytes([msg[pos + 2], msg[pos + 3]]);
    Some((name, qtype, qclass))
}

/// The smallest TTL across the answer section, or `None` if there are no
/// answer records (used to bound how long a response may be cached).
pub fn min_ttl(msg: &[u8]) -> Option<u32> {
    if msg.len() < HEADER_LEN {
        return None;
    }
    let qdcount = u16::from_be_bytes([msg[4], msg[5]]);
    let ancount = u16::from_be_bytes([msg[6], msg[7]]);

    let mut pos = HEADER_LEN;
    for _ in 0..qdcount {
        pos = skip_name(msg, pos)?;
        pos += 4;
        if pos > msg.len() {
            return None;
        }
    }

    let mut min: Option<u32> = None;
    for _ in 0..ancount {
        pos = skip_name(msg, pos)?;
        if pos + 10 > msg.len() {
            break;
        }
        let ttl = u32::from_be_bytes([msg[pos + 4], msg[pos + 5], msg[pos + 6], msg[pos + 7]]);
        let rdlength = u16::from_be_bytes([msg[pos + 8], msg[pos + 9]]) as usize;
        pos += 10 + rdlength;
        min = Some(min.map_or(ttl, |m| m.min(ttl)));
    }
    min
}

/// Extract every IPv4 address from the answer section of a DNS response.
///
/// Returns an empty vector for a query, an answer with no `A` records, or any
/// malformed message.
pub fn a_records(msg: &[u8]) -> Vec<Ipv4Addr> {
    let mut out = Vec::new();
    if msg.len() < HEADER_LEN {
        return out;
    }
    let qdcount = u16::from_be_bytes([msg[4], msg[5]]);
    let ancount = u16::from_be_bytes([msg[6], msg[7]]);

    let mut pos = HEADER_LEN;
    // Skip each question: QNAME + QTYPE(2) + QCLASS(2).
    for _ in 0..qdcount {
        pos = match skip_name(msg, pos) {
            Some(p) => p,
            None => return out,
        };
        pos += 4;
        if pos > msg.len() {
            return out;
        }
    }

    // Walk the answer RRs.
    for _ in 0..ancount {
        pos = match skip_name(msg, pos) {
            Some(p) => p,
            None => return out,
        };
        // TYPE(2) CLASS(2) TTL(4) RDLENGTH(2) = 10 bytes of fixed fields.
        if pos + 10 > msg.len() {
            return out;
        }
        let rtype = u16::from_be_bytes([msg[pos], msg[pos + 1]]);
        let rclass = u16::from_be_bytes([msg[pos + 2], msg[pos + 3]]);
        let rdlength = u16::from_be_bytes([msg[pos + 8], msg[pos + 9]]) as usize;
        pos += 10;
        if pos + rdlength > msg.len() {
            return out;
        }
        if rtype == TYPE_A && rclass == CLASS_IN && rdlength == 4 {
            out.push(Ipv4Addr::new(
                msg[pos],
                msg[pos + 1],
                msg[pos + 2],
                msg[pos + 3],
            ));
        }
        pos += rdlength;
    }
    out
}

/// Extract every IPv6 address from the answer section of a DNS response.
pub fn aaaa_records(msg: &[u8]) -> Vec<Ipv6Addr> {
    let mut out = Vec::new();
    if msg.len() < HEADER_LEN {
        return out;
    }
    let qdcount = u16::from_be_bytes([msg[4], msg[5]]);
    let ancount = u16::from_be_bytes([msg[6], msg[7]]);

    let mut pos = HEADER_LEN;
    for _ in 0..qdcount {
        pos = match skip_name(msg, pos) {
            Some(p) => p,
            None => return out,
        };
        pos += 4;
        if pos > msg.len() {
            return out;
        }
    }
    for _ in 0..ancount {
        pos = match skip_name(msg, pos) {
            Some(p) => p,
            None => return out,
        };
        if pos + 10 > msg.len() {
            return out;
        }
        let rtype = u16::from_be_bytes([msg[pos], msg[pos + 1]]);
        let rclass = u16::from_be_bytes([msg[pos + 2], msg[pos + 3]]);
        let rdlength = u16::from_be_bytes([msg[pos + 8], msg[pos + 9]]) as usize;
        pos += 10;
        if pos + rdlength > msg.len() {
            return out;
        }
        if rtype == TYPE_AAAA && rclass == CLASS_IN && rdlength == 16 {
            if let Ok(octets) = <[u8; 16]>::try_from(&msg[pos..pos + 16]) {
                out.push(Ipv6Addr::from(octets));
            }
        }
        pos += rdlength;
    }
    out
}

/// Synthesize a response to `query` with `addrs` as answers (filtered by
/// qtype: `A` → IPv4, `AAAA` → IPv6). Other types get an empty NOERROR.
///
/// Copies the question, sets QR/RD/RA, and uses a name pointer at `0xC0 0x0C`.
pub fn build_response(query: &[u8], addrs: &[IpAddr], ttl: u32) -> Option<Vec<u8>> {
    let (_, qtype, qclass) = question(query)?;
    let qend = skip_name(query, HEADER_LEN)? + 4;
    if qend > query.len() {
        return None;
    }
    let mut m = query[..qend].to_vec();
    m[2] = 0x81; // QR + RD
    m[3] = 0x80; // RA, RCODE=0
    m[4..6].copy_from_slice(&1u16.to_be_bytes());
    m[8..12].copy_from_slice(&[0, 0, 0, 0]); // NSCOUNT / ARCOUNT
    if qclass != CLASS_IN {
        m[6..8].copy_from_slice(&0u16.to_be_bytes());
        return Some(m);
    }
    let records: Vec<IpAddr> = addrs
        .iter()
        .copied()
        .filter(|a| match qtype {
            TYPE_A => a.is_ipv4(),
            TYPE_AAAA => a.is_ipv6(),
            _ => false,
        })
        .collect();
    m[6..8].copy_from_slice(&(records.len() as u16).to_be_bytes());
    for addr in records {
        m.extend_from_slice(&[0xC0, 0x0C]);
        match addr {
            IpAddr::V4(ip) => {
                m.extend_from_slice(&TYPE_A.to_be_bytes());
                m.extend_from_slice(&CLASS_IN.to_be_bytes());
                m.extend_from_slice(&ttl.to_be_bytes());
                m.extend_from_slice(&4u16.to_be_bytes());
                m.extend_from_slice(&ip.octets());
            }
            IpAddr::V6(ip) => {
                m.extend_from_slice(&TYPE_AAAA.to_be_bytes());
                m.extend_from_slice(&CLASS_IN.to_be_bytes());
                m.extend_from_slice(&ttl.to_be_bytes());
                m.extend_from_slice(&16u16.to_be_bytes());
                m.extend_from_slice(&ip.octets());
            }
        }
    }
    Some(m)
}

/// Synthesize an NXDOMAIN reply that echoes `query`'s question.
pub fn build_nxdomain(query: &[u8]) -> Option<Vec<u8>> {
    let qend = skip_name(query, HEADER_LEN)? + 4;
    if qend > query.len() {
        return None;
    }
    let mut m = query[..qend].to_vec();
    m[2] = 0x81; // QR + RD
    m[3] = 0x83; // RA + NXDOMAIN
    m[4..6].copy_from_slice(&1u16.to_be_bytes());
    m[6..12].copy_from_slice(&[0; 6]);
    Some(m)
}

/// Read a (possibly compressed) name starting at `pos`, returning the dot-joined
/// lower-cased name and the offset just past the name *in the original stream*
/// (i.e. past the first pointer if one is encountered).
fn read_name(msg: &[u8], start: usize) -> Option<(String, usize)> {
    let mut labels: Vec<String> = Vec::new();
    let mut pos = start;
    let mut jumped = false;
    let mut after_ptr = start;
    let mut budget = msg.len(); // guard against pointer loops

    loop {
        if pos >= msg.len() || budget == 0 {
            return None;
        }
        budget -= 1;
        let len = msg[pos];
        match len & 0xC0 {
            0x00 => {
                if len == 0 {
                    pos += 1;
                    if !jumped {
                        after_ptr = pos;
                    }
                    break;
                }
                let l = len as usize;
                let s = pos + 1;
                let e = s + l;
                if e > msg.len() {
                    return None;
                }
                labels.push(String::from_utf8_lossy(&msg[s..e]).to_ascii_lowercase());
                pos = e;
            }
            0xC0 => {
                if pos + 1 >= msg.len() {
                    return None;
                }
                let ptr = (((len & 0x3F) as usize) << 8) | msg[pos + 1] as usize;
                if !jumped {
                    after_ptr = pos + 2;
                    jumped = true;
                }
                if ptr >= msg.len() {
                    return None;
                }
                pos = ptr;
            }
            _ => return None, // 0x40 / 0x80 are reserved
        }
    }
    Some((labels.join("."), after_ptr))
}

/// Skip over a (possibly compressed) name, returning the offset just past it.
fn skip_name(msg: &[u8], start: usize) -> Option<usize> {
    let mut pos = start;
    loop {
        if pos >= msg.len() {
            return None;
        }
        let len = msg[pos];
        match len & 0xC0 {
            0x00 => {
                if len == 0 {
                    return Some(pos + 1);
                }
                pos += 1 + len as usize;
            }
            0xC0 => {
                // A pointer is always the end of the name; it is 2 bytes wide.
                return if pos + 1 < msg.len() {
                    Some(pos + 2)
                } else {
                    None
                };
            }
            _ => return None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Build a DNS query for `name` (type A).
    fn query(name: &str) -> Vec<u8> {
        let mut m = vec![0x12, 0x34, 0x01, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0];
        for label in name.split('.') {
            m.push(label.len() as u8);
            m.extend_from_slice(label.as_bytes());
        }
        m.push(0);
        m.extend_from_slice(&TYPE_A.to_be_bytes());
        m.extend_from_slice(&CLASS_IN.to_be_bytes());
        m
    }

    #[test]
    fn reads_question_name() {
        assert_eq!(
            question_name(&query("www.google.com")).as_deref(),
            Some("www.google.com")
        );
        assert_eq!(
            question_name(&query("EXAMPLE.com")).as_deref(),
            Some("example.com")
        );
        assert_eq!(question_name(b"short"), None);
    }

    /// Build a response to `q` with the given (ip, ttl) A records.
    fn response_with(q: &[u8], records: &[([u8; 4], u32)]) -> Vec<u8> {
        let mut m = q.to_vec();
        m[2] = 0x81;
        m[3] = 0x80;
        m[6] = (records.len() >> 8) as u8;
        m[7] = records.len() as u8;
        for (ip, ttl) in records {
            m.extend_from_slice(&[0xC0, 0x0C]); // name pointer
            m.extend_from_slice(&TYPE_A.to_be_bytes());
            m.extend_from_slice(&CLASS_IN.to_be_bytes());
            m.extend_from_slice(&ttl.to_be_bytes());
            m.extend_from_slice(&4u16.to_be_bytes());
            m.extend_from_slice(ip);
        }
        m
    }

    #[test]
    fn build_query_round_trips() {
        let q = build_query(0xABCD, "www.example.com");
        assert_eq!(&q[0..2], &[0xAB, 0xCD]); // id
        let (name, qtype, qclass) = question(&q).unwrap();
        assert_eq!(name, "www.example.com");
        assert_eq!((qtype, qclass), (TYPE_A, CLASS_IN));
    }

    #[test]
    fn reads_question_tuple() {
        let (name, qtype, qclass) = question(&query("a.b.example.com")).unwrap();
        assert_eq!(name, "a.b.example.com");
        assert_eq!((qtype, qclass), (TYPE_A, CLASS_IN));
        assert!(question(b"short").is_none());
    }

    #[test]
    fn min_ttl_picks_smallest() {
        let m = response_with(
            &query("example.com"),
            &[([1, 2, 3, 4], 300), ([5, 6, 7, 8], 60)],
        );
        assert_eq!(min_ttl(&m), Some(60));
        assert_eq!(min_ttl(&query("example.com")), None); // a query has no answers
    }

    #[test]
    fn build_response_filters_by_qtype() {
        let q = query("laptop.svpn");
        let addrs = [
            IpAddr::V4(Ipv4Addr::new(10, 9, 0, 7)),
            IpAddr::V6("fd07:7::7".parse().unwrap()),
        ];
        let resp = build_response(&q, &addrs, 30).unwrap();
        assert_eq!(a_records(&resp), vec![Ipv4Addr::new(10, 9, 0, 7)]);
        assert!(aaaa_records(&resp).is_empty());
        assert_eq!(min_ttl(&resp), Some(30));

        // AAAA question.
        let mut q6 = q.clone();
        let n = q6.len();
        q6[n - 4..n - 2].copy_from_slice(&TYPE_AAAA.to_be_bytes());
        let resp6 = build_response(&q6, &addrs, 30).unwrap();
        assert!(a_records(&resp6).is_empty());
        assert_eq!(
            aaaa_records(&resp6),
            vec!["fd07:7::7".parse::<Ipv6Addr>().unwrap()]
        );
    }

    #[test]
    fn build_nxdomain_sets_rcode() {
        let q = query("nope.svpn");
        let r = build_nxdomain(&q).unwrap();
        assert_eq!(r[3] & 0x0f, 3);
        assert_eq!(question_name(&r).as_deref(), Some("nope.svpn"));
        assert!(a_records(&r).is_empty());
    }

    #[test]
    fn extracts_a_records_with_compression() {
        // Response: header (ancount=2), one question, two A answers that point
        // back to the question name via a compression pointer (0xC0 0x0C).
        let mut m = query("example.com");
        m[2] = 0x81; // QR=1, RD=1
        m[3] = 0x80; // RA=1
        m[6] = 0x00;
        m[7] = 0x02; // ANCOUNT = 2
        for ip in [[93, 184, 216, 34], [1, 2, 3, 4]] {
            m.extend_from_slice(&[0xC0, 0x0C]); // name pointer -> offset 12
            m.extend_from_slice(&TYPE_A.to_be_bytes());
            m.extend_from_slice(&CLASS_IN.to_be_bytes());
            m.extend_from_slice(&300u32.to_be_bytes()); // TTL
            m.extend_from_slice(&4u16.to_be_bytes()); // RDLENGTH
            m.extend_from_slice(&ip);
        }
        let ips = a_records(&m);
        assert_eq!(
            ips,
            vec![Ipv4Addr::new(93, 184, 216, 34), Ipv4Addr::new(1, 2, 3, 4)]
        );
    }

    #[test]
    fn ignores_non_a_and_malformed() {
        assert!(a_records(&query("example.com")).is_empty()); // query, no answers
        assert!(a_records(b"").is_empty());
        assert!(a_records(b"\x00\x00\x00\x00\xff\xff\xff\xff").is_empty()); // bogus counts
    }
}