zenith-ebpf 0.1.0

//! 双 Bank 热更新管理器
//!
//! 核心功能：
//! 1. 双 Bank 状态管理
//! 2. 版本控制
//! 3. 原子切换逻辑（基于 bpf_link 的 Link::update_prog 实现零丢包）

use crate::error::ManagerError;
use crate::xdp_prog::{XdpLoader, XdpProgramName};
use std::sync::atomic::{AtomicU32, Ordering};

/// Bank 标识
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum BankId {
    /// Bank A（主用）
    A,
    /// Bank B（备用）
    B,
}

impl BankId {
    /// 将 `BankId` 转为 u32 编码（A=0, B=1），用于与 eBPF Map 通信
    pub fn as_u32(&self) -> u32 {
        match self {
            BankId::A => 0,
            BankId::B => 1,
        }
    }

    /// 从 u32 编码反解 `BankId`，未知值返回 `None`
    pub fn from_u32(id: u32) -> Option<Self> {
        match id {
            0 => Some(BankId::A),
            1 => Some(BankId::B),
            _ => None,
        }
    }

    /// 返回对侧 Bank（A ↔ B），用于双 Bank 切换时获取备用 Bank
    pub fn other(&self) -> BankId {
        match self {
            BankId::A => BankId::B,
            BankId::B => BankId::A,
        }
    }
}

impl std::fmt::Display for BankId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            BankId::A => write!(f, "BankA"),
            BankId::B => write!(f, "BankB"),
        }
    }
}

/// 单个 Bank 的状态
struct BankState {
    /// Bank 标识
    id: BankId,
    /// eBPF 程序加载器
    loader: Option<XdpLoader>,
    /// 用户态协议栈版本号
    user_version: u64,
    /// 是否为活动 Bank
    is_active: bool,
}

impl std::fmt::Debug for BankState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BankState")
            .field("id", &self.id)
            .field("user_version", &self.user_version)
            .field("is_active", &self.is_active)
            .field("has_loader", &self.loader.is_some())
            .finish()
    }
}

/// 双 Bank 管理器
///
/// 支持两层热更新：
/// 1. **内核态热更新**：通过 `Link::update_prog()` 原子替换 XDP 程序
/// 2. **用户态热更新**：通过原子替换 Rust 侧逻辑
///
/// # 设计约束
/// - 单 Owner：每个 Bank 的 loader 由管理器独占
/// - 热切换：通过 `bpf_link__update_program` 实现内核级零丢包替换
/// - 版本追踪：全局版本号 + 切换计数
pub struct DualBankManager {
    /// 当前活动的 Bank
    active_bank: BankState,
    /// 备用 Bank
    standby_bank: BankState,
    /// 全局版本号
    global_version: AtomicU32,
    /// 切换统计
    switch_count: AtomicU32,
}

impl DualBankManager {
    /// 创建双 Bank 管理器（兼容模式：任一 Bank 加载失败降级为 `None` 并 `tracing::warn!` 告警）
    ///
    /// # 行为
    /// - 非严格模式：A/B 任一或两者加载失败 → 构造函数仍返回 `Ok(Self)`，对应 Bank 的 loader 为 `None`
    /// - **降级不静默**：每次降级都产生 `tracing::warn!` 日志，且加载状态可由
    ///   [`Self::ebpf_loaded_banks`] 直接上报监控面板
    /// - 返回 `Ok` **绝不等于 eBPF 已就绪**：调用方必须显式检查
    ///   [`Self::is_ebpf_fully_loaded`]（双轨）或 [`Self::is_ebpf_active_bank_loaded`]（单轨），
    ///   否则 `switch` / `attach_xdp` 的失败会被推迟到运行时才暴露
    ///
    /// # 严格模式（推荐）
    /// 如要求 fail-closed（加载失败即拒绝启动），请使用 [`Self::new_strict`]。
    pub fn new() -> Result<Self, ManagerError> {
        // NET-016：默认 Bank 必须走 xdp_main 完整过滤链（L2/L3/L4 校验 + 协议
        // 白名单 + fail_closed），禁止默认加载绕过过滤链的 Redirect/Stats。
        let bank_a_loader = match XdpLoader::load(XdpProgramName::Main) {
            Ok(loader) => {
                tracing::info!("Bank A loaded successfully (xdp_main)");
                Some(loader)
            }
            Err(e) => {
                tracing::warn!(error = %e, "Failed to load Bank A");
                None
            }
        };

        let bank_b_loader = match XdpLoader::load(XdpProgramName::Main) {
            Ok(loader) => {
                tracing::info!("Bank B loaded successfully (xdp_main)");
                Some(loader)
            }
            Err(e) => {
                tracing::warn!(error = %e, "Failed to load Bank B");
                None
            }
        };

        Ok(Self {
            active_bank: BankState {
                id: BankId::A,
                loader: bank_a_loader,
                user_version: 1,
                is_active: true,
            },
            standby_bank: BankState {
                id: BankId::B,
                loader: bank_b_loader,
                user_version: 1,
                is_active: false,
            },
            global_version: AtomicU32::new(1),
            switch_count: AtomicU32::new(0),
        })
    }

    /// 创建双 Bank 管理器（严格 fail-closed 模式：任一 Bank 加载失败 → 返回错误）
    ///
    /// # Fail-Closed 语义
    /// 本函数保证：只要返回 `Ok(Self)`，两个 Bank 的 XDP 字节码 **均已通过**
    /// libbpf-rs 真实 `open_memory + load`；任一 Bank 加载失败 **立即** 返回
    /// `Err(ManagerError::Load)`，系统宁可拒绝启动，也绝不带着"半加载 eBPF"
    /// 运行——与 [`Self::new`] 的兼容降级模式严格互斥，不存在中间状态。
    ///
    /// 与 [`Self::new`] 最大区别：
    /// - 无任何降级；没有加载成功的 eBPF 程序 = ManagerError::Load
    /// - 调用方无需再检查 [`Self::is_ebpf_fully_loaded`]，
    ///   可直接在后续 `switch / attach_xdp` 中使用 loader
    ///
    /// # 错误
    /// - `ManagerError::Load`：Bank A 或 Bank B 的 `XdpLoader::load` 失败，错误信息包含具体 Bank 与 libbpf 原始描述
    pub fn new_strict() -> Result<Self, ManagerError> {
        // NET-016：默认 Bank 走 xdp_main 完整过滤链，禁止绕过过滤链的 Redirect/Stats
        let bank_a_loader =
            XdpLoader::load(XdpProgramName::Main).map_err(|e| {
                ManagerError::Load(format!("Bank A (xdp_main) load failed: {e}"))
            })?;
        let bank_b_loader = XdpLoader::load(XdpProgramName::Main).map_err(|e| {
            ManagerError::Load(format!("Bank B (xdp_main) load failed: {e}"))
        })?;

        Ok(Self {
            active_bank: BankState {
                id: BankId::A,
                loader: Some(bank_a_loader),
                user_version: 1,
                is_active: true,
            },
            standby_bank: BankState {
                id: BankId::B,
                loader: Some(bank_b_loader),
                user_version: 1,
                is_active: false,
            },
            global_version: AtomicU32::new(1),
            switch_count: AtomicU32::new(0),
        })
    }

    /// 活动/备用两个 Bank 的 XDP Loader 是否 **均已真实加载成功**
    ///
    /// 严格判定：`active_bank.loader.is_some() && standby_bank.loader.is_some()`。
    /// 用于区分"兼容模式构造成功但实际无 eBPF 内核态程序"的情形。
    #[inline]
    pub fn is_ebpf_fully_loaded(&self) -> bool {
        self.active_bank.loader.is_some() && self.standby_bank.loader.is_some()
    }

    /// 当前活动 Bank 的 XDP Loader 是否已真实加载成功
    ///
    /// 这是数据面"能否用当前 Bank 转发/统计"的最小充分条件：
    /// - `true`：活动 Bank 已持有真实通过 libbpf-rs 加载的 BPF object；可尝试 attach_xdp
    /// - `false`：活动 Bank 已在构造时被静默降级为 None；任何 kernel attach 操作均会失败
    #[inline]
    pub fn is_ebpf_active_bank_loaded(&self) -> bool {
        self.active_bank.loader.is_some()
    }

    /// 获取两个 Bank 的加载状态元组 `(A_loaded, B_loaded)`，供监控面板直接上报
    #[inline]
    pub fn ebpf_loaded_banks(&self) -> (bool, bool) {
        (
            matches!(self.active_bank.id, BankId::A) && self.active_bank.loader.is_some()
                || matches!(self.standby_bank.id, BankId::A) && self.standby_bank.loader.is_some(),
            matches!(self.active_bank.id, BankId::B) && self.active_bank.loader.is_some()
                || matches!(self.standby_bank.id, BankId::B) && self.standby_bank.loader.is_some(),
        )
    }

    /// 原子切换 Bank
    ///
    /// 通过 `Link::update_prog()` 实现内核级零丢包热切换：
    /// 1. 确认两个 Bank 都已加载
    /// 2. 若活动 Bank 已挂载 bpf_link，通过 `bpf_link__update_program`
    ///    将该 Link 原子指向备用 Bank 的 XDP 程序（内核级零丢包）
    /// 3. 交换用户态 Bank 状态，bpf_link 跟随新的活动 Bank
    ///
    /// # Fail-Closed
    /// 内核态原子替换失败时直接返回 `Err`，**不交换任何状态**，
    /// 旧程序保持在线，杜绝"用户态已切换但内核仍跑旧程序"的分裂状态。
    ///
    /// # 参数
    /// * `ifindex` - 网卡接口索引（仅用于日志；内核替换通过既有 bpf_link 完成）
    ///
    /// # 返回
    /// * `Result<BankId, ManagerError>` - 切换后的活动 Bank ID
    pub fn switch(&mut self, ifindex: i32) -> Result<BankId, ManagerError> {
        // 验证状态
        if self.active_bank.loader.is_none() || self.standby_bank.loader.is_none() {
            return Err(ManagerError::InvalidState);
        }

        // NET-016：默认 Bank 必须走 xdp_main 完整过滤链。switch 前校验活动程序
        // 为 xdp_main，否则 fail-closed——禁止切换到绕过过滤链的 Redirect/Stats。
        let active_prog = self
            .active_bank
            .loader
            .as_ref()
            .map(XdpLoader::program_name);
        if active_prog != Some(XdpProgramName::Main) {
            return Err(ManagerError::NotXdpMain(format!(
                "活动 Bank ({}) 程序为 {:?}，非 xdp_main（完整过滤链被绕过），拒绝切换",
                self.active_bank.id, active_prog
            )));
        }

        // 内核态原子替换：活动 Bank 已挂载 bpf_link 时，
        // 用 update_prog 将 Link 原子指向备用 Bank 的程序。
        // 失败时直接 Err，旧程序保持在线（fail-closed，禁止 let _ = 吞错）。
        let has_link = self
            .active_bank
            .loader
            .as_ref()
            .is_some_and(|l| l.active_link().is_some());
        if has_link {
            let active = self
                .active_bank
                .loader
                .as_mut()
                .ok_or(ManagerError::InvalidState)?;
            let standby = self
                .standby_bank
                .loader
                .as_mut()
                .ok_or(ManagerError::InvalidState)?;
            active.update_link_prog(standby).map_err(|e| {
                ManagerError::SwitchFailed(format!(
                    "bpf_link 原子替换失败（旧程序仍在线，未切换）: {e}"
                ))
            })?;
        } else {
            tracing::debug!(
                ifindex,
                "活动 Bank 未挂载 bpf_link，仅执行用户态 Bank 切换"
            );
        }

        // 交换 Bank 状态（用户态）
        let new_active_id = self.standby_bank.id;
        let old_active = std::mem::replace(
            &mut self.active_bank,
            BankState {
                id: self.standby_bank.id,
                loader: self.standby_bank.loader.take(),
                user_version: self.standby_bank.user_version,
                is_active: true,
            },
        );

        self.standby_bank = BankState {
            id: old_active.id,
            loader: old_active.loader,
            user_version: old_active.user_version,
            is_active: false,
        };

        // bpf_link 跟随活动 Bank：从原活动（现备用）转移至新活动 Bank，
        // 保证下次 switch 仍能对同一 Link 执行 update_prog。
        let link = self
            .standby_bank
            .loader
            .as_mut()
            .and_then(XdpLoader::take_link);
        if link.is_some()
            && let Some(active) = self.active_bank.loader.as_mut() {
            active.set_link(link);
        } else if link.is_none() {
            tracing::warn!(
                ifindex,
                "bpf_link 转移失败：备用 Bank 无可转移的 link（take_link 返回 None），\
                 下次 switch 将无法对同一 Link 执行 update_prog"
            );
        }

        // 更新统计
        self.switch_count.fetch_add(1, Ordering::SeqCst);

        Ok(new_active_id)
    }

    /// 获取当前活动 Bank ID
    pub fn active_bank_id(&self) -> BankId {
        self.active_bank.id
    }

    /// 获取备用 Bank ID
    pub fn standby_bank_id(&self) -> BankId {
        self.standby_bank.id
    }

    /// 获取活动 Bank 的用户态版本
    pub fn active_user_version(&self) -> u64 {
        self.active_bank.user_version
    }

    /// 获取备用 Bank 的用户态版本
    pub fn standby_user_version(&self) -> u64 {
        self.standby_bank.user_version
    }

    /// 获取切换次数
    pub fn switch_count(&self) -> u32 {
        self.switch_count.load(Ordering::SeqCst)
    }

    /// 获取全局版本号
    pub fn global_version(&self) -> u32 {
        self.global_version.load(Ordering::SeqCst)
    }

    /// 增加全局版本号
    ///
    /// 语义：`fetch_add` 返回旧值，`+ 1` 即新版本号（与后续 `global_version()`
    /// 读取一致）。u32 版本号在每次热更新 +1 的速率下回绕需 > 百年，
    /// 且回绕仅影响版本展示语义，不影响 bpf_link 原子替换正确性。
    pub fn increment_version(&self) -> u32 {
        self.global_version.fetch_add(1, Ordering::SeqCst) + 1
    }

    /// 获取活动 Bank 的 loader 引用
    pub fn active_loader(&self) -> Option<&XdpLoader> {
        self.active_bank.loader.as_ref()
    }

    /// 获取活动 Bank 的 loader 可变引用
    pub fn active_loader_mut(&mut self) -> Option<&mut XdpLoader> {
        self.active_bank.loader.as_mut()
    }

    /// 获取备用 Bank 的 loader 引用
    pub fn standby_loader(&self) -> Option<&XdpLoader> {
        self.standby_bank.loader.as_ref()
    }

    /// 获取备用 Bank 的 loader 可变引用
    pub fn standby_loader_mut(&mut self) -> Option<&mut XdpLoader> {
        self.standby_bank.loader.as_mut()
    }
}

impl std::fmt::Debug for DualBankManager {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DualBankManager")
            .field("active", &self.active_bank)
            .field("standby", &self.standby_bank)
            .field("switch_count", &self.switch_count)
            .field("global_version", &self.global_version)
            .finish()
    }
}

impl Drop for DualBankManager {
    fn drop(&mut self) {
        tracing::debug!("DualBankManager dropped");
    }
}

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

    #[test]
    fn test_bank_id_values() {
        assert_eq!(BankId::A.as_u32(), 0);
        assert_eq!(BankId::B.as_u32(), 1);
        assert_eq!(BankId::from_u32(0), Some(BankId::A));
        assert_eq!(BankId::from_u32(1), Some(BankId::B));
        assert_eq!(BankId::from_u32(2), None);
    }

    #[test]
    fn test_bank_id_other() {
        assert_eq!(BankId::A.other(), BankId::B);
        assert_eq!(BankId::B.other(), BankId::A);
    }

    #[test]
    fn test_bank_id_display() {
        assert_eq!(format!("{}", BankId::A), "BankA");
        assert_eq!(format!("{}", BankId::B), "BankB");
    }

    #[test]
    fn test_dual_bank_manager_creation() {
        let mgr = match DualBankManager::new() {
            Ok(mgr) => mgr,
            Err(e) => {
                // 构造本身理论上不返回 Err（兼容模式），防御性 skip 并注明原因
                eprintln!("SKIP: DualBankManager::new() 返回 Err: {e}");
                return;
            }
        };
        // 用户态状态字段为确定性断言（与内核能力无关）
        assert_eq!(mgr.active_bank_id(), BankId::A);
        assert_eq!(mgr.standby_bank_id(), BankId::B);
        assert_eq!(mgr.switch_count(), 0);
        assert_eq!(mgr.global_version(), 1);

        if mgr.is_ebpf_fully_loaded() {
            // 双轨真实加载成功：loaded banks 必须非空且版本计数语义正确
            let (a_loaded, b_loaded) = mgr.ebpf_loaded_banks();
            assert!(a_loaded, "is_ebpf_fully_loaded 为真时 Bank A loader 必须存在");
            assert!(b_loaded, "is_ebpf_fully_loaded 为真时 Bank B loader 必须存在");
            assert!(mgr.active_loader().is_some());
            assert!(mgr.standby_loader().is_some());
            assert_eq!(mgr.active_user_version(), 1);
            assert_eq!(mgr.standby_user_version(), 1);
        } else {
            // 环境无 eBPF 权限（CI/容器内 libbpf 加载被内核拒绝）时，
            // 加载分支按兼容模式降级为 None：明确注明 skip，绝不用恒真断言掩饰
            eprintln!(
                "SKIP: eBPF 字节码加载失败（权限/内核不支持），loaded banks = {:?}，仅校验用户态状态",
                mgr.ebpf_loaded_banks()
            );
        }
    }

    #[test]
    fn test_dual_bank_manager_versioning() {
        let result = DualBankManager::new();
        if let Ok(mgr) = result {
            let v1 = mgr.global_version();
            let v2 = mgr.increment_version();
            assert!(v2 > v1);
        }
    }

    #[test]
    fn test_bank_id_debug() {
        let debug_str = format!("{:?}", BankId::A);
        assert!(!debug_str.is_empty());
        let debug_str = format!("{:?}", BankId::B);
        assert!(!debug_str.is_empty());
    }

    #[test]
    fn test_bank_id_clone_copy() {
        let a = BankId::A;
        let b = a;
        assert_eq!(a, b);
        assert_eq!(a.as_u32(), b.as_u32());
    }

    #[test]
    fn test_bank_id_unique() {
        assert_ne!(BankId::A, BankId::B);
        assert_eq!(BankId::A, BankId::A);
        assert_eq!(BankId::B, BankId::B);
    }

    #[test]
    fn test_bank_id_eq() {
        assert_eq!(BankId::A, BankId::A);
        assert_eq!(BankId::B, BankId::B);
        assert_ne!(BankId::A, BankId::B);
    }

    #[test]
    fn test_bank_id_from_u32_boundary() {
        assert_eq!(BankId::from_u32(0), Some(BankId::A));
        assert_eq!(BankId::from_u32(1), Some(BankId::B));
        assert_eq!(BankId::from_u32(u32::MAX), None);
        assert_eq!(BankId::from_u32(2), None);
        assert_eq!(BankId::from_u32(100), None);
    }

    #[test]
    fn test_bank_id_other_involution() {
        assert_eq!(BankId::A.other().other(), BankId::A);
        assert_eq!(BankId::B.other().other(), BankId::B);
    }

    #[test]
    fn test_dual_bank_manager_debug() {
        if let Ok(mgr) = DualBankManager::new() {
            let debug_str = format!("{:?}", mgr);
            assert!(!debug_str.is_empty());
            assert!(debug_str.contains("DualBankManager"));
        }
    }

    #[test]
    fn test_dual_bank_manager_initial_versions() {
        if let Ok(mgr) = DualBankManager::new() {
            assert_eq!(mgr.active_user_version(), 1);
            assert_eq!(mgr.standby_user_version(), 1);
            assert_eq!(mgr.global_version(), 1);
        }
    }

    #[test]
    fn test_dual_bank_manager_increment_version_multiple() {
        if let Ok(mgr) = DualBankManager::new() {
            let v1 = mgr.global_version();
            let v2 = mgr.increment_version();
            let v3 = mgr.increment_version();
            let v4 = mgr.increment_version();
            assert_eq!(v2, v1 + 1);
            assert_eq!(v3, v2 + 1);
            assert_eq!(v4, v3 + 1);
        }
    }

    #[test]
    fn test_dual_bank_manager_bank_ids_consistent() {
        if let Ok(mgr) = DualBankManager::new() {
            let active = mgr.active_bank_id();
            let standby = mgr.standby_bank_id();
            assert_ne!(active, standby);
            assert_eq!(active.other(), standby);
            assert_eq!(standby.other(), active);
        }
    }

    #[test]
    fn test_manager_error_display() {
        use crate::error::ManagerError;
        let e = ManagerError::InvalidState;
        assert!(format!("{}", e).contains("Invalid bank state"));
        let e = ManagerError::Load("test".to_string());
        assert!(format!("{}", e).contains("Load error"));
    }

    /// NET-016：活动 Bank 默认必须走 xdp_main 完整过滤链（若双轨加载成功）
    #[test]
    fn test_bank_active_program_is_xdp_main() {
        let mgr = match DualBankManager::new() {
            Ok(mgr) => mgr,
            Err(e) => {
                eprintln!("SKIP: DualBankManager::new() 返回 Err: {e}");
                return;
            }
        };
        if mgr.is_ebpf_fully_loaded() {
            let active = mgr
                .active_loader()
                .expect("双轨加载成功时活动 loader 必须存在");
            assert_eq!(
                active.program_name(),
                XdpProgramName::Main,
                "NET-016：默认活动 Bank 必须为 xdp_main"
            );
            assert!(
                active.is_full_filter_chain(),
                "NET-016：活动程序必须是完整过滤链 xdp_main"
            );
        } else {
            eprintln!(
                "SKIP: eBPF 未真实加载（{}/{}），无法校验活动程序为 xdp_main",
                mgr.ebpf_loaded_banks().0,
                mgr.ebpf_loaded_banks().1
            );
        }
    }
}