zenith-net 0.1.0

//! 协议预期配置 — 非预期包丢弃策略
//!
//! 在 Worker 数据面准入放行后执行二次校验：检查数据包是否符合预期协议/端口/
//! TTL/分片策略。不符合预期的包静默丢弃（回收至 Completion Ring），避免
//! 非预期流量消耗用户态处理资源。
//!
//! 本模块为 zenith-net 内部独立定义，不依赖 zenith-fingerprint crate，
//! 保证 zenith-net 的依赖树最小化。

use smallvec::SmallVec;

/// 协议预期配置 — 非预期包丢弃策略
#[derive(Debug, Clone)]
pub struct ProtocolExpectation {
    /// 允许的 TCP 端口集 (空 = 全部允许)
    pub tcp_allowed_ports: SmallVec<[u16; 16]>,
    /// 允许的 UDP 端口集
    pub udp_allowed_ports: SmallVec<[u16; 16]>,
    /// 是否丢弃无连接的非 SYN TCP 包
    pub drop_stray_tcp: bool,
    /// 是否丢弃无 QUIC 连接的 UDP 包
    pub drop_stray_udp: bool,
    /// 允许的 IP 协议号位图（u64，与 eBPF 侧 ExpectationConfig 一致）
    pub allowed_protocols: u64,
    /// 分片策略: 0=不限制 1=丢弃重叠分片(不支持,按 fail-closed 拒绝) 2=丢弃所有分片
    pub fragment_policy: u8,
    /// 最小 TTL (低于此值的包丢弃, 0=不检查)
    pub min_ttl: u8,
}

impl Default for ProtocolExpectation {
    fn default() -> Self {
        Self {
            tcp_allowed_ports: SmallVec::new(),
            udp_allowed_ports: SmallVec::new(),
            drop_stray_tcp: false,
            drop_stray_udp: false,
            // NET-013：统一白名单默认值（TCP/UDP/ICMP/ICMPv6），单一权威来源
            allowed_protocols: zenith_foundation::net::DEFAULT_PROTO_WHITELIST,
            fragment_policy: 0,
            min_ttl: 0,
        }
    }
}

impl ProtocolExpectation {
    /// 判断数据包是否符合协议预期
    ///
    /// # 参数
    /// * `proto` - IP 协议号 (6=TCP, 17=UDP, 1=ICMP, ...)
    /// * `dst_port` - 目的端口 (TCP/UDP 有效，其他协议为 0)
    /// * `ttl` - IP TTL 值
    /// * `is_fragment` - 是否为 IP 分片
    ///
    /// # 返回
    /// `true` = 符合预期（放行）；`false` = 不符合预期（丢弃）
    pub fn is_packet_expected(&self, proto: u8, dst_port: u16, ttl: u8, is_fragment: bool) -> bool {
        // 协议号白名单检查（使用 u64 位图，与 eBPF 侧 ExpectationConfig 一致）
        if proto < 64 {
            let bit = 1u64 << proto;
            if (self.allowed_protocols & bit) == 0 {
                return false;
            }
        } else {
            // proto >= 64 在 IP 协议号中未定义，默认拒绝
            return false;
        }
        // TCP 端口白名单
        if !self.tcp_allowed_ports.is_empty() && proto == 6 && !self.tcp_allowed_ports.contains(&dst_port) {
            return false;
        }
        // UDP 端口白名单
        if !self.udp_allowed_ports.is_empty() && proto == 17 && !self.udp_allowed_ports.contains(&dst_port) {
            return false;
        }
        // 最小 TTL 检查
        if self.min_ttl > 0 && ttl < self.min_ttl {
            return false;
        }
        // 分片策略：NET-010 修复。策略 1（丢弃重叠分片）未实现，按 fail-closed
        // 处理为与策略 2 相同（丢弃所有分片），绝不静默放行重叠分片。
        if is_fragment && (self.fragment_policy == 1 || self.fragment_policy == 2) {
            return false;
        }
        true
    }

    /// 校验配置合法性（fail-closed，NET-010）
    ///
    /// 分片策略 1（丢弃重叠分片）内核态与用户态均未实现有状态重叠检测，
    /// 为避免静默放行重叠分片，直接拒绝该配置值。
    ///
    /// # 错误
    /// * `NetError::InvalidOperation`：`fragment_policy == 1`（不支持的策略）
    pub fn validate(&self) -> std::result::Result<(), crate::NetError> {
        if self.fragment_policy == 1 {
            return Err(crate::NetError::InvalidOperation {
                reason: "fragment_policy=1 (丢弃重叠分片) 未实现，fail-closed 拒绝。\
                         请使用 0（不限制）或 2（丢弃所有分片）"
                    .to_string(),
            });
        }
        Ok(())
    }
}

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

    #[test]
    fn test_default_allows_tcp_udp() {
        let exp = ProtocolExpectation::default();
        // TCP to any port
        assert!(exp.is_packet_expected(6, 443, 64, false));
        // UDP to any port
        assert!(exp.is_packet_expected(17, 53, 64, false));
        // NET-013：默认白名单统一为 TCP/UDP/ICMP/ICMPv6，ICMP 现在默认放行
        assert!(exp.is_packet_expected(1, 0, 64, false));
    }

    #[test]
    fn test_tcp_port_whitelist() {
        let mut exp = ProtocolExpectation::default();
        exp.tcp_allowed_ports.push(80);
        exp.tcp_allowed_ports.push(443);
        // Allowed port
        assert!(exp.is_packet_expected(6, 80, 64, false));
        assert!(exp.is_packet_expected(6, 443, 64, false));
        // Non-allowed port
        assert!(!exp.is_packet_expected(6, 8080, 64, false));
        // Empty port list means all allowed (but we pushed ports, so non-empty)
    }

    #[test]
    fn test_udp_port_whitelist() {
        let mut exp = ProtocolExpectation::default();
        exp.udp_allowed_ports.push(53);
        // Allowed
        assert!(exp.is_packet_expected(17, 53, 64, false));
        // Not allowed
        assert!(!exp.is_packet_expected(17, 123, 64, false));
    }

    #[test]
    fn test_min_ttl() {
        let mut exp = ProtocolExpectation::default();
        exp.min_ttl = 10;
        // TTL above threshold
        assert!(exp.is_packet_expected(6, 80, 64, false));
        // TTL at threshold
        assert!(exp.is_packet_expected(6, 80, 10, false));
        // TTL below threshold
        assert!(!exp.is_packet_expected(6, 80, 5, false));
    }

    #[test]
    fn test_fragment_policy_drop_all() {
        let mut exp = ProtocolExpectation::default();
        exp.fragment_policy = 2;
        // Non-fragment OK
        assert!(exp.is_packet_expected(6, 80, 64, false));
        // Fragment dropped
        assert!(!exp.is_packet_expected(6, 80, 64, true));
    }

    #[test]
    fn test_fragment_policy_allow() {
        let exp = ProtocolExpectation::default();
        // Default policy = 0 (no restriction)
        assert!(exp.is_packet_expected(6, 80, 64, true));
        assert!(exp.is_packet_expected(6, 80, 64, false));
    }

    #[test]
    fn test_protocol_not_in_bitmap() {
        let exp = ProtocolExpectation::default();
        // Protocol 2 (IGMP) < 64 and not in default bitmap (TCP+UDP only)
        assert!(!exp.is_packet_expected(2, 0, 64, false));
        // Protocol 47 (GRE) >= 32 也不再绕过位图检查：不在默认位图中 → 拒绝
        assert!(!exp.is_packet_expected(47, 0, 64, false));
    }

    #[test]
    fn test_empty_port_lists_allow_all() {
        let exp = ProtocolExpectation::default();
        // Empty tcp_allowed_ports = all TCP ports allowed
        assert!(exp.is_packet_expected(6, 1, 64, false));
        assert!(exp.is_packet_expected(6, 65535, 64, false));
        // Empty udp_allowed_ports = all UDP ports allowed
        assert!(exp.is_packet_expected(17, 1, 64, false));
        assert!(exp.is_packet_expected(17, 65535, 64, false));
    }

    #[test]
    fn test_fragment_policy_1_fail_closed_runtime() {
        // NET-010：策略 1 未实现，运行时按 fail-closed 处理（丢弃所有分片）
        let mut exp = ProtocolExpectation::default();
        exp.fragment_policy = 1;
        // 非分片放行
        assert!(exp.is_packet_expected(6, 80, 64, false));
        // 分片丢弃（fail-closed，绝不静默放行重叠分片）
        assert!(!exp.is_packet_expected(6, 80, 64, true));
    }

    #[test]
    fn test_fragment_policy_1_rejected_by_validate() {
        // NET-010：配置校验必须拒绝策略 1
        let mut exp = ProtocolExpectation::default();
        exp.fragment_policy = 1;
        assert!(exp.validate().is_err());
        // 合法策略 0 与 2 通过校验
        exp.fragment_policy = 0;
        assert!(exp.validate().is_ok());
        exp.fragment_policy = 2;
        assert!(exp.validate().is_ok());
    }
}