runbound 0.4.4

// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright (C) 2024-2026 RedLemonBe — https://github.com/redlemonbe/Runbound
// XDP worker: poll loop that reads raw Ethernet frames from the AF_XDP RX
// ring, processes local-zone DNS queries entirely in user space, and writes
// DNS responses back to the TX ring.
//
// Each worker owns one XskSocket (one NIC queue).  Workers run on dedicated
// OS threads so the hot path never contends with the Tokio executor.
//
// Fast path (local query answered here):
//   poll() → consume_rx() → parse eth/ip/udp/dns → LocalZoneSet lookup →
//   build response → craft reply frame → enqueue_tx() → kick if needed →
//   return frames to fill ring
//
// Slow path (recursive / unknown name):
//   The XDP program was configured with XDP_PASS fallback, so these packets
//   never reached the AF_XDP socket — hickory-server handles them normally.
//   Within the worker, a query whose name isn't found locally returns None
//   from process_packet(); the frame is recycled without a TX response.

use std::net::IpAddr;
use std::sync::Arc;

use arc_swap::ArcSwap;
use hickory_proto::{
    op::{Message, MessageType, OpCode, ResponseCode},
    serialize::binary::{BinDecodable, BinEncodable, BinEncoder},
};

use crate::dns::acl::{Acl, AclAction};
use crate::dns::local::{LocalZoneSet, ZoneAction};
use crate::dns::RateLimiter;
use super::loader::XdpHandle;
use super::socket::{
    XskSocket, create_xsk_socket, get_rx_queue_count, iface_index,
};
use super::umem::{XdpDesc, FRAME_SIZE};

const ETH_HDR: usize = 14;
const IPV4_HDR_MIN: usize = 20;
const IPV6_HDR: usize = 40;
const UDP_HDR: usize = 8;

const ETH_P_IP:   u16 = 0x0800;
const ETH_P_IPV6: u16 = 0x86DD;
const PROTO_UDP:   u8 = 17;

/// Load the XDP program onto `iface`, open one AF_XDP socket per RX queue,
/// and spawn a dedicated OS thread for each.  Returns the `XdpHandle` that
/// must be kept alive for as long as the fast path should remain active.
///
/// The `rate_limiter` is shared with the normal hickory-server path so that
/// the per-IP budget is consumed identically whether the packet is handled by
/// the XDP fast path or the kernel / Tokio slow path.
pub fn start_xdp(
    iface:        &str,
    zones:        Arc<ArcSwap<LocalZoneSet>>,
    rate_limiter: Arc<RateLimiter>,
    acl:          Arc<Acl>,
) -> Result<XdpHandle, String> {
    let ifidx = iface_index(iface)
        .ok_or_else(|| format!("interface {iface} not found"))?;

    let mut handle = XdpHandle::load(iface)?;

    let queue_count = get_rx_queue_count(iface).max(1);
    tracing::info!(iface = %iface, queues = queue_count, "Starting XDP workers");

    for q in 0..queue_count {
        let sock = unsafe { create_xsk_socket(ifidx, q, true) }
            .unwrap_or_else(|_| unsafe {
                create_xsk_socket(ifidx, q, false)
                    .expect("AF_XDP socket creation failed even in copy mode")
            });

        handle.register_socket(q, sock.fd)?;

        let z   = Arc::clone(&zones);
        let rl  = Arc::clone(&rate_limiter);
        let acl = Arc::clone(&acl);
        std::thread::Builder::new()
            .name(format!("xdp-{iface}-q{q}"))
            .spawn(move || xdp_worker(sock, z, rl, acl))
            .map_err(|e| format!("thread spawn: {e}"))?;
    }

    Ok(handle)
}

/// Poll loop for one NIC queue. Runs until the socket fd is closed.
fn xdp_worker(
    mut sock:     XskSocket,
    zones:        Arc<ArcSwap<LocalZoneSet>>,
    rate_limiter: Arc<RateLimiter>,
    acl:          Arc<Acl>,
) {
    use libc::{poll, pollfd, POLLIN};

    loop {
        sock.umem.reclaim_tx();

        let mut pfd = pollfd { fd: sock.fd, events: POLLIN, revents: 0 };
        let ret = unsafe { poll(&mut pfd, 1, 1 /* ms timeout */) };
        if ret < 0 {
            break;
        }

        let rxds = sock.rx.consume_rx();
        if rxds.is_empty() {
            continue;
        }

        let snapshot = zones.load();
        let mut tx_descs: Vec<XdpDesc> = Vec::with_capacity(rxds.len());
        let mut rx_addrs: Vec<u64>     = Vec::with_capacity(rxds.len());

        for desc in &rxds {
            rx_addrs.push(desc.addr);

            if let Some(tx_addr) = sock.umem.tx_free.pop_front() {
                let (rx_frame, tx_frame) = unsafe {
                    let rx = std::slice::from_raw_parts(
                        sock.umem.area.add(desc.addr as usize),
                        desc.len as usize,
                    );
                    let tx = std::slice::from_raw_parts_mut(
                        sock.umem.area.add(tx_addr as usize),
                        FRAME_SIZE as usize,
                    );
                    (rx, tx)
                };

                // Rate-limit check before processing: extract source IP from
                // the raw frame and consume a token from the shared bucket.
                // Dropped frames are silently recycled — no REFUSED response
                // is crafted in the XDP path (matches `deny` ACL semantics).
                let src_ip = extract_src_ip(rx_frame);
                if src_ip.map(|ip| !rate_limiter.check(ip)).unwrap_or(false) {
                    sock.umem.tx_free.push_back(tx_addr);
                    continue;
                }

                match process_packet(rx_frame, tx_frame, &snapshot, &acl, src_ip) {
                    Some(tx_len) => tx_descs.push(XdpDesc {
                        addr: tx_addr,
                        len:  tx_len as u32,
                        options: 0,
                    }),
                    None => {
                        sock.umem.tx_free.push_back(tx_addr);
                    }
                }
            }
        }

        // Return consumed RX frames to the kernel fill ring.
        sock.umem.fill.enqueue_batch(&rx_addrs);

        if !tx_descs.is_empty() {
            sock.tx.enqueue_tx(&tx_descs);

            // Kick the driver if it set the NEED_WAKEUP flag.
            if sock.tx.needs_wakeup() {
                unsafe {
                    libc::sendto(
                        sock.fd,
                        std::ptr::null(),
                        0,
                        libc::MSG_DONTWAIT,
                        std::ptr::null(),
                        0,
                    );
                }
            }
        }
    }
}

/// Extract the source IP address from a raw Ethernet frame (IPv4 or IPv6).
/// Returns None for non-IP frames or frames that are too short.
#[inline]
fn extract_src_ip(rx: &[u8]) -> Option<IpAddr> {
    if rx.len() < ETH_HDR { return None; }
    let ethertype = u16::from_be_bytes([rx[12], rx[13]]);
    match ethertype {
        ETH_P_IP => {
            if rx.len() < ETH_HDR + 20 { return None; }
            let src: [u8; 4] = rx[ETH_HDR + 12..ETH_HDR + 16].try_into().ok()?;
            Some(IpAddr::V4(std::net::Ipv4Addr::from(src)))
        }
        ETH_P_IPV6 => {
            if rx.len() < ETH_HDR + 40 { return None; }
            let src: [u8; 16] = rx[ETH_HDR + 8..ETH_HDR + 24].try_into().ok()?;
            Some(IpAddr::V6(std::net::Ipv6Addr::from(src)))
        }
        _ => None,
    }
}

/// Parse a raw Ethernet frame, answer the DNS query from `zones`, and write
/// the response frame into `tx`.  Returns the number of bytes written on
/// success, or `None` if this frame should not receive an XDP reply.
fn process_packet(
    rx:     &[u8],
    tx:     &mut [u8],
    zones:  &LocalZoneSet,
    acl:    &Acl,
    src_ip: Option<IpAddr>,
) -> Option<usize> {
    // ── Ethernet ─────────────────────────────────────────────────────────────
    if rx.len() < ETH_HDR { return None; }
    let ethertype = u16::from_be_bytes([rx[12], rx[13]]);

    let (ip_off, is_v6) = match ethertype {
        ETH_P_IP   => (ETH_HDR, false),
        ETH_P_IPV6 => (ETH_HDR, true),
        _          => return None,
    };

    // ── IP ───────────────────────────────────────────────────────────────────
    let (udp_off, ip_hdr_len, src_ip_off, dst_ip_off, ip_len_off) = if !is_v6 {
        if rx.len() < ip_off + IPV4_HDR_MIN { return None; }
        if rx[ip_off + 9] != PROTO_UDP { return None; }
        let ihl = (rx[ip_off] & 0x0F) as usize * 4;
        if ihl < 20 || ihl > 60 { return None; }
        (ip_off + ihl, ihl, ip_off + 12, ip_off + 16, ip_off + 2)
    } else {
        if rx.len() < ip_off + IPV6_HDR { return None; }
        if rx[ip_off + 6] != PROTO_UDP { return None; }
        (ip_off + IPV6_HDR, IPV6_HDR, ip_off + 8, ip_off + 24, ip_off + 4)
    };

    // ── UDP ──────────────────────────────────────────────────────────────────
    if rx.len() < udp_off + UDP_HDR { return None; }
    let src_port = u16::from_be_bytes([rx[udp_off],     rx[udp_off + 1]]);
    let dst_port = u16::from_be_bytes([rx[udp_off + 2], rx[udp_off + 3]]);
    if dst_port != 53 { return None; }

    let dns_off = udp_off + UDP_HDR;
    if rx.len() <= dns_off { return None; }
    let dns_in = &rx[dns_off..];

    // ── DNS ──────────────────────────────────────────────────────────────────
    let mut dns_out: Vec<u8> = Vec::with_capacity(512);
    if !answer_dns(dns_in, zones, acl, src_ip, &mut dns_out) {
        return None; // not a local query or ACL deny — let it fall through / drop
    }

    // ── Build reply frame ────────────────────────────────────────────────────
    let reply_len = dns_off + dns_out.len();
    if reply_len > tx.len() { return None; }

    // Ethernet: swap src ↔ dst MAC
    tx[0..6].copy_from_slice(&rx[6..12]);
    tx[6..12].copy_from_slice(&rx[0..6]);
    tx[12..14].copy_from_slice(&rx[12..14]);

    if !is_v6 {
        // IPv4: copy then fix length, swap src/dst, recompute checksum
        tx[ip_off..ip_off + ip_hdr_len].copy_from_slice(&rx[ip_off..ip_off + ip_hdr_len]);
        let new_tot = (ip_hdr_len + UDP_HDR + dns_out.len()) as u16;
        tx[ip_len_off..ip_len_off + 2].copy_from_slice(&new_tot.to_be_bytes());

        let src: [u8; 4] = rx[src_ip_off..src_ip_off + 4].try_into().ok()?;
        let dst: [u8; 4] = rx[dst_ip_off..dst_ip_off + 4].try_into().ok()?;
        tx[ip_off + 12..ip_off + 16].copy_from_slice(&dst);
        tx[ip_off + 16..ip_off + 20].copy_from_slice(&src);

        // Clear then recompute IPv4 header checksum
        tx[ip_off + 10..ip_off + 12].fill(0);
        let cksum = ipv4_checksum(&tx[ip_off..ip_off + ip_hdr_len]);
        tx[ip_off + 10..ip_off + 12].copy_from_slice(&cksum.to_be_bytes());
    } else {
        // IPv6: copy, set payload length, swap src/dst
        tx[ip_off..ip_off + IPV6_HDR].copy_from_slice(&rx[ip_off..ip_off + IPV6_HDR]);
        let payload_len = (UDP_HDR + dns_out.len()) as u16;
        tx[ip_len_off..ip_len_off + 2].copy_from_slice(&payload_len.to_be_bytes());

        let src: [u8; 16] = rx[src_ip_off..src_ip_off + 16].try_into().ok()?;
        let dst: [u8; 16] = rx[dst_ip_off..dst_ip_off + 16].try_into().ok()?;
        tx[ip_off + 8..ip_off + 24].copy_from_slice(&dst);
        tx[ip_off + 24..ip_off + 40].copy_from_slice(&src);
    }

    // UDP: swap ports, set length
    let udp_len = (UDP_HDR + dns_out.len()) as u16;
    tx[udp_off..udp_off + 2].copy_from_slice(&dst_port.to_be_bytes()); // src = 53
    tx[udp_off + 2..udp_off + 4].copy_from_slice(&src_port.to_be_bytes());
    tx[udp_off + 4..udp_off + 6].copy_from_slice(&udp_len.to_be_bytes());

    // Compute UDP checksum using the reply frame already in tx
    tx[udp_off + 6..udp_off + 8].fill(0);
    let cksum = if !is_v6 {
        let si: [u8; 4] = tx[ip_off + 12..ip_off + 16].try_into().ok()?;
        let di: [u8; 4] = tx[ip_off + 16..ip_off + 20].try_into().ok()?;
        udp_checksum_v4(&si, &di, &tx[udp_off..udp_off + UDP_HDR + dns_out.len()])
    } else {
        let si: [u8; 16] = tx[ip_off + 8..ip_off + 24].try_into().ok()?;
        let di: [u8; 16] = tx[ip_off + 24..ip_off + 40].try_into().ok()?;
        udp_checksum_v6(&si, &di, &tx[udp_off..udp_off + UDP_HDR + dns_out.len()])
    };
    tx[udp_off + 6..udp_off + 8].copy_from_slice(&cksum.to_be_bytes());

    // DNS payload
    tx[dns_off..dns_off + dns_out.len()].copy_from_slice(&dns_out);

    Some(reply_len)
}

/// Parse `query_bytes` as a DNS query, look it up in `zones`, write the
/// serialised response into `out`.  Returns false if the query should be
/// forwarded to the kernel (non-local name, not a standard query, ACL deny, etc.).
///
/// ACL enforcement in the XDP path:
///   Allow  → proceed normally.
///   Deny   → silent drop (return false, no TX frame crafted).
///   Refuse → craft a REFUSED response and return true so the TX path sends it.
fn answer_dns(
    query_bytes: &[u8],
    zones:       &LocalZoneSet,
    acl:         &Acl,
    src_ip:      Option<IpAddr>,
    out:         &mut Vec<u8>,
) -> bool {
    let msg = match Message::from_bytes(query_bytes) {
        Ok(m) => m,
        Err(_) => return false,
    };
    if msg.message_type() != MessageType::Query { return false; }
    if msg.op_code() != OpCode::Query { return false; }

    let q = match msg.queries().first() {
        Some(q) => q,
        None    => return false,
    };

    // ── ACL check ─────────────────────────────────────────────────────────
    // Applied before zone lookup so that Deny/Refuse clients cannot probe
    // local zone membership even in the XDP fast path.
    if let Some(ip) = src_ip {
        match acl.check(ip) {
            AclAction::Allow  => {}
            AclAction::Deny   => return false, // silent drop — no response
            AclAction::Refuse => {
                // Craft a minimal REFUSED response and send it.
                let mut refused = Message::new();
                refused.set_id(msg.id());
                refused.set_message_type(MessageType::Response);
                refused.set_op_code(OpCode::Query);
                refused.set_response_code(ResponseCode::Refused);
                refused.set_recursion_desired(msg.recursion_desired());
                refused.add_query(q.clone());
                let mut enc = BinEncoder::new(out);
                return refused.emit(&mut enc).is_ok();
            }
        }
    }

    let name  = q.name();
    let rtype = q.query_type();

    // ANY queries go to the normal server (which returns NOTIMP per RFC 8482)
    if rtype == hickory_proto::rr::RecordType::ANY { return false; }

    let mut resp = Message::new();
    resp.set_id(msg.id());
    resp.set_message_type(MessageType::Response);
    resp.set_op_code(OpCode::Query);
    resp.set_recursion_desired(msg.recursion_desired());
    resp.set_recursion_available(false);
    resp.add_query(q.clone());

    match zones.find(name) {
        Some(ZoneAction::Refuse) => {
            resp.set_response_code(ResponseCode::Refused);
            resp.set_authoritative(false);
        }
        Some(ZoneAction::NxDomain) => {
            resp.set_response_code(ResponseCode::NXDomain);
            resp.set_authoritative(true);
        }
        Some(ZoneAction::Static) | Some(ZoneAction::Redirect) => {
            resp.set_authoritative(true);
            let records = zones.local_records(name, rtype);
            if !records.is_empty() {
                resp.set_response_code(ResponseCode::NoError);
                for r in records {
                    resp.add_answer(r.clone());
                }
            } else if zones.name_has_records(name) {
                // NODATA — name exists, wrong type (RFC 2308)
                resp.set_response_code(ResponseCode::NoError);
            } else {
                resp.set_response_code(ResponseCode::NXDomain);
            }
        }
        // Name not in any local zone — forward to kernel / hickory-server
        None => return false,
    }

    let mut enc = BinEncoder::new(out);
    resp.emit(&mut enc).is_ok()
}

// ── Checksum helpers ──────────────────────────────────────────────────────────

fn ones_complement_sum(data: &[u8]) -> u32 {
    let mut sum: u32 = 0;
    let mut i = 0;
    while i + 1 < data.len() {
        sum += u16::from_be_bytes([data[i], data[i + 1]]) as u32;
        i += 2;
    }
    if data.len() % 2 == 1 {
        sum += (data[data.len() - 1] as u32) << 8;
    }
    sum
}

fn fold_checksum(mut s: u32) -> u16 {
    while s >> 16 != 0 {
        s = (s & 0xFFFF) + (s >> 16);
    }
    let r = !(s as u16);
    if r == 0 { 0xFFFF } else { r } // RFC 768: 0 is transmitted as all-ones
}

fn ipv4_checksum(header: &[u8]) -> u16 {
    fold_checksum(ones_complement_sum(header))
}

fn udp_checksum_v4(src: &[u8; 4], dst: &[u8; 4], udp: &[u8]) -> u16 {
    let udp_len = udp.len() as u32;
    let s = ones_complement_sum(src)
          + ones_complement_sum(dst)
          + PROTO_UDP as u32
          + udp_len
          + ones_complement_sum(udp);
    fold_checksum(s)
}

fn udp_checksum_v6(src: &[u8; 16], dst: &[u8; 16], udp: &[u8]) -> u16 {
    // IPv6 pseudo-header: src(16) + dst(16) + UDP length(4) + zeros(3) + next-header(1)
    let udp_len = udp.len() as u32;
    let udp_len_bytes = udp_len.to_be_bytes();
    let s = ones_complement_sum(src)
          + ones_complement_sum(dst)
          + ones_complement_sum(&udp_len_bytes)
          + PROTO_UDP as u32
          + ones_complement_sum(udp);
    fold_checksum(s)
}