neiam /neiam-portal

//! Minimal captive-portal DNS server.
//!
//! Answers every A query with the SoftAP gateway IP so that any hostname a
//! client looks up resolves to this device, which triggers the OS
//! captive-portal browser.
//!
//! Only `std::net` is used, so this module builds and tests on the host.

use std::net::UdpSocket;

/// DNS `TYPE` for an IPv4 address record.
const TYPE_A: u16 = 1;
/// Fixed offset of the question section (immediately after the 12-byte header).
const QUESTION_OFFSET: usize = 12;

/// Serve DNS on `0.0.0.0:53`, replying to all A queries with `ip`. Never
/// returns; intended to own a thread.
pub fn serve(ip: [u8; 4]) {
    let socket = match UdpSocket::bind("0.0.0.0:53") {
        Ok(s) => s,
        Err(e) => {
            log::error!("captive DNS: bind :53 failed: {e}");
            return;
        }
    };
    log::info!(
        "captive DNS serving -> {}.{}.{}.{}",
        ip[0],
        ip[1],
        ip[2],
        ip[3]
    );

    let mut buf = [0u8; 512];
    loop {
        let (len, peer) = match socket.recv_from(&mut buf) {
            Ok(v) => v,
            Err(e) => {
                log::warn!("captive DNS recv error: {e}");
                continue;
            }
        };

        if let Some(reply) = build_reply(&buf[..len], ip) {
            let _ = socket.send_to(&reply, peer);
        }
    }
}

/// Build a reply pointing the queried name at `ip`.
///
/// Echoes the question section and, for A queries, appends a single answer
/// using a name pointer (`0xC00C`) back to the question.
///
/// Non-A queries (notably AAAA, which phones interleave with A) get a
/// well-formed `NOERROR` with **no** answer. This differs from the original
/// `nfc-522` implementation, which replied to every query type with a type-A
/// answer record — a malformed response that a strict resolver may discard.
/// Returning an empty NOERROR is what captive portals normally do and lets the
/// client fall through to its A lookup.
///
/// Returns `None` for anything too short to be a question, which the caller
/// simply drops.
pub fn build_reply(query: &[u8], ip: [u8; 4]) -> Option<Vec<u8>> {
    // Need at least a header (12 bytes) plus one question.
    if query.len() < QUESTION_OFFSET {
        return None;
    }
    // QDCOUNT == 0 would leave the 0xC00C answer pointer dangling.
    let qdcount = u16::from_be_bytes([query[4], query[5]]);
    if qdcount == 0 {
        return None;
    }

    let qtype = question_type(query)?;

    let mut resp = query.to_vec();
    // Flags: QR=1 (response), Opcode=0, AA=1, RD copied, RA=0; second byte RCODE=0.
    resp[2] = 0x81;
    resp[3] = 0x80;
    // Zero out NSCOUNT/ARCOUNT to avoid dangling sections.
    resp[8] = 0x00;
    resp[9] = 0x00;
    resp[10] = 0x00;
    resp[11] = 0x00;

    if qtype != TYPE_A {
        // NOERROR with ANCOUNT = 0.
        resp[6] = 0x00;
        resp[7] = 0x00;
        return Some(resp);
    }

    // ANCOUNT = 1.
    resp[6] = 0x00;
    resp[7] = 0x01;
    resp.extend_from_slice(&[
        0xC0, 0x0C, // name pointer -> offset 12 (the question)
        0x00, 0x01, // TYPE = A
        0x00, 0x01, // CLASS = IN
        0x00, 0x00, 0x00, 0x3C, // TTL = 60
        0x00, 0x04, // RDLENGTH = 4
    ]);
    resp.extend_from_slice(&ip);

    Some(resp)
}

/// Read `QTYPE` of the first question, walking the length-prefixed QNAME.
///
/// Returns `None` if the name is truncated or uses compression (a pointer in a
/// question is not something a normal client sends, and following it here would
/// mean handling loops).
fn question_type(query: &[u8]) -> Option<u16> {
    let mut i = QUESTION_OFFSET;
    loop {
        let len = *query.get(i)? as usize;
        if len == 0 {
            i += 1;
            break;
        }
        if len & 0xC0 != 0 {
            return None;
        }
        i += 1 + len;
    }
    let hi = *query.get(i)?;
    let lo = *query.get(i + 1)?;
    // QCLASS must also be present for the question to be well-formed.
    query.get(i + 3)?;
    Some(u16::from_be_bytes([hi, lo]))
}

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

    const IP: [u8; 4] = [192, 168, 71, 1];

    /// Build a minimal single-question query for `name` with the given qtype.
    fn query(name: &str, qtype: u16) -> Vec<u8> {
        let mut q = vec![
            0x12, 0x34, // ID
            0x01, 0x00, // flags: RD
            0x00, 0x01, // QDCOUNT = 1
            0x00, 0x00, // ANCOUNT
            0x00, 0x00, // NSCOUNT
            0x00, 0x00, // ARCOUNT
        ];
        for label in name.split('.') {
            q.push(label.len() as u8);
            q.extend_from_slice(label.as_bytes());
        }
        q.push(0);
        q.extend_from_slice(&qtype.to_be_bytes());
        q.extend_from_slice(&1u16.to_be_bytes()); // QCLASS = IN
        q
    }

    #[test]
    fn answers_a_query_with_the_gateway_ip() {
        let r = build_reply(&query("captive.apple.com", TYPE_A), IP).unwrap();
        // Echoes the ID.
        assert_eq!(&r[0..2], &[0x12, 0x34]);
        // QR + AA set, RCODE 0.
        assert_eq!(r[2], 0x81);
        assert_eq!(r[3], 0x80);
        // One answer.
        assert_eq!(u16::from_be_bytes([r[6], r[7]]), 1);
        // The answer is the last 16 bytes: ptr, type, class, ttl, rdlen, ip.
        let answer = &r[r.len() - 16..];
        assert_eq!(&answer[0..2], &[0xC0, 0x0C]);
        assert_eq!(u16::from_be_bytes([answer[2], answer[3]]), TYPE_A);
        assert_eq!(&answer[12..], &IP);
    }

    #[test]
    fn aaaa_query_gets_empty_noerror() {
        // Regression guard for the behaviour change vs nfc-522: an AAAA query
        // must NOT come back with a type-A answer record.
        let r = build_reply(&query("example.com", 28), IP).unwrap();
        assert_eq!(r[3] & 0x0F, 0, "RCODE must be NOERROR");
        assert_eq!(u16::from_be_bytes([r[6], r[7]]), 0, "no answers");
        // Nothing appended beyond the echoed question.
        assert_eq!(r.len(), query("example.com", 28).len());
    }

    #[test]
    fn clears_authority_and_additional_counts() {
        let mut q = query("a.b", TYPE_A);
        q[8] = 0x00;
        q[9] = 0x01; // pretend NSCOUNT = 1
        q[10] = 0x00;
        q[11] = 0x01; // pretend ARCOUNT = 1
        let r = build_reply(&q, IP).unwrap();
        assert_eq!(&r[8..12], &[0, 0, 0, 0]);
    }

    #[test]
    fn rejects_runt_packets() {
        assert!(build_reply(&[], IP).is_none());
        assert!(build_reply(&[0u8; 11], IP).is_none());
    }

    #[test]
    fn rejects_zero_question_count() {
        let mut q = query("a.b", TYPE_A);
        q[4] = 0;
        q[5] = 0;
        assert!(build_reply(&q, IP).is_none());
    }

    #[test]
    fn rejects_truncated_name() {
        // Claims a 9-byte label but supplies 3.
        let mut q = vec![0x12, 0x34, 0x01, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0];
        q.extend_from_slice(&[9, b'a', b'b', b'c']);
        assert!(build_reply(&q, IP).is_none());
    }

    #[test]
    fn rejects_missing_qtype() {
        // Name terminates but QTYPE/QCLASS are absent.
        let mut q = vec![0x12, 0x34, 0x01, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0];
        q.extend_from_slice(&[1, b'a', 0]);
        assert!(build_reply(&q, IP).is_none());
    }

    #[test]
    fn rejects_compressed_question_name() {
        let mut q = vec![0x12, 0x34, 0x01, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0];
        q.extend_from_slice(&[0xC0, 0x0C, 0, 1, 0, 1]);
        assert!(build_reply(&q, IP).is_none());
    }

    #[test]
    fn handles_root_query() {
        // QNAME = "." (just the zero byte).
        let mut q = vec![0x12, 0x34, 0x01, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0];
        q.extend_from_slice(&[0, 0, 1, 0, 1]);
        let r = build_reply(&q, IP).unwrap();
        assert_eq!(u16::from_be_bytes([r[6], r[7]]), 1);
    }
}