netring 0.25.0

//! `XdpProgram` and `XdpAttachment` — RAII guards for a loaded
//! XDP program and its interface attachment.

use aya::Ebpf;
use aya::maps::xdp::XskMap;
use aya::programs::{Xdp, XdpFlags as AyaXdpFlags, links::FdLink};

use crate::afxdp::XdpSocket;
use crate::error::Error;

use super::LoaderError;

bitflags::bitflags! {
    /// Flags controlling how an XDP program attaches to an interface.
    /// Mirrors `IFLA_XDP_FLAGS` in the kernel.
    #[derive(Debug, Clone, Copy)]
    pub struct XdpFlags: u32 {
        /// Generic SKB-mode (works on every driver, slower).
        const SKB_MODE = 1 << 1;
        /// Native driver mode (preferred when the driver supports it).
        const DRV_MODE = 1 << 2;
        /// Hardware offload (Netronome SmartNIC etc.).
        const HW_MODE  = 1 << 3;
        /// Replace any existing program. Without this, attach fails
        /// with `EBUSY` if a program is already attached.
        const REPLACE  = 1 << 4;
    }
}

impl XdpFlags {
    /// Convert to aya's `XdpFlags`.
    pub(crate) fn to_aya(self) -> AyaXdpFlags {
        let mut out = AyaXdpFlags::default();
        if self.contains(XdpFlags::SKB_MODE) {
            out |= AyaXdpFlags::SKB_MODE;
        }
        if self.contains(XdpFlags::DRV_MODE) {
            out |= AyaXdpFlags::DRV_MODE;
        }
        if self.contains(XdpFlags::HW_MODE) {
            out |= AyaXdpFlags::HW_MODE;
        }
        if self.contains(XdpFlags::REPLACE) {
            out |= AyaXdpFlags::REPLACE;
        }
        out
    }
}

/// Key for the table-driven filter program's `filter_map` (0.25 W1a). Layout
/// must match `struct filter_key` in `filter_redirect.bpf.c`: a host-order L4
/// port, the IP protocol byte, and an explicit zero pad so the 4-byte key has
/// no uninitialised slack (HASH map keys are compared bytewise).
#[repr(C)]
#[derive(Clone, Copy)]
pub(crate) struct FilterKey {
    pub(crate) port: u16,
    pub(crate) proto: u8,
    pub(crate) pad: u8,
}

// SAFETY: `FilterKey` is plain-old-data — `#[repr(C)]`, 4 bytes with no padding
// (the trailing `pad` field fills the slack), and valid for any bit pattern.
unsafe impl aya::Pod for FilterKey {}

/// A loaded XDP program plus its embedded `BPF_MAP_TYPE_XSKMAP`.
///
/// Construct via [`super::default_program`] (redirect-all) or
/// [`super::filter_program`] (table-driven). The wrapped `Ebpf` instance keeps
/// the program and its maps alive.
pub struct XdpProgram {
    bpf: Ebpf,
    /// Name of the program inside the loaded object. Used to look up
    /// the program after the `Ebpf` has been moved.
    program_name: String,
    /// Name of the XSKMAP inside the loaded object.
    map_name: String,
}

impl XdpProgram {
    pub(crate) fn new(bpf: Ebpf, program_name: &str, map_name: &str) -> Self {
        Self {
            bpf,
            program_name: program_name.to_string(),
            map_name: map_name.to_string(),
        }
    }

    /// Wrap a caller-loaded `aya::Ebpf` for use with netring's XDP
    /// attach + register + RAII teardown.
    ///
    /// Use this when you've compiled your own XDP program (e.g. via
    /// `aya-bpf` and `bpf-linker`) and want netring to handle the
    /// kernel attach + AF_XDP socket registration. Your program
    /// must:
    ///
    /// 1. Define a `BPF_MAP_TYPE_XSKMAP` and use it in a
    ///    `bpf_redirect_map(&xsks_map, ctx->rx_queue_index, ...)`
    ///    call.
    /// 2. Be loaded by aya (`aya::Ebpf::load(...)`) before you call
    ///    this.
    ///
    /// `program_name` is the BPF program symbol name (the Rust
    /// `#[xdp(name = "...")]` or C `SEC("xdp")` function name).
    /// `map_name` is the XSKMAP's variable name in your source.
    ///
    /// ```no_run,ignore
    /// use aya::Ebpf;
    /// use netring::xdp::{XdpFlags, XdpProgram};
    ///
    /// let bpf = Ebpf::load(MY_BYTECODE)?;
    /// let mut prog = XdpProgram::from_aya(bpf, "my_xdp", "xsks_map");
    /// // build an XdpSocket as usual, then:
    /// prog.register(queue_id, &xsk)?;
    /// let _attachment = prog.attach("eth0", XdpFlags::DRV_MODE)?;
    /// ```
    pub fn from_aya(bpf: Ebpf, program_name: &str, map_name: &str) -> Self {
        Self::new(bpf, program_name, map_name)
    }

    /// Borrow the program's `XskMap` for setting/unsetting socket fds.
    pub fn xsk_map(&mut self) -> Result<XskMap<&mut aya::maps::MapData>, Error> {
        let map = self
            .bpf
            .map_mut(&self.map_name)
            .ok_or(LoaderError::MapMissing("xsks_map"))?;
        XskMap::try_from(map).map_err(|e| LoaderError::Aya(e.to_string()).into())
    }

    /// Register an AF_XDP socket on the program's XSKMAP at
    /// `queue_id`. Convenience wrapper around `xsk_map().set(...)`.
    pub fn register(&mut self, queue_id: u32, xsk: &XdpSocket) -> Result<(), Error> {
        use std::os::fd::AsRawFd;
        let mut map = self.xsk_map()?;
        let fd: i32 = xsk.as_raw_fd();
        map.set(queue_id, fd, 0)
            .map_err(|e| LoaderError::Aya(e.to_string()).into())
    }

    /// Insert or remove a `{proto, port}` entry in the table-driven filter
    /// program's `filter_map` (0.25 W1a / S5).
    ///
    /// `on = true` marks the tuple as interesting — frames whose IP protocol is
    /// `proto` (`6` TCP / `17` UDP) and whose source *or* destination L4 port is
    /// `port` are redirected into the AF_XDP socket in-kernel; everything else
    /// is `XDP_PASS`ed up the normal stack ("shed"). `on = false` removes the
    /// entry. The `port` is the plain host-order number (e.g. `443`).
    ///
    /// Only meaningful for a program loaded via [`super::filter_program`] (which
    /// embeds `filter_map`); on the redirect-all program this returns
    /// [`LoaderError::MapMissing`].
    pub fn set_filter(&mut self, proto: u8, port: u16, on: bool) -> Result<(), Error> {
        use aya::maps::HashMap as AyaHashMap;
        let map = self
            .bpf
            .map_mut("filter_map")
            .ok_or(LoaderError::MapMissing("filter_map"))?;
        let mut hm: AyaHashMap<_, FilterKey, u8> =
            AyaHashMap::try_from(map).map_err(|e| LoaderError::Aya(e.to_string()))?;
        let key = FilterKey {
            port,
            proto,
            pad: 0,
        };
        if on {
            hm.insert(key, 1u8, 0)
                .map_err(|e| LoaderError::Aya(e.to_string()))?;
        } else {
            hm.remove(&key)
                .map_err(|e| LoaderError::Aya(e.to_string()))?;
        }
        Ok(())
    }

    /// Attach the program to interface `iface` with the given mode.
    /// On Drop of the returned [`XdpAttachment`], the program is
    /// detached.
    pub fn attach(mut self, iface: &str, flags: XdpFlags) -> Result<XdpAttachment, Error> {
        let prog: &mut Xdp = self
            .bpf
            .program_mut(&self.program_name)
            .ok_or(LoaderError::SymbolMissing("xdp_sock_prog"))?
            .try_into()
            .map_err(|e: aya::programs::ProgramError| LoaderError::Aya(e.to_string()))?;
        prog.load().map_err(|e| LoaderError::Aya(e.to_string()))?;
        let link_id = prog
            .attach(iface, flags.to_aya())
            .map_err(|e| LoaderError::Aya(e.to_string()))?;
        // Take ownership of the link so we can detach in Drop.
        let xdp_link = prog
            .take_link(link_id)
            .map_err(|e| LoaderError::Aya(e.to_string()))?;
        let owned: FdLink = xdp_link
            .try_into()
            .map_err(|e: aya::programs::links::LinkError| LoaderError::Aya(e.to_string()))?;
        Ok(XdpAttachment {
            _link: owned,
            _program: self,
        })
    }
}

impl std::fmt::Debug for XdpProgram {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("XdpProgram")
            .field("program_name", &self.program_name)
            .field("map_name", &self.map_name)
            .finish()
    }
}

/// RAII guard for an attached XDP program. Drop detaches the program
/// from the interface; the program and map are then unloaded when
/// the inner `Ebpf` drops.
pub struct XdpAttachment {
    _link: FdLink,
    _program: XdpProgram,
}

impl XdpAttachment {
    /// Borrow the program's XSKMAP after attachment, e.g. to register
    /// additional sockets onto it.
    pub fn xsk_map(&mut self) -> Result<XskMap<&mut aya::maps::MapData>, Error> {
        self._program.xsk_map()
    }

    /// Update the table-driven filter program's `filter_map` after attach
    /// (0.25 W1a). Delegates to [`XdpProgram::set_filter`] — see there for the
    /// `{proto, port}` semantics. Lets a running capture add/remove kernel-side
    /// `{proto,port}` interest without re-attaching the program.
    pub fn set_filter(&mut self, proto: u8, port: u16, on: bool) -> Result<(), Error> {
        self._program.set_filter(proto, port, on)
    }
}

impl std::fmt::Debug for XdpAttachment {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("XdpAttachment")
            .field("program", &self._program)
            .finish()
    }
}