OpenPXE/crates/dhcp-proxy/src/server.rs

//! UDP listener loop for the DHCP proxy. Accepts on :67 (and :4011 on a
//! second socket) and dispatches each datagram through the pure reply logic.

use crate::escalation::DriverEscalation;
use crate::reply::{build_reply, decide, BootDirective, ReplyContext};
use dhcproto::v4::{DhcpOption, Message, OptionCode};
use dhcproto::{Decodable, Decoder, Encodable, Encoder};
use openpxe_core::{ClientArch, ClientEvent, ClientRegistry, DriverMode, FirmwareClass};
use socket2::{Domain, Protocol, Socket, Type};
use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4};
use std::sync::Arc;
use tokio::net::UdpSocket;

pub struct DhcpProxyServer {
    bind: IpAddr,
    dhcp_port: u16,
    pxe_port: u16,
    our_ip: Ipv4Addr,
    public_base_url: String,
    clients: Arc<ClientRegistry>,
    metrics: openpxe_core::Metrics,
    /// Automatic per-MAC NIC driver-mode escalation (v0.6.1). Shared across
    /// the :67 and :4011 listener tasks via the server `Arc`.
    escalation: DriverEscalation,
}

impl DhcpProxyServer {
    pub fn new(
        bind: IpAddr,
        dhcp_port: u16,
        pxe_port: u16,
        our_ip: Ipv4Addr,
        public_base_url: String,
        clients: Arc<ClientRegistry>,
        metrics: openpxe_core::Metrics,
    ) -> Self {
        Self {
            bind,
            dhcp_port,
            pxe_port,
            our_ip,
            public_base_url,
            clients,
            metrics,
            escalation: DriverEscalation::new(),
        }
    }

    pub async fn run(self) -> anyhow::Result<()> {
        let dhcp_sock = bind_udp(self.bind, self.dhcp_port, true)?;
        let pxe_sock = bind_udp(self.bind, self.pxe_port, false)?;
        tracing::info!(
            target: "openpxe::dhcp",
            "DHCP proxy listening on {}:{} and :{}",
            self.bind, self.dhcp_port, self.pxe_port
        );

        let ctx = Arc::new(self);
        let c1 = ctx.clone();
        let c2 = ctx.clone();
        let a = tokio::spawn(async move { c1.serve_loop(dhcp_sock, "67").await });
        let b = tokio::spawn(async move { c2.serve_loop(pxe_sock, "4011").await });
        let _ = tokio::try_join!(a, b)?;
        Ok(())
    }

    async fn serve_loop(&self, sock: UdpSocket, label: &'static str) -> anyhow::Result<()> {
        let mut buf = vec![0u8; 4096];
        loop {
            let (n, from) = match sock.recv_from(&mut buf).await {
                Ok(v) => v,
                Err(e) => {
                    tracing::warn!(target: "openpxe::dhcp", port=label, "recv error: {e}");
                    continue;
                }
            };
            if let Err(e) = self.handle_datagram(&sock, &buf[..n], from, label).await {
                tracing::warn!(target: "openpxe::dhcp", port=label, "handle error: {e}");
            }
        }
    }

    async fn handle_datagram(
        &self,
        sock: &UdpSocket,
        data: &[u8],
        from: SocketAddr,
        label: &'static str,
    ) -> anyhow::Result<()> {
        let request = Message::decode(&mut Decoder::new(data))?;

        let vendor_class = request
            .opts()
            .get(OptionCode::ClassIdentifier)
            .and_then(|o| {
                if let DhcpOption::ClassIdentifier(v) = o {
                    Some(v.as_slice())
                } else {
                    None
                }
            });
        let user_class = request.opts().get(OptionCode::UserClass).and_then(|o| {
            if let DhcpOption::UserClass(v) = o {
                Some(v.as_slice())
            } else {
                None
            }
        });
        let class = FirmwareClass::classify(vendor_class, user_class);
        if matches!(class, FirmwareClass::Other) {
            // Not a PXE client (e.g. a regular DHCP DISCOVER from a phone).
            // Silently ignore — we are a proxy, we only speak to PXE clients.
            return Ok(());
        }

        // dhcproto types option 93 as an enum that drops unknown codes;
        // re-parse from the raw wire bytes so firmware quirks like 0x0009
        // come through intact.
        let raw_arch = extract_raw_arch(data).unwrap_or(0);
        let arch = ClientArch::from_option_93(raw_arch);

        let chaddr = request.chaddr();
        let mac = format_mac(chaddr);
        self.clients.record(
            &mac,
            None,
            Some(arch),
            match label {
                "4011" => ClientEvent::PxeBootServerRequest,
                _ => ClientEvent::DhcpDiscover,
            },
        );

        // Automatic NIC driver-mode selection (v0.6.1). The default is
        // firmware-net (snponly/undionly). A successful iPXE handoff confirms
        // the current mode works for this MAC; a fresh firmware boot whose
        // predecessor never handed off escalates the MAC to iPXE's built-in
        // NIC drivers. No operator toggle — the firmware path is unchanged so
        // hardware that already boots never regresses.
        let driver_mode = match class {
            FirmwareClass::IpxeUserClass => {
                self.escalation.mark_ipxe_success(&mac);
                DriverMode::Firmware // unused: this path serves the HTTP script
            }
            FirmwareClass::PxeClient | FirmwareClass::HttpClient => self
                .escalation
                .mode_for_firmware_attempt(&mac, label == "67"),
            // Unreachable: FirmwareClass::Other returned above.
            FirmwareClass::Other => DriverMode::Firmware,
        };

        let ctx = ReplyContext {
            request: &request,
            our_ip: self.our_ip,
            arch,
            class,
            driver_mode,
            public_base_url: &self.public_base_url,
        };
        let directive = decide(&ctx);
        if matches!(directive, BootDirective::Ignore) {
            self.metrics.record_dhcp_decline();
            tracing::debug!(
                target: "openpxe::dhcp",
                mac=%mac, arch=?arch, "ignoring — no bootfile for arch"
            );
            return Ok(());
        }
        self.metrics.record_dhcp_reply(arch.as_str());

        let Some(reply) = build_reply(&ctx, &directive) else {
            return Ok(());
        };
        let mut out = Vec::with_capacity(512);
        reply.encode(&mut Encoder::new(&mut out))?;

        let dest = reply_destination(&request, from);
        sock.send_to(&out, dest).await?;
        tracing::info!(
            target: "openpxe::dhcp",
            mac=%mac, arch=arch.as_str(), class=?class, driver=?driver_mode, dest=%dest, directive=?directive,
            "PXE reply sent"
        );
        Ok(())
    }
}

/// Choose where to send the reply. DHCP semantics (RFC 2131 §4.1):
///   1. If the request came via a relay agent (`giaddr` != 0), reply to
///      that agent on port 67. The relay will forward to the client.
///   2. If the client already has an IP (`ciaddr`), unicast there on :68.
///   3. If the broadcast flag is set in the BOOTP flags (bit 15), the
///      client cannot receive unicast frames yet — we MUST broadcast.
///   4. Otherwise, per the spec we MAY unicast to `chaddr` if we ARP-inject,
///      but since we don't craft raw frames (proxy mode, no NET_RAW), we
///      fall back to broadcast which every client accepts.
///   5. Special case for the PXE Boot Server port 4011: reply to the
///      source address/port exactly — this is a unicast query and the
///      client expects a unicast answer there.
fn reply_destination(request: &Message, from: SocketAddr) -> SocketAddr {
    // (1) relayed request
    let giaddr = request.giaddr();
    if giaddr != Ipv4Addr::UNSPECIFIED {
        return SocketAddr::V4(SocketAddrV4::new(giaddr, 67));
    }
    // (5) PXE Boot Server discovery is unicast
    if from.port() == 4011 {
        return from;
    }
    // (2) client has an IP and has NOT requested broadcast-only
    let ciaddr = request.ciaddr();
    let bflag = request.flags().broadcast();
    if ciaddr != Ipv4Addr::UNSPECIFIED && !bflag {
        return SocketAddr::V4(SocketAddrV4::new(ciaddr, 68));
    }
    // (3, 4) broadcast to 255.255.255.255:68
    SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::BROADCAST, 68))
}

fn bind_udp(bind: IpAddr, port: u16, broadcast: bool) -> anyhow::Result<UdpSocket> {
    let domain = match bind {
        IpAddr::V4(_) => Domain::IPV4,
        IpAddr::V6(_) => Domain::IPV6,
    };
    let sock = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP))?;
    sock.set_reuse_address(true)?;
    #[cfg(unix)]
    sock.set_reuse_port(true)?;
    if broadcast {
        sock.set_broadcast(true)?;
    }
    sock.set_nonblocking(true)?;
    let addr: SocketAddr = SocketAddr::new(bind, port);
    sock.bind(&addr.into())?;
    let std_sock: std::net::UdpSocket = sock.into();
    Ok(UdpSocket::from_std(std_sock)?)
}

fn format_mac(chaddr: &[u8]) -> String {
    use std::fmt::Write;
    // One allocation — this runs for every PXE datagram we answer.
    let mut s = String::with_capacity(17);
    for (i, b) in chaddr.iter().take(6).enumerate() {
        if i > 0 {
            s.push(':');
        }
        let _ = write!(s, "{b:02x}");
    }
    s
}

/// Walk raw DHCP options looking for option 93 (Client System Architecture)
/// and return the first 2-byte big-endian value. This bypasses dhcproto's
/// typed decoding because some firmwares emit values outside the IANA table
/// that the typed decoder may drop.
fn extract_raw_arch(packet: &[u8]) -> Option<u16> {
    // DHCPv4 fixed header is 240 bytes including the 4-byte magic cookie.
    // Options start at offset 240.
    let opts = packet.get(240..)?;
    let mut i = 0;
    while i < opts.len() {
        let code = opts[i];
        if code == 0xff {
            return None;
        } // END
        if code == 0x00 {
            i += 1;
            continue;
        } // PAD
        i += 1;
        if i >= opts.len() {
            return None;
        }
        let len = opts[i] as usize;
        i += 1;
        if code == 93 && len >= 2 && i + 2 <= opts.len() {
            return Some(u16::from_be_bytes([opts[i], opts[i + 1]]));
        }
        i += len;
    }
    None
}

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

    #[test]
    fn extracts_arch_from_raw_options() {
        // Minimal BOOTP header + magic cookie + option 93 (arch)=0x0007 + END.
        let mut pkt = vec![0u8; 240];
        pkt[236..240].copy_from_slice(&[99, 130, 83, 99]); // magic cookie
        pkt.extend_from_slice(&[53, 1, 1]); // option 53 DHCPDISCOVER
        pkt.extend_from_slice(&[93, 2, 0x00, 0x07]);
        pkt.push(0xff);
        assert_eq!(extract_raw_arch(&pkt), Some(0x0007));
    }
}