zenith_linux/

umem.rs

//! UMEM 内存管理
//!
//! 负责 AF_XDP 所需的用户态内存区域管理：
//! - mmap 连续物理内存
//! - 锁定页面（mlock）防止交换
//! - HugePage 优化（可选）
//! - 提供安全的内存访问接口
//!
//! # unsafe 使用
//! 本文件封装 mmap/munmap/mlock/sysconf 等原生系统调用与裸指针操作。
//! 是 crate 内部四个被精确放开 `#[allow(unsafe_code)]` 的文件之一。

#![allow(unsafe_code)]

use crate::error::{LinuxError, Result, UmemError};
use std::sync::atomic::{AtomicU64, Ordering};

/// UMEM 配置
#[derive(Debug, Clone)]
pub struct UmemConfig {
    /// 总大小（字节），必须是 PAGE_SIZE 的倍数
    pub size: usize,
    /// 是否使用 HugePage
    pub hugepage: bool,
    /// 是否锁定内存
    pub locked: bool,
    /// 共享内存标志（用于多个 XSK 共享）
    pub shared: bool,
}

impl Default for UmemConfig {
    fn default() -> Self {
        Self {
            size: 0,
            hugepage: false,
            locked: true,
            shared: false,
        }
    }
}

/// 内存区域信息
#[derive(Debug, Clone)]
pub struct UmemRegion {
    /// 起始地址（虚拟地址）
    pub addr: *mut u8,
    /// 大小（字节）
    pub size: usize,
    /// 是否使用 HugePage
    pub hugepage: bool,
    /// 是否已锁定
    pub locked: bool,
    /// 页偏移（相对于起始地址）
    pub page_offset: u64,
}

// SAFETY: UmemRegion 仅持有 mmap 元数据（基址/长度/标志），无内部可变性。
// 真实共享模型：UmemManager 经 `Arc<UmemManager>` 被多 Worker 共享（XskConfig
// 的 shared_umem / 多队列场景），mmap 区地址在 munmap 前固定；munmap 仅在
// UmemManager Drop 中原子 CAS(1→0) 成功后执行唯一一次，此后 incref 因
// ref_count==0 fail-closed，绝不在活跃引用下提前解除映射。
// 数据竞争由"按帧分片独占"契约在结构层面排除：各 Worker 只访问自己
// 持有的帧区间（frame index 互斥分片，由描述符引擎位图保证单 Owner），
// 同一帧同一时刻不存在两个活跃借用；与内核 DMA 的顺序由 ring 索引的
// Acquire/Release 原子语义建立。因此元数据可安全跨线程传递/引用。
unsafe impl Send for UmemRegion {}
unsafe impl Sync for UmemRegion {}

/// UMEM 管理器
///
/// 管理 AF_XDP 所需的用户态内存区域。
/// 提供 mmap、lock、HugePage 等功能。
pub struct UmemManager {
    /// 内存区域
    region: Option<UmemRegion>,
    /// 配置
    config: UmemConfig,
    /// 引用计数（用于共享模式）
    ref_count: AtomicU64,
    /// 是否已初始化
    initialized: bool,
}

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

impl UmemManager {
    /// 创建 UMEM 管理器
    ///
    /// # 参数
    /// * `config` - UMEM 配置
    ///
    /// # 返回
    /// * `Result<Self>` - 管理器实例
    pub fn new(config: UmemConfig) -> Result<Self> {
        if config.size == 0 {
            return Err(LinuxError::Umem(UmemError::InsufficientSize {
                actual: 0,
                required: 4096,
            }));
        }

        // 检查对齐（复用 crate 级 page_size() 封装：失败回退 4096，无裸 sysconf）
        let page_size = crate::page_size();
        if !config.size.is_multiple_of(page_size) {
            return Err(LinuxError::Umem(UmemError::NotAligned {
                actual: config.size,
                expected: page_size,
            }));
        }

        Ok(Self {
            region: None,
            config,
            ref_count: AtomicU64::new(0),
            initialized: false,
        })
    }

    /// 创建内存映射
    ///
    /// # 返回
    /// * `Result<()>` - 成功或错误
    pub fn create(&mut self) -> Result<()> {
        if self.initialized {
            return Err(LinuxError::Umem(UmemError::AlreadyCreated));
        }

        // 复用 crate 级 page_size() 封装（sysconf 失败时回退 4096）；
        // sysconf 裸调用只保留在 lib.rs 的唯一 wrapper 内，此处不得绕过
        let page_size = crate::page_size();

        // 确定映射标志
        let mut flags = libc::MAP_PRIVATE | libc::MAP_ANONYMOUS;
        if self.config.shared {
            flags = libc::MAP_SHARED | libc::MAP_ANONYMOUS;
        }

        // HugePage 优化
        let region_size = if self.config.hugepage {
            // 尝试 2MB HugePage
            let hugepage_size = 2 * 1024 * 1024;
            // 向上取整到 HugePage 大小（checked 算术：size + hugepage_size - 1
            // 在 size 接近 usize::MAX 时可能溢出，Fail-Closed 返回错误）
            self.config
                .size
                .checked_add(hugepage_size - 1)
                .map(|v| v & !(hugepage_size - 1))
                .ok_or_else(|| {
                    LinuxError::InsufficientResources(
                        "UMEM size 对齐 HugePage 时溢出".to_string(),
                    )
                })?
        } else {
            self.config.size
        };

        if self.config.hugepage {
            flags |= libc::MAP_HUGETLB;
        }

        // SAFETY: mmap 参数均为有效值：
        // - addr: null_mut()，让内核选择映射地址
        // - length: region_size，已按 HugePage 大小对齐
        // - prot: PROT_READ | PROT_WRITE，允许读写
        // - flags: MAP_PRIVATE/MAP_SHARED + MAP_ANONYMOUS，按需可选 MAP_HUGETLB
        // - fd: -1，匿名映射无需文件描述符
        // - offset: 0，匿名映射无偏移
        // 返回值已检查 MAP_FAILED，失败则立即返回错误。
        let addr = unsafe {
            libc::mmap(
                std::ptr::null_mut(),
                region_size,
                libc::PROT_READ | libc::PROT_WRITE,
                flags,
                -1,
                0,
            )
        };

        if addr == libc::MAP_FAILED {
            let err = std::io::Error::last_os_error();
            return Err(LinuxError::Umem(UmemError::MmapFailed(format!(
                "mmap failed: {}",
                err
            ))));
        }

        let addr = addr as *mut u8;

        // 锁定内存（防止交换）
        if self.config.locked {
            // SAFETY: addr 来自上方成功的 mmap 返回值（已检查非 MAP_FAILED），
            // region_size 与 mmap 时使用的大小完全一致，保证 mlock 操作在有效映射区域内。
            let lock_result = unsafe { libc::mlock(addr as *const libc::c_void, region_size) };
            if lock_result != 0 {
                let err = std::io::Error::last_os_error();
                // SAFETY: addr 为刚才成功 mmap 的区域，region_size 与 mmap 时一致。
                // mlock 失败后立即 munmap 清理，此时映射区域仍有效（未被 munmap 过）。
                unsafe {
                    libc::munmap(addr as *mut libc::c_void, region_size);
                }
                return Err(LinuxError::Umem(UmemError::LockFailed(format!(
                    "mlock failed: {}",
                    err
                ))));
            }
        }

        // SAFETY: addr 为有效 mmap 区域，region_size 与 mmap 一致。
        // MADV_HUGEPAGE 建议内核使用 HugePage，仅在启用 hugepage 时调用，
        // 避免在非 HugePage 场景下产生无意义的 madvise 系统调用开销。
        if self.config.hugepage {
            let madvise_ret = unsafe {
                libc::madvise(addr as *mut libc::c_void, region_size, libc::MADV_HUGEPAGE)
            };
            if madvise_ret != 0 {
                tracing::debug!("madvise(MADV_HUGEPAGE) 失败: {}", std::io::Error::last_os_error());
            }
        }

        // 计算页偏移
        let page_offset = (addr as usize / page_size) as u64;

        self.region = Some(UmemRegion {
            addr,
            size: region_size,
            hugepage: self.config.hugepage,
            locked: self.config.locked,
            page_offset,
        });

        self.initialized = true;
        self.ref_count.store(1, Ordering::SeqCst);

        Ok(())
    }

    /// 获取内存区域
    ///
    /// # 返回
    /// * `Option<&UmemRegion>` - 内存区域引用
    pub fn region(&self) -> Option<&UmemRegion> {
        self.region.as_ref()
    }

    /// 获取起始地址
    ///
    /// # 返回
    /// * `*mut u8` - 起始地址
    pub fn as_ptr(&self) -> *mut u8 {
        self.region.as_ref().map_or(std::ptr::null_mut(), |r| r.addr)
    }

    /// 获取 UMEM 映射区的只读切片（checked 边界，越界返回 None）
    ///
    /// 用于真实 AF_XDP 模式：内核 DMA 将收包数据写入本映射区，
    /// 数据面经本方法零拷贝读取帧内容。
    ///
    /// # 安全
    /// 返回切片的生命周期绑定 `&self`；映射区所有权由 `UmemManager` 持有，
    ///  munmap 仅发生在 drop（调用方以 `Arc<UmemManager>` 保证使用期间存活）。
    #[inline]
    pub fn slice(&self, offset: usize, len: usize) -> Option<&[u8]> {
        let end = offset.checked_add(len)?;
        if end > self.config.size {
            return None;
        }
        let addr = self.as_ptr();
        if addr.is_null() {
            return None;
        }
        // SAFETY: offset+len <= self.config.size（上方已校验），addr 为 create()
        // 成功 mmap 的有效区域基址；切片生命周期绑定 &self，映射区在
        // UmemManager drop（munmap）前始终有效；内核 DMA 并发写入本区域是
        // AF_XDP 的固有语义，读取侧观察到的是字节级快照，无 UB。
        Some(unsafe { std::slice::from_raw_parts(addr.cast_const().add(offset), len) })
    }

    /// 获取 UMEM 映射区的可变切片（checked 边界，越界返回 None）
    ///
    /// 用于真实 AF_XDP 模式的 TX 路径：出站帧内容写入本映射区后由内核发送。
    ///
    /// # 安全
    /// 接收 `&self` 借用并返回 `&mut [u8]`，技术上违反 Rust 别名规则
    /// （aliasing rules）。UMEM 为与内核共享的可变内存（mmap MAP_SHARED），
    /// 语义上等价于 `UnsafeCell<[u8]>`。调用方（单 Owner Worker）在数据面
    /// 单线程顺序访问，保证同一帧区域无并发活跃借用，满足 SAFETY 不变量。
    ///
    /// # Safety
    /// 调用者必须确保同一帧区域不会被两个线程同时通过 `slice_mut` 获取可变引用。
    /// UMEM 区域是 mmap 映射的共享内存，不由 Rust 所有权模型管理，
    /// 帧级别的互斥由 Fill Ring / Completion Ring 的单消费者语义保证。
    #[allow(clippy::mut_from_ref)]
    #[inline]
    pub fn slice_mut(&self, offset: usize, len: usize) -> Option<&mut [u8]> {
        let end = offset.checked_add(len)?;
        if end > self.config.size {
            return None;
        }
        let addr = self.as_ptr();
        if addr.is_null() {
            return None;
        }
        // SAFETY: offset+len <= self.config.size（上方已校验），addr 为有效
        // mmap 基址；UMEM 为内核共享可变内存，单 Owner Worker 在数据面单
        // 线程顺序访问，保证同一帧区域无并发活跃借用。
        Some(unsafe { std::slice::from_raw_parts_mut(addr.add(offset), len) })
    }

    /// 获取页偏移
    ///
    /// # 返回
    /// * `u64` - 页偏移
    pub fn page_offset(&self) -> u64 {
        self.region.as_ref().map_or(0, |r| r.page_offset)
    }

    /// 获取大小
    ///
    /// # 返回
    /// * `usize` - 大小（字节）
    pub fn size(&self) -> usize {
        self.config.size
    }

    /// 是否已初始化
    ///
    /// # 返回
    /// * `bool` - 是否已初始化
    pub fn is_initialized(&self) -> bool {
        self.initialized
    }

    /// 增加引用计数（共享模式）
    ///
    /// # Fail-Closed（返回错误而非 panic）
    /// - 若 `UmemManager` 尚未调用 `create()`，返回 `UmemError::NotCreated`
    ///   （避免对未映射内存建立引用）。
    /// - 若引用计数已为 0（说明最后一个引用已触发 `munmap`），返回
    ///   `UmemError::MunmapFailed`（防止 0→1 复活导致 UAF）。
    /// - 若 `fetch_add` 后溢出 `u64::MAX`，Rust 原子为 wrapping，但 `u64` 的空间
    ///   对任何实际场景均足够（2^64 次），不可能溢出。
    ///
    /// # 返回
    /// * `Result<u64>` - 递增后的新引用计数
    pub fn incref(&self) -> Result<u64> {
        // 1) 未初始化：不允许建立引用
        if !self.initialized {
            return Err(LinuxError::Umem(UmemError::NotCreated));
        }
        // 2) CAS 循环：原子地检查 ref_count != 0 并递增，消除 TOCTOU 竞态窗口。
        // 旧实现使用 load(Acquire) + fetch_add(SeqCst) 两步操作，存在
        // load 与 fetch_add 之间的竞态窗口（另一线程可能在两步之间
        // 将 ref_count 从 >0 减到 0 并触发 munmap）。CAS 循环保证
        // 「检查非零」与「递增」在同一原子操作中完成。
        loop {
            let current = self.ref_count.load(Ordering::Acquire);
            if current == 0 {
                return Err(LinuxError::Umem(UmemError::MunmapFailed(
                    "incref: region already munmapped (ref_count == 0, use-after-free prevented)"
                        .to_string(),
                )));
            }
            match self.ref_count.compare_exchange(
                current,
                current + 1,
                Ordering::SeqCst,
                Ordering::Acquire,
            ) {
                Ok(_) => return Ok(current + 1),
                Err(_) => continue,
            }
        }
    }

    /// 减少引用计数
    ///
    /// 返回递减后的新值。递减到 0 **不会**在此处执行 `munmap`；`munmap`
    /// 的唯一入口在 `Drop::drop` 中通过原子 CAS（1 → 0）触发，
    /// 确保每个映射区域精确释放一次且仅一次（无 double-free）。
    pub fn decref(&self) -> u64 {
        // SeqCst：与 Drop 中的 CAS 以及其他线程的 incref/decref 全序。
        self.ref_count.fetch_sub(1, Ordering::SeqCst) - 1
    }
}

impl Drop for UmemManager {
    fn drop(&mut self) {
        if let Some(region) = &self.region {
            // GAP-1 修复：将原来的 `load + if ref_count <= 1` TOCTOU 窗口
            // 替换为一条原子 RMW（compare_exchange 1 → 0）。
            //
            // 成功路径（CAS Ok(1)）：
            //   - RMW 之前 ref_count 精确等于 1
            //   - 本线程成功将其置为 0
            //   - 任何并发的 incref（incref 已使用 CAS 循环，无 TOCTOU 窗口）：
            //     * CAS 之前就完成了递增 → ref_count 变为 ≥2，Drop 的 CAS 失败
            //     * CAS 之后到达 incref → 观察到 ref_count == 0 → incref 返回 Err（Fail-Closed）
            //   - 任何并发的 decref：
            //     * CAS 之前完成 → ref_count < 1 → CAS 失败（安全不释放）
            //     * CAS 之后到达 → ref_count == 0，decref 正常减为 -1（wrapping 到 u64::MAX）
            //       但这是调用者未按「incref/decref 严格成对」契约使用的违例；
            //       实际释放只会发生一次，不会 double-munmap。
            //
            // 失败路径（CAS Err(actual)）：
            //   - actual > 1：仍有活跃引用，跳过 munmap（由未来最后一次 Drop/CAS 触发）
            //   - actual == 0：已被另一条 Drop 路径成功 CAS 并 munmap，跳过
            //   两种情况均安全，绝不 double-free。
            //
            // AcqRel 成功顺序：对 munmap 前的所有写入（本线程 + 其他线程之前通过 ref_count 同步的）
            //   建立 happens-before；失败使用 Acquire 即可（不修改内存）。
            let cas_result = self.ref_count.compare_exchange(
                1,
                0,
                Ordering::AcqRel,
                Ordering::Acquire,
            );
            if cas_result.is_ok() {
                // SAFETY: CAS(1 → 0) 成功保证：
                //   (a) 本映射从未被 munmap 过（ref_count 到 0 是首次，且此后 incref 断言阻止复活）
                //   (b) 没有其他并发 Drop 会再次走到 munmap 分支（它们的 CAS 必失败）
                //   (c) region.addr / region.size 来自之前成功的 mmap，长度与 mmap 时完全匹配
                //   → 精确 munmap 一次，无 double-free，无越界释放。
                unsafe {
                    libc::munmap(region.addr as *mut libc::c_void, region.size);
                }
            }
        }
    }
}

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

    #[test]
    fn test_umem_config_default() {
        let config = UmemConfig::default();
        assert_eq!(config.size, 0);
        assert!(!config.hugepage);
        assert!(config.locked);
        assert!(!config.shared);
    }

    #[test]
    fn test_umem_manager_creation() {
        let config = UmemConfig {
            size: 0,
            ..Default::default()
        };
        let result = UmemManager::new(config);
        assert!(result.is_err());
    }

    #[test]
    fn test_umem_manager_create() {
        let config = UmemConfig {
            size: 4096 * 256, // 1MB
            hugepage: false,
            locked: false, // 测试时不锁定（需要权限）
            shared: false,
        };
        let mut manager = UmemManager::new(config).unwrap();
        assert!(!manager.is_initialized());

        let result = manager.create();
        assert!(result.is_ok());
        assert!(manager.is_initialized());
        assert!(!manager.as_ptr().is_null());
        assert_eq!(manager.size(), 4096 * 256);
    }

    #[test]
    fn test_umem_alignment_check() {
        let config = UmemConfig {
            size: 4096 + 1, // 未对齐
            ..Default::default()
        };
        let result = UmemManager::new(config);
        assert!(result.is_err());
    }

    #[test]
    fn test_umem_double_create() {
        let config = UmemConfig {
            size: 4096 * 64,
            hugepage: false,
            locked: false,
            shared: false,
        };
        let mut manager = UmemManager::new(config).unwrap();
        manager.create().unwrap();
        let result = manager.create();
        assert!(result.is_err());
    }

    #[test]
    fn test_umem_size_zero_rejected() {
        let config = UmemConfig {
            size: 0,
            ..Default::default()
        };
        let result = UmemManager::new(config);
        assert!(result.is_err());
    }

    #[test]
    fn test_umem_unaligned_size_rejected() {
        let page_size = crate::page_size();
        let config = UmemConfig {
            size: page_size + 1,
            ..Default::default()
        };
        let result = UmemManager::new(config);
        assert!(result.is_err());
    }

    #[test]
    fn test_umem_aligned_size_accepted() {
        let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) } as usize;
        for multiplier in [1, 2, 4, 8, 16, 32, 64, 128, 256] {
            let config = UmemConfig {
                size: page_size * multiplier,
                hugepage: false,
                locked: false,
                shared: false,
            };
            let result = UmemManager::new(config);
            assert!(
                result.is_ok(),
                "Size {} ({} * {}) should be valid",
                page_size * multiplier,
                page_size,
                multiplier
            );
        }
    }

    #[test]
    fn test_umem_region_access_after_create() {
        let config = UmemConfig {
            size: 4096 * 128,
            hugepage: false,
            locked: false,
            shared: false,
        };
        let mut manager = UmemManager::new(config).unwrap();
        assert!(manager.region().is_none());
        assert!(manager.as_ptr().is_null());
        assert_eq!(manager.page_offset(), 0);

        manager.create().unwrap();
        assert!(manager.region().is_some());
        assert!(!manager.as_ptr().is_null());

        let region = manager.region().unwrap();
        assert!(!region.addr.is_null());
        assert!(region.size >= 4096 * 128);
        assert!(!region.hugepage);
        assert!(!region.locked);
    }

    #[test]
    fn test_umem_size_method_returns_config_size() {
        let config = UmemConfig {
            size: 4096 * 64,
            hugepage: false,
            locked: false,
            shared: false,
        };
        let manager = UmemManager::new(config.clone()).unwrap();
        assert_eq!(manager.size(), config.size);
    }

    #[test]
    fn test_umem_ref_counting() {
        let config = UmemConfig {
            size: 4096 * 64,
            hugepage: false,
            locked: false,
            shared: false,
        };
        let mut manager = UmemManager::new(config).unwrap();
        manager.create().unwrap();

        let ref1 = manager.incref().unwrap();
        assert_eq!(ref1, 2);

        let ref2 = manager.incref().unwrap();
        assert_eq!(ref2, 3);

        let dec1 = manager.decref();
        assert_eq!(dec1, 2);

        let dec2 = manager.decref();
        assert_eq!(dec2, 1);
    }

    #[test]
    fn test_umem_incref_before_create_fails() {
        // Fail-Closed：未 create() 时 incref 返回 Err（NotCreated），不 panic
        let config = UmemConfig {
            size: 4096 * 16,
            hugepage: false,
            locked: false,
            shared: false,
        };
        let manager = UmemManager::new(config).unwrap();
        let result = manager.incref();
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            LinuxError::Umem(UmemError::NotCreated)
        ));
    }

    #[test]
    fn test_umem_hugepage_align_overflow_fails() {
        // Fail-Closed：size 页对齐但接近 usize::MAX（= 2^64 - 4096），
        // HugePage 对齐时 size + 2MB - 1 溢出，必须返回错误而非 wrap
        let config = UmemConfig {
            size: usize::MAX - 4095,
            hugepage: true,
            locked: false,
            shared: false,
        };
        let mut manager = UmemManager::new(config).unwrap();
        let result = manager.create();
        assert!(result.is_err(), "HugePage 对齐溢出必须 Fail-Closed");
    }

    #[test]
    fn test_umem_config_clone() {
        let config = UmemConfig {
            size: 4096 * 32,
            hugepage: true,
            locked: true,
            shared: true,
        };
        let cloned = config.clone();
        assert_eq!(cloned.size, config.size);
        assert_eq!(cloned.hugepage, config.hugepage);
        assert_eq!(cloned.locked, config.locked);
        assert_eq!(cloned.shared, config.shared);
    }

    #[test]
    fn test_umem_manager_debug_format() {
        let config = UmemConfig {
            size: 4096 * 16,
            hugepage: false,
            locked: false,
            shared: false,
        };
        let manager = UmemManager::new(config).unwrap();
        let debug = format!("{:?}", manager);
        assert!(debug.contains("UmemManager"));
        assert!(debug.contains("initialized"));
        assert!(debug.contains("ref_count"));
    }

    #[test]
    fn test_umem_initialized_state() {
        let config = UmemConfig {
            size: 4096 * 32,
            hugepage: false,
            locked: false,
            shared: false,
        };
        let mut manager = UmemManager::new(config).unwrap();
        assert!(!manager.is_initialized());

        manager.create().unwrap();
        assert!(manager.is_initialized());
    }

    #[test]
    fn test_umem_large_allocation() {
        let config = UmemConfig {
            size: 4096 * 1024, // 4MB
            hugepage: false,
            locked: false,
            shared: false,
        };
        let mut manager = match UmemManager::new(config) {
            Ok(m) => m,
            Err(_) => return, // 跳过如果内存不足
        };

        let result = manager.create();
        if result.is_ok() {
            assert!(manager.is_initialized());
            assert!(!manager.as_ptr().is_null());
        }
    }

    #[test]
    fn test_umem_shared_config_flag() {
        let config = UmemConfig {
            size: 4096 * 32,
            hugepage: false,
            locked: false,
            shared: true,
        };
        let mut manager = UmemManager::new(config).unwrap();
        let result = manager.create();
        assert!(result.is_ok());
        assert!(manager.is_initialized());
    }

    #[test]
    fn test_umem_region_send_sync() {
        fn assert_send_sync<T: Send + Sync>() {}
        assert_send_sync::<UmemRegion>();
    }
}