zenith-net 0.1.0

//! UDP 双模式实现
//!
//! 支持两种 UDP 工作模式：
//! 1. 无状态模式：直接转发，无会话跟踪（高性能）
//! 2. 有状态模式：维护 UDP 会话（防火墙/NAT 场景）
//!
//! 设计约束：
//! - 预分配会话表，硬上限
//! - 单 Owner，零堆分配热路径
//! - O(1) 会话查找
//! - 超时自动回收

use crate::error::NetError;
use crate::packet::{IpVersion};
use crate::source_admission::IpAddr;

/// UDP 工作模式
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum UdpMode {
    /// 无状态（直接转发，零会话开销）
    Stateless,
    /// 有状态（跟踪会话，超时回收）
    Stateful,
}

/// UDP 会话（有状态模式）
#[derive(Debug, Clone, Copy)]
pub struct UdpSession {
    /// 源 IP
    pub src_ip: IpAddr,
    /// 源端口
    pub src_port: u16,
    /// 目标 IP
    pub dst_ip: IpAddr,
    /// 目标端口
    pub dst_port: u16,
    /// IP 版本
    pub ip_version: IpVersion,
    /// 创建时间戳
    pub created_at: u64,
    /// 最近活动时间戳
    pub last_active: u64,
    /// 入站包计数
    pub inbound_count: u64,
    /// 出站包计数
    pub outbound_count: u64,
    /// 是否活跃
    pub active: bool,
}

impl UdpSession {
    /// 创建空会话（未激活）
    #[inline]
    pub fn empty() -> Self {
        Self {
            src_ip: IpAddr::V4([0; 4]),
            src_port: 0,
            dst_ip: IpAddr::V4([0; 4]),
            dst_port: 0,
            ip_version: IpVersion::V4,
            created_at: 0,
            last_active: 0,
            inbound_count: 0,
            outbound_count: 0,
            active: false,
        }
    }

    /// 激活会话
    #[inline]
    pub fn activate(
        &mut self,
        src_ip: IpAddr,
        src_port: u16,
        dst_ip: IpAddr,
        dst_port: u16,
        ip_version: IpVersion,
        now: u64,
    ) {
        self.src_ip = src_ip;
        self.src_port = src_port;
        self.dst_ip = dst_ip;
        self.dst_port = dst_port;
        self.ip_version = ip_version;
        self.created_at = now;
        self.last_active = now;
        self.inbound_count = 0;
        self.outbound_count = 0;
        self.active = true;
    }

    /// 更新活动时间
    #[inline]
    pub fn touch(&mut self, now: u64) {
        self.last_active = now;
    }

    /// 记录入站包
    #[inline]
    pub fn record_inbound(&mut self) {
        self.inbound_count += 1;
    }

    /// 记录出站包
    #[inline]
    pub fn record_outbound(&mut self) {
        self.outbound_count += 1;
    }

    /// 是否超时
    #[inline]
    pub fn is_expired(&self, now: u64, timeout_ms: u64) -> bool {
        self.active && now.saturating_sub(self.last_active) > timeout_ms
    }

    /// 是否匹配四元组
    #[inline]
    pub fn matches(
        &self,
        src_ip: IpAddr,
        src_port: u16,
        dst_ip: IpAddr,
        dst_port: u16,
        ip_version: IpVersion,
    ) -> bool {
        self.active
            && self.src_ip == src_ip
            && self.src_port == src_port
            && self.dst_ip == dst_ip
            && self.dst_port == dst_port
            && self.ip_version == ip_version
    }

    /// 获取反向匹配键（响应包查找）
    #[inline]
    pub fn reverse_matches(
        &self,
        src_ip: IpAddr,
        src_port: u16,
        dst_ip: IpAddr,
        dst_port: u16,
        ip_version: IpVersion,
    ) -> bool {
        self.active
            && self.src_ip == dst_ip
            && self.src_port == dst_port
            && self.dst_ip == src_ip
            && self.dst_port == src_port
            && self.ip_version == ip_version
    }
}

/// UDP 统计信息
#[derive(Debug, Clone, Copy, Default)]
pub struct UdpStats {
    /// 总会话数
    pub total_sessions: u64,
    /// 活跃会话数
    pub active_sessions: u64,
    /// 入站包计数
    pub inbound_packets: u64,
    /// 出站包计数
    pub outbound_packets: u64,
    /// 会话超时回收数
    pub session_timeouts: u64,
    /// 拒绝包数
    pub rejected: u64,
}

/// 最大 UDP 会话数（硬上限）
pub const MAX_SESSIONS: usize = 32768;

/// 哈希桶数量（必须为 2 的幂）
pub const SESSION_BUCKETS: usize = 8192;

/// UDP 会话表（预分配，零堆分配）
///
/// # 与 [`crate::transport::tcp::ConnectionTable`] 的关系（泛型化评估结论）
/// 两者同为「预分配槽数组 + u32 桶开放寻址 + 线性探测」骨架，但此处
/// **有意保持单态化副本以保内联**，不做泛型化，原因：
/// - 插入语义不同：TCP 侧插入时会复用指向 CLOSED 槽的桶，本表无此逻辑；
/// - 槽类型、空闲判定（`!active` vs `state == Closed`）、键提取方式均不同，
///   泛型化只能以闭包/trait 钩子参数化，虽静态分发但阻碍编译期内联优化；
/// - 桶数常量不同（[`SESSION_BUCKETS`] = 8192 vs TCP 侧 16384），
///   且本表在无状态模式下容量为 0（TCP 侧最小为 1）。
///
/// 两侧共享的部分（IP 折叠哈希、五元组哈希）已提取至
/// [`crate::transport::hash_flow_tuple`]。
#[derive(Debug)]
pub struct UdpSessionTable {
    /// 模式
    mode: UdpMode,
    /// 会话槽数组（预分配）
    sessions: Vec<UdpSession>,
    /// 哈希桶（存储会话索引，0 表示空）
    buckets: Vec<u32>,
    /// 已使用会话数
    count: usize,
    /// 最大容量
    capacity: usize,
    /// 统计信息
    stats: UdpStats,
}

impl UdpSessionTable {
    /// 创建会话表
    ///
    /// # 参数
    /// * `mode` - UDP 工作模式
    /// * `capacity` - 最大会话数（硬上限，仅对有状态模式有效）
    pub fn new(mode: UdpMode, capacity: usize) -> Self {
        let cap = if mode == UdpMode::Stateful {
            capacity.clamp(1, MAX_SESSIONS)
        } else {
            0 // 无状态模式不需要会话表
        };
        let mut sessions = Vec::with_capacity(cap);
        for _ in 0..cap {
            sessions.push(UdpSession::empty());
        }
        let buckets = vec![0u32; SESSION_BUCKETS];
        Self {
            mode,
            sessions,
            buckets,
            count: 0,
            capacity: cap,
            stats: UdpStats::default(),
        }
    }

    /// 获取模式
    #[inline]
    pub fn mode(&self) -> UdpMode {
        self.mode
    }

    /// 获取统计信息
    #[inline]
    pub fn stats(&self) -> UdpStats {
        self.stats
    }

    /// 获取当前会话数
    #[inline]
    pub fn count(&self) -> usize {
        self.count
    }

    /// 计算哈希（使用折叠哈希）
    ///
    /// 实现委托给 transport 层共享的 [`crate::transport::hash_flow_tuple`]
    /// （与 TCP 连接表同源，禁止复制实现）。
    #[inline]
    fn compute_hash(
        &self,
        src_ip: IpAddr,
        src_port: u16,
        dst_ip: IpAddr,
        dst_port: u16,
        ip_version: IpVersion,
    ) -> u64 {
        super::hash_flow_tuple(&src_ip, src_port, &dst_ip, dst_port, ip_version)
    }

    /// 查找会话（O(1)）
    ///
    /// # 返回
    /// * `Option<usize>` - 会话索引
    #[inline]
    pub fn find_session(
        &self,
        src_ip: IpAddr,
        src_port: u16,
        dst_ip: IpAddr,
        dst_port: u16,
        ip_version: IpVersion,
    ) -> Option<usize> {
        if self.mode == UdpMode::Stateless {
            return None; // 无状态模式直接返回 None
        }

        let hash = self.compute_hash(src_ip, src_port, dst_ip, dst_port, ip_version);
        let bucket_idx = (hash as usize) & (SESSION_BUCKETS - 1);
        let mut i = bucket_idx;
        let mut probe = 0;

        loop {
            let raw = self.buckets[i];
            if raw == 0 {
                return None;
            }
            let real_idx = (raw as usize) - 1;
            let session = &self.sessions[real_idx];
            if session.matches(src_ip, src_port, dst_ip, dst_port, ip_version) {
                return Some(real_idx);
            }
            probe += 1;
            if probe >= SESSION_BUCKETS {
                return None;
            }
            i = (i + 1) & (SESSION_BUCKETS - 1);
        }
    }

    /// 查找反向会话（用于响应包路由）
    #[inline]
    pub fn find_reverse_session(
        &self,
        src_ip: IpAddr,
        src_port: u16,
        dst_ip: IpAddr,
        dst_port: u16,
        ip_version: IpVersion,
    ) -> Option<usize> {
        if self.mode == UdpMode::Stateless {
            return None;
        }

        // 反转四元组查找
        self.find_session(dst_ip, dst_port, src_ip, src_port, ip_version)
    }

    /// 创建新会话
    ///
    /// # 返回
    /// * `Result<usize, NetError>` - 会话索引
    pub fn create_session(
        &mut self,
        src_ip: IpAddr,
        src_port: u16,
        dst_ip: IpAddr,
        dst_port: u16,
        ip_version: IpVersion,
        now: u64,
    ) -> Result<usize, NetError> {
        // 无状态模式不应调用此方法
        if self.mode == UdpMode::Stateless {
            return Err(NetError::InvalidOperation {
                reason: "Cannot create session in stateless mode".to_string(),
            });
        }

        // 检查容量
        if self.count >= self.capacity {
            return Err(NetError::ConnectionTableFull {
                capacity: self.capacity,
            });
        }

        // 查找空闲槽（fail-closed：满表返回错误，绝不静默复用槽 0）
        let slot_idx = self.find_free_session().ok_or(NetError::ConnectionTableFull {
            capacity: self.capacity,
        })?;

        // 插入哈希表
        let hash = self.compute_hash(src_ip, src_port, dst_ip, dst_port, ip_version);
        let bucket_idx = (hash as usize) & (SESSION_BUCKETS - 1);
        let mut i = bucket_idx;
        let mut probe = 0;

        loop {
            if self.buckets[i] == 0 {
                self.buckets[i] = (slot_idx as u32) + 1;
                break;
            }
            probe += 1;
            if probe >= SESSION_BUCKETS {
                return Err(NetError::HashTableFull);
            }
            i = (i + 1) & (SESSION_BUCKETS - 1);
        }

        // 激活会话
        self.sessions[slot_idx].activate(
            src_ip, src_port, dst_ip, dst_port, ip_version, now,
        );
        self.count += 1;
        self.stats.total_sessions += 1;
        self.stats.active_sessions += 1;

        Ok(slot_idx)
    }

    /// 查找空闲会话槽
    ///
    /// # 返回
    /// * `Some(idx)` - 空闲槽索引
    /// * `None` - 表已满（调用方须 fail-closed，禁止静默复用槽 0）
    #[inline]
    fn find_free_session(&self) -> Option<usize> {
        (0..self.capacity).find(|&i| !self.sessions[i].active)
    }

    /// 获取会话引用
    #[inline]
    pub fn get_session(&self, idx: usize) -> Option<&UdpSession> {
        if idx < self.sessions.len() {
            Some(&self.sessions[idx])
        } else {
            None
        }
    }

    /// 获取会话可变引用
    #[inline]
    pub fn get_session_mut(&mut self, idx: usize) -> Option<&mut UdpSession> {
        if idx < self.sessions.len() {
            Some(&mut self.sessions[idx])
        } else {
            None
        }
    }

    /// 处理入站 UDP 包
    ///
    /// # 参数
    /// * `src_ip` - 源 IP
    /// * `src_port` - 源端口
    /// * `dst_ip` - 目标 IP
    /// * `dst_port` - 目标端口
    /// * `ip_version` - IP 版本
    /// * `now` - 当前时间戳
    ///
    /// # 返回
    /// * `UdpAction` - 应采取的动作
    pub fn handle_inbound(
        &mut self,
        src_ip: IpAddr,
        src_port: u16,
        dst_ip: IpAddr,
        dst_port: u16,
        ip_version: IpVersion,
        now: u64,
    ) -> UdpAction {
        self.stats.inbound_packets += 1;

        match self.mode {
            UdpMode::Stateless => UdpAction::Forward,
            UdpMode::Stateful => {
                // 查找现有会话
                if let Some(idx) = self.find_session(src_ip, src_port, dst_ip, dst_port, ip_version) {
                    if let Some(session) = self.get_session_mut(idx) {
                        session.touch(now);
                        session.record_inbound();
                    }
                    UdpAction::Forward
                } else {
                    // 创建新会话
                    match self.create_session(
                        src_ip, src_port, dst_ip, dst_port, ip_version, now,
                    ) {
                        Ok(idx) => {
                            // 新会话也记录入站
                            if let Some(session) = self.get_session_mut(idx) {
                                session.record_inbound();
                            }
                            UdpAction::Forward
                        }
                        Err(_) => {
                            self.stats.rejected += 1;
                            UdpAction::Drop
                        }
                    }
                }
            }
        }
    }

    /// 处理出站 UDP 包
    ///
    /// # 参数
    /// * `src_ip` - 源 IP
    /// * `src_port` - 源端口
    /// * `dst_ip` - 目标 IP
    /// * `dst_port` - 目标端口
    /// * `ip_version` - IP 版本
    /// * `now` - 当前时间戳
    ///
    /// # 返回
    /// * `UdpAction` - 应采取的动作
    pub fn handle_outbound(
        &mut self,
        src_ip: IpAddr,
        src_port: u16,
        dst_ip: IpAddr,
        dst_port: u16,
        ip_version: IpVersion,
        now: u64,
    ) -> UdpAction {
        self.stats.outbound_packets += 1;

        match self.mode {
            UdpMode::Stateless => UdpAction::Forward,
            UdpMode::Stateful => {
                // 查找反向会话（用于 NAT 映射）
                if let Some(idx) =
                    self.find_reverse_session(src_ip, src_port, dst_ip, dst_port, ip_version)
                {
                    if let Some(session) = self.get_session_mut(idx) {
                        session.touch(now);
                        session.record_outbound();
                    }
                    UdpAction::Forward
                } else {
                    // 没有对应会话，可能是新出站流
                    match self.create_session(
                        src_ip, src_port, dst_ip, dst_port, ip_version, now,
                    ) {
                        Ok(idx) => {
                            if let Some(session) = self.get_session_mut(idx) {
                                session.record_outbound();
                            }
                            UdpAction::Forward
                        }
                        Err(_) => {
                            self.stats.rejected += 1;
                            UdpAction::Drop
                        }
                    }
                }
            }
        }
    }

    /// 基于时间戳的超时扫描（兼容性方法）
    ///
    /// # Arguments
    /// * `now` - 当前时间戳
    /// * `timeout_ms` - 超时阈值
    pub fn scan_timeouts(&mut self, now: u64, timeout_ms: u64) -> usize {
        if self.mode == UdpMode::Stateless {
            return 0;
        }

        let mut expired_count = 0;
        for session in self.sessions.iter_mut() {
            if session.is_expired(now, timeout_ms) {
                session.active = false;
                expired_count += 1;
            }
        }

        if expired_count > 0 {
            self.count = self.count.saturating_sub(expired_count);
            self.stats.session_timeouts += expired_count as u64;
            self.stats.active_sessions = self.stats.active_sessions.saturating_sub(expired_count as u64);
            self.rebuild_hash();
        }

        expired_count
    }

    /// 重建哈希表
    fn rebuild_hash(&mut self) {
        self.buckets.fill(0);
        for i in 0..self.sessions.len() {
            if self.sessions[i].active {
                let s = &self.sessions[i];
                let hash =
                    self.compute_hash(s.src_ip, s.src_port, s.dst_ip, s.dst_port, s.ip_version);
                let bucket_idx = (hash as usize) & (SESSION_BUCKETS - 1);
                let mut idx = bucket_idx;
                let mut probe = 0;
                loop {
                    if self.buckets[idx] == 0 {
                        self.buckets[idx] = (i as u32) + 1;
                        break;
                    }
                    probe += 1;
                    if probe >= SESSION_BUCKETS {
                        break;
                    }
                    idx = (idx + 1) & (SESSION_BUCKETS - 1);
                }
            }
        }
    }

    /// 遍历活跃会话
    pub fn iter_active(&self) -> impl Iterator<Item = &UdpSession> {
        self.sessions.iter().filter(|s| s.active)
    }
}

/// UDP 动作
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum UdpAction {
    /// 直接转发
    Forward,
    /// 丢弃
    Drop,
}

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

    fn v4(ip: [u8; 4]) -> IpAddr {
        IpAddr::V4(ip)
    }

    #[test]
    fn test_udp_session_activate() {
        let mut session = UdpSession::empty();
        session.activate(v4([10, 0, 0, 1]), 8080, v4([10, 0, 0, 2]), 1234, IpVersion::V4, 1000);

        assert!(session.active);
        assert_eq!(session.src_ip, v4([10, 0, 0, 1]));
        assert_eq!(session.src_port, 8080);
        assert_eq!(session.dst_ip, v4([10, 0, 0, 2]));
        assert_eq!(session.dst_port, 1234);
        assert_eq!(session.created_at, 1000);
        assert_eq!(session.inbound_count, 0);
    }

    #[test]
    fn test_udp_session_matches() {
        let mut session = UdpSession::empty();
        session.activate(v4([10, 0, 0, 1]), 8080, v4([10, 0, 0, 2]), 1234, IpVersion::V4, 1000);

        // 精确匹配
        assert!(session.matches(v4([10, 0, 0, 1]), 8080, v4([10, 0, 0, 2]), 1234, IpVersion::V4));

        // 不匹配
        assert!(!session.matches(v4([10, 0, 0, 1]), 9090, v4([10, 0, 0, 2]), 1234, IpVersion::V4));

        // 反向匹配
        assert!(session.reverse_matches(v4([10, 0, 0, 2]), 1234, v4([10, 0, 0, 1]), 8080, IpVersion::V4));
    }

    #[test]
    fn test_udp_session_expiry() {
        let mut session = UdpSession::empty();
        session.activate(v4([10, 0, 0, 1]), 8080, v4([10, 0, 0, 2]), 1234, IpVersion::V4, 1000);

        // 未超时
        assert!(!session.is_expired(1500, 1000));

        // 超时
        assert!(session.is_expired(2100, 1000));

        // 未激活的会话永不超时
        let empty = UdpSession::empty();
        assert!(!empty.is_expired(99999, 0));
    }

    #[test]
    fn test_udp_session_table_stateless() {
        let mut table = UdpSessionTable::new(UdpMode::Stateless, 1024);

        assert_eq!(table.mode(), UdpMode::Stateless);

        // 无状态模式应直接 Forward
        let action = table.handle_inbound(
            v4([10, 0, 0, 1]),
            8080,
            v4([10, 0, 0, 2]),
            1234,
            IpVersion::V4,
            1000,
        );
        assert_eq!(action, UdpAction::Forward);
        assert_eq!(table.stats().inbound_packets, 1);

        // 查找会话应返回 None
        assert!(table.find_session(
            v4([10, 0, 0, 1]),
            8080,
            v4([10, 0, 0, 2]),
            1234,
            IpVersion::V4,
        ).is_none());
    }

    #[test]
    fn test_udp_session_table_stateful() {
        let mut table = UdpSessionTable::new(UdpMode::Stateful, 1024);

        assert_eq!(table.mode(), UdpMode::Stateful);

        // 第一次包：创建会话
        let action = table.handle_inbound(
            v4([10, 0, 0, 1]),
            8080,
            v4([10, 0, 0, 2]),
            1234,
            IpVersion::V4,
            1000,
        );
        assert_eq!(action, UdpAction::Forward);
        assert_eq!(table.count(), 1);
        assert_eq!(table.stats().total_sessions, 1);

        // 第二次相同包：找到已有会话
        let action2 = table.handle_inbound(
            v4([10, 0, 0, 1]),
            8080,
            v4([10, 0, 0, 2]),
            1234,
            IpVersion::V4,
            1001,
        );
        assert_eq!(action2, UdpAction::Forward);
        assert_eq!(table.count(), 1); // 不应增加

        // 验证会话信息已更新
        let idx = table.find_session(
            v4([10, 0, 0, 1]),
            8080,
            v4([10, 0, 0, 2]),
            1234,
            IpVersion::V4,
        ).unwrap();
        let session = table.get_session(idx).unwrap();
        assert_eq!(session.inbound_count, 2);
        assert_eq!(session.last_active, 1001);
    }

    #[test]
    fn test_udp_session_table_outbound() {
        let mut table = UdpSessionTable::new(UdpMode::Stateful, 1024);

        // 创建入站会话
        table.handle_inbound(
            v4([10, 0, 0, 1]),
            8080,
            v4([10, 0, 0, 2]),
            1234,
            IpVersion::V4,
            1000,
        );

        // 出站包（从会话目标返回）应找到反向会话
        let action = table.handle_outbound(
            v4([10, 0, 0, 2]),
            1234,
            v4([10, 0, 0, 1]),
            8080,
            IpVersion::V4,
            1002,
        );
        assert_eq!(action, UdpAction::Forward);

        // 验证出站计数
        let idx = table.find_session(
            v4([10, 0, 0, 1]),
            8080,
            v4([10, 0, 0, 2]),
            1234,
            IpVersion::V4,
        ).unwrap();
        let session = table.get_session(idx).unwrap();
        assert_eq!(session.outbound_count, 1);
    }

    #[test]
    fn test_udp_session_timeout_scan() {
        let mut table = UdpSessionTable::new(UdpMode::Stateful, 1024);

        // 创建会话
        table.handle_inbound(
            v4([10, 0, 0, 1]),
            8080,
            v4([10, 0, 0, 2]),
            1234,
            IpVersion::V4,
            1000,
        );
        table.handle_inbound(
            v4([10, 0, 0, 3]),
            9090,
            v4([10, 0, 0, 4]),
            5678,
            IpVersion::V4,
            1000,
        );

        assert_eq!(table.count(), 2);

        // 超时扫描
        let expired = table.scan_timeouts(2000, 500); // 500ms 超时
        assert_eq!(expired, 2);
        assert_eq!(table.count(), 0);
        assert_eq!(table.stats().session_timeouts, 2);
    }

    #[test]
    fn test_udp_session_table_full() {
        let mut table = UdpSessionTable::new(UdpMode::Stateful, 4);

        // 填满会话表
        for i in 0..4u16 {
            table.handle_inbound(
                v4([10, 0, 0, i as u8]),
                8000 + i,
                v4([10, 0, 0, 2]),
                1234,
                IpVersion::V4,
                1000 + i as u64,
            );
        }
        assert_eq!(table.count(), 4);

        // 第 5 个应被拒绝
        let action = table.handle_inbound(
            v4([10, 0, 0, 100]),
            9999,
            v4([10, 0, 0, 2]),
            1234,
            IpVersion::V4,
            2000,
        );
        assert_eq!(action, UdpAction::Drop);
        assert_eq!(table.stats().rejected, 1);
    }

    #[test]
    fn test_udp_session_iterator() {
        let mut table = UdpSessionTable::new(UdpMode::Stateful, 1024);

        table.handle_inbound(
            v4([10, 0, 0, 1]),
            8080,
            v4([10, 0, 0, 2]),
            1234,
            IpVersion::V4,
            1000,
        );
        table.handle_inbound(
            v4([10, 0, 0, 3]),
            9090,
            v4([10, 0, 0, 4]),
            5678,
            IpVersion::V4,
            1000,
        );

        let active: Vec<&UdpSession> = table.iter_active().collect();
        assert_eq!(active.len(), 2);
    }

    #[test]
    fn test_udp_session_empty() {
        let session = UdpSession::empty();
        assert!(!session.active);
        assert_eq!(session.inbound_count, 0);
        assert_eq!(session.outbound_count, 0);
        assert_eq!(session.created_at, 0);
        assert_eq!(session.last_active, 0);
        assert!(!session.is_expired(1000, 100));
    }

    #[test]
    fn test_udp_session_touch() {
        let mut session = UdpSession::empty();
        session.activate(v4([10, 0, 0, 1]), 8080, v4([10, 0, 0, 2]), 1234, IpVersion::V4, 1000);
        assert_eq!(session.last_active, 1000);

        session.touch(2000);
        assert_eq!(session.last_active, 2000);
    }

    #[test]
    fn test_udp_session_record_counts() {
        let mut session = UdpSession::empty();
        session.activate(v4([10, 0, 0, 1]), 8080, v4([10, 0, 0, 2]), 1234, IpVersion::V4, 1000);

        session.record_inbound();
        session.record_inbound();
        session.record_outbound();

        assert_eq!(session.inbound_count, 2);
        assert_eq!(session.outbound_count, 1);
    }

    #[test]
    fn test_udp_session_not_active_no_match() {
        let session = UdpSession::empty();
        assert!(!session.matches(v4([10, 0, 0, 1]), 8080, v4([10, 0, 0, 2]), 1234, IpVersion::V4));
        assert!(!session.reverse_matches(v4([10, 0, 0, 2]), 1234, v4([10, 0, 0, 1]), 8080, IpVersion::V4));
    }

    #[test]
    fn test_udp_session_ipv6() {
        let mut session = UdpSession::empty();
        let ip1 = IpAddr::V6([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]);
        let ip2 = IpAddr::V6([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2]);
        session.activate(ip1, 8080, ip2, 1234, IpVersion::V6, 1000);

        assert!(session.matches(ip1, 8080, ip2, 1234, IpVersion::V6));
        assert!(session.reverse_matches(ip2, 1234, ip1, 8080, IpVersion::V6));
        assert!(!session.matches(ip1, 8080, ip2, 1234, IpVersion::V4));
    }

    #[test]
    fn test_udp_session_table_mode() {
        let stateless = UdpSessionTable::new(UdpMode::Stateless, 1024);
        assert_eq!(stateless.mode(), UdpMode::Stateless);

        let stateful = UdpSessionTable::new(UdpMode::Stateful, 1024);
        assert_eq!(stateful.mode(), UdpMode::Stateful);
    }

    #[test]
    fn test_udp_session_table_find_reverse() {
        let mut table = UdpSessionTable::new(UdpMode::Stateful, 1024);

        table.handle_inbound(
            v4([10, 0, 0, 1]),
            8080,
            v4([10, 0, 0, 2]),
            1234,
            IpVersion::V4,
            1000,
        );

        let forward = table.find_session(
            v4([10, 0, 0, 1]),
            8080,
            v4([10, 0, 0, 2]),
            1234,
            IpVersion::V4,
        );
        assert!(forward.is_some());

        let reverse = table.find_reverse_session(
            v4([10, 0, 0, 2]),
            1234,
            v4([10, 0, 0, 1]),
            8080,
            IpVersion::V4,
        );
        assert!(reverse.is_some());
        assert_eq!(forward, reverse);
    }

    #[test]
    fn test_udp_session_table_get_out_of_bounds() {
        let table = UdpSessionTable::new(UdpMode::Stateful, 1024);
        assert!(table.get_session(1024).is_none());
        assert!(table.get_session(9999).is_none());
    }

    #[test]
    fn test_udp_session_table_get_mut_out_of_bounds() {
        let mut table = UdpSessionTable::new(UdpMode::Stateful, 1024);
        assert!(table.get_session_mut(1024).is_none());
        assert!(table.get_session_mut(9999).is_none());
    }

    #[test]
    fn test_udp_stats_default() {
        let stats = UdpStats::default();
        assert_eq!(stats.total_sessions, 0);
        assert_eq!(stats.active_sessions, 0);
        assert_eq!(stats.inbound_packets, 0);
        assert_eq!(stats.outbound_packets, 0);
        assert_eq!(stats.session_timeouts, 0);
        assert_eq!(stats.rejected, 0);
    }

    #[test]
    fn test_udp_session_table_stats_initial() {
        let table = UdpSessionTable::new(UdpMode::Stateful, 1024);
        let stats = table.stats();
        assert_eq!(stats.total_sessions, 0);
        assert_eq!(stats.active_sessions, 0);
    }

    #[test]
    fn test_udp_stateless_outbound() {
        let mut table = UdpSessionTable::new(UdpMode::Stateless, 1024);

        let action = table.handle_outbound(
            v4([10, 0, 0, 2]),
            1234,
            v4([10, 0, 0, 1]),
            8080,
            IpVersion::V4,
            1000,
        );
        assert_eq!(action, UdpAction::Forward);
        assert_eq!(table.stats().outbound_packets, 1);
    }

    #[test]
    fn test_udp_stateful_outbound_no_session() {
        let mut table = UdpSessionTable::new(UdpMode::Stateful, 1024);

        let action = table.handle_outbound(
            v4([10, 0, 0, 2]),
            1234,
            v4([10, 0, 0, 1]),
            8080,
            IpVersion::V4,
            1000,
        );
        assert_eq!(action, UdpAction::Forward);
        assert_eq!(table.count(), 1);
    }

    #[test]
    fn test_udp_session_table_zero_capacity() {
        let table = UdpSessionTable::new(UdpMode::Stateful, 0);
        assert_eq!(table.capacity, 1);
    }

    #[test]
    fn test_udp_session_timeout_boundary() {
        let mut session = UdpSession::empty();
        session.activate(v4([10, 0, 0, 1]), 8080, v4([10, 0, 0, 2]), 1234, IpVersion::V4, 1000);
        assert!(!session.is_expired(1500, 500));
        assert!(session.is_expired(1500, 499));
        assert!(session.is_expired(1501, 500));
        assert!(!session.is_expired(1499, 500));
        assert!(!session.is_expired(1000, 0));
        assert!(session.is_expired(1001, 0));
    }

    #[test]
    fn test_udp_session_table_count() {
        let mut table = UdpSessionTable::new(UdpMode::Stateful, 1024);
        assert_eq!(table.count(), 0);

        table.handle_inbound(
            v4([10, 0, 0, 1]),
            8080,
            v4([10, 0, 0, 2]),
            1234,
            IpVersion::V4,
            1000,
        );
        assert_eq!(table.count(), 1);

        table.handle_inbound(
            v4([10, 0, 0, 3]),
            9090,
            v4([10, 0, 0, 4]),
            5678,
            IpVersion::V4,
            1000,
        );
        assert_eq!(table.count(), 2);
    }
}