bext_waf/

lib.rs

//! `bext-waf` — Web Application Firewall for the bext server.
//!
//! Provides IP filtering (CIDR), geo-blocking, request inspection (SQLi/XSS/traversal/scanner),
//! bot detection, DDoS mitigation, enhanced rate limiting, and audit logging.
//!
//! All regex patterns are compiled once via `OnceLock` — zero per-request compilation cost.

pub mod audit;
pub mod bot;
pub mod ddos;
pub mod geo;
pub mod ip_filter;
pub mod rate_limit;
pub mod rules;

use std::collections::HashMap;
use std::net::IpAddr;

use chrono::Utc;
use serde::{Deserialize, Serialize};

// Re-export key types.
pub use audit::{WafAuditLog, WafAuditStats, WafEvent};
pub use bot::{BotConfig, BotDetector, BotMode};
pub use ddos::{DdosConfig, DdosGuard};
pub use geo::{GeoBlocker, GeoConfig, GeoMode};
pub use ip_filter::{IpFilter, IpFilterConfig, IpFilterMode};
pub use rate_limit::{EnhancedRateLimiter, RateLimitRule};
pub use rules::custom::{CustomRule, CustomRuleAction, MatchConfig};
pub use rules::{RuleConfig, RuleEngine};

/// The decision made by the WAF for a given request.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum WafDecision {
    /// Allow the request to proceed.
    Allow,
    /// Block the request.
    Block {
        status: u16,
        reason: String,
        rule: String,
    },
    /// Rate-limit the request.
    RateLimit { retry_after: u64 },
    /// Issue a challenge (e.g. JS challenge for bot detection).
    Challenge { html: String },
}

/// A protocol-agnostic representation of an HTTP request.
#[derive(Debug, Clone)]
pub struct WafRequest {
    pub client_ip: IpAddr,
    pub method: String,
    pub path: String,
    pub query: Option<String>,
    pub headers: HashMap<String, String>,
    pub body: Option<String>,
    pub user_agent: Option<String>,
}

/// WAF statistics.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct WafStats {
    pub total_requests: u64,
    pub allowed: u64,
    pub blocked: u64,
    pub rate_limited: u64,
    pub challenged: u64,
    pub audit: WafAuditStats,
}

/// Full WAF configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WafConfig {
    /// Whether the WAF is enabled.
    #[serde(default = "default_true")]
    pub enabled: bool,
    #[serde(default)]
    pub ip_filter: IpFilterConfig,
    #[serde(default)]
    pub geo: GeoConfig,
    #[serde(default)]
    pub rules: RuleConfig,
    #[serde(default)]
    pub custom_rules: Vec<CustomRule>,
    #[serde(default)]
    pub bot: BotConfig,
    #[serde(default)]
    pub ddos: DdosConfig,
    #[serde(default)]
    pub rate_limit_rules: Vec<RateLimitRule>,
    /// Path to MaxMind GeoLite2-Country.mmdb file.
    #[serde(default)]
    pub geoip_db_path: Option<String>,
}

fn default_true() -> bool {
    true
}

impl Default for WafConfig {
    fn default() -> Self {
        Self {
            enabled: true,
            ip_filter: IpFilterConfig::default(),
            geo: GeoConfig::default(),
            rules: RuleConfig::default(),
            custom_rules: Vec::new(),
            bot: BotConfig::default(),
            ddos: DdosConfig::default(),
            rate_limit_rules: Vec::new(),
            geoip_db_path: None,
        }
    }
}

/// The main WAF engine composing all sub-systems.
pub struct WafEngine {
    enabled: bool,
    ip_filter: IpFilter,
    geo_blocker: GeoBlocker,
    rule_engine: RuleEngine,
    bot_detector: BotDetector,
    ddos_guard: DdosGuard,
    rate_limiter: EnhancedRateLimiter,
    audit_log: WafAuditLog,
    stats: parking_lot::Mutex<WafStats>,
}

impl WafEngine {
    /// Create a new WAF engine from configuration.
    pub fn new(config: WafConfig) -> anyhow::Result<Self> {
        let geo_blocker = match config.geoip_db_path {
            Some(ref path) if config.geo.enabled => GeoBlocker::new(config.geo.clone(), path)?,
            _ => GeoBlocker::disabled(),
        };

        Ok(Self {
            enabled: config.enabled,
            ip_filter: IpFilter::new(config.ip_filter),
            geo_blocker,
            rule_engine: RuleEngine::new(config.rules, config.custom_rules),
            bot_detector: BotDetector::new(config.bot),
            ddos_guard: DdosGuard::new(config.ddos),
            rate_limiter: EnhancedRateLimiter::new(config.rate_limit_rules),
            audit_log: WafAuditLog::new(),
            stats: parking_lot::Mutex::new(WafStats::default()),
        })
    }

    /// Check a request through all WAF layers.
    /// Returns the first non-Allow decision, or `WafDecision::Allow`.
    pub fn check(&self, req: &WafRequest) -> WafDecision {
        self.check_impl(req, false).0
    }

    /// Check a request and also return rate-limit response headers if applicable.
    pub fn check_with_headers(&self, req: &WafRequest) -> (WafDecision, Vec<(String, String)>) {
        self.check_impl(req, true)
    }

    /// Shared implementation for `check()` and `check_with_headers()`.
    fn check_impl(
        &self,
        req: &WafRequest,
        return_headers: bool,
    ) -> (WafDecision, Vec<(String, String)>) {
        if !self.enabled {
            return (WafDecision::Allow, vec![]);
        }

        {
            let mut s = self.stats.lock();
            s.total_requests += 1;
        }

        // 1. IP filter (fastest check -- pure network match).
        if let Some(decision) = self.ip_filter.check(req.client_ip) {
            self.record_decision(req, &decision);
            return (decision, vec![]);
        }

        // 2. Geo-blocking (uses real_ip_header if configured).
        if let Some(decision) = self.geo_blocker.check_request(req) {
            self.record_decision(req, &decision);
            return (decision, vec![]);
        }

        // 3. DDoS guard (connection limits, body/header size).
        if let Some(decision) = self.ddos_guard.check(req) {
            self.record_decision(req, &decision);
            return (decision, vec![]);
        }

        // 4. Rate limiting.
        if let Some((decision, headers)) = self.rate_limiter.check(req) {
            self.record_decision(req, &decision);
            let hdrs = if return_headers { headers } else { vec![] };
            return (decision, hdrs);
        }

        // 5. Bot detection.
        if let Some(decision) = self.bot_detector.check(req) {
            self.record_decision(req, &decision);
            return (decision, vec![]);
        }

        // 6. Rule engine (SQLi, XSS, traversal, shell injection, protocol violations, scanner, custom rules).
        if let Some(decision) = self.rule_engine.inspect(req) {
            self.record_decision(req, &decision);
            return (decision, vec![]);
        }

        // All checks passed.
        {
            let mut s = self.stats.lock();
            s.allowed += 1;
        }
        (WafDecision::Allow, vec![])
    }

    /// Record a DDoS connection start.
    pub fn record_connection(&self, ip: IpAddr) {
        self.ddos_guard.record_connection(ip);
    }

    /// Record a DDoS connection end.
    pub fn release_connection(&self, ip: IpAddr) {
        self.ddos_guard.release_connection(ip);
    }

    /// Get current WAF statistics.
    pub fn stats(&self) -> WafStats {
        let mut s = self.stats.lock().clone();
        s.audit = self.audit_log.stats();
        s
    }

    /// Get the audit log.
    pub fn audit_log(&self) -> &WafAuditLog {
        &self.audit_log
    }

    /// Get recent audit events.
    pub fn recent_events(&self, count: usize) -> Vec<WafEvent> {
        self.audit_log.recent(count)
    }

    /// Get Prometheus metrics.
    pub fn prometheus_metrics(&self) -> String {
        let s = self.stats.lock();
        let mut out = self.audit_log.format_prometheus();

        out.push_str("# HELP waf_requests_total Total requests processed\n");
        out.push_str("# TYPE waf_requests_total counter\n");
        out.push_str(&format!("waf_requests_total {}\n", s.total_requests));

        out.push_str("# HELP waf_requests_allowed Total requests allowed\n");
        out.push_str("# TYPE waf_requests_allowed counter\n");
        out.push_str(&format!("waf_requests_allowed {}\n", s.allowed));

        out
    }

    /// Hot-reload the IP filter configuration.
    pub fn reload_ip_filter(&self, config: IpFilterConfig) {
        self.ip_filter.reload(config);
    }

    /// Periodically clean up stale state (connection tracking, rate-limit buckets).
    pub fn cleanup(&self) {
        self.ddos_guard.cleanup(300); // 5 minutes
        self.rate_limiter
            .cleanup(std::time::Duration::from_secs(600)); // 10 minutes
    }

    fn record_decision(&self, req: &WafRequest, decision: &WafDecision) {
        let (action, rule, reason) = match decision {
            WafDecision::Allow => return,
            WafDecision::Block { rule, reason, .. } => ("block", rule.as_str(), reason.as_str()),
            WafDecision::RateLimit { retry_after: _ } => {
                let mut s = self.stats.lock();
                s.rate_limited += 1;
                ("rate_limit", "rate_limit", "")
            }
            WafDecision::Challenge { .. } => {
                let mut s = self.stats.lock();
                s.challenged += 1;
                ("challenge", "bot_detection", "JS challenge issued")
            }
        };

        // Update block stats (rate_limit and challenge already updated above).
        if action == "block" {
            let mut s = self.stats.lock();
            s.blocked += 1;
        }

        self.audit_log.record(WafEvent {
            timestamp: Utc::now(),
            client_ip: req.client_ip,
            path: req.path.clone(),
            rule: rule.to_string(),
            action: action.to_string(),
            details: reason.to_string(),
            request_id: None,
        });
    }
}

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

    fn make_req(ip: &str, method: &str, path: &str) -> WafRequest {
        WafRequest {
            client_ip: ip.parse().unwrap(),
            method: method.into(),
            path: path.into(),
            query: None,
            headers: {
                let mut h = HashMap::new();
                h.insert("Accept".into(), "text/html".into());
                h.insert("Accept-Language".into(), "en-US".into());
                h.insert("Accept-Encoding".into(), "gzip".into());
                h
            },
            body: None,
            user_agent: Some("Mozilla/5.0 Chrome/120.0".into()),
        }
    }

    fn make_req_minimal(ip: &str, path: &str) -> WafRequest {
        WafRequest {
            client_ip: ip.parse().unwrap(),
            method: "GET".into(),
            path: path.into(),
            query: None,
            headers: HashMap::new(),
            body: None,
            user_agent: None,
        }
    }

    #[test]
    fn disabled_engine_allows_all() {
        let config = WafConfig {
            enabled: false,
            ..Default::default()
        };
        let engine = WafEngine::new(config).unwrap();
        let req = make_req_minimal("10.0.0.1", "/../../../etc/passwd");
        assert_eq!(engine.check(&req), WafDecision::Allow);
    }

    #[test]
    fn default_config_allows_clean_request() {
        let engine = WafEngine::new(WafConfig::default()).unwrap();
        let req = make_req("10.0.0.1", "GET", "/api/users");
        assert_eq!(engine.check(&req), WafDecision::Allow);
    }

    #[test]
    fn blocks_sqli_in_path() {
        let engine = WafEngine::new(WafConfig::default()).unwrap();
        let mut req = make_req("10.0.0.1", "GET", "/search");
        req.query = Some("q=1 UNION SELECT * FROM users".into());
        match engine.check(&req) {
            WafDecision::Block { rule, .. } => assert_eq!(rule, "sql_injection"),
            other => panic!("expected Block, got {other:?}"),
        }
    }

    #[test]
    fn blocks_xss_in_body() {
        let engine = WafEngine::new(WafConfig::default()).unwrap();
        let mut req = make_req("10.0.0.1", "POST", "/comment");
        req.body = Some("<script>alert(1)</script>".into());
        match engine.check(&req) {
            WafDecision::Block { rule, .. } => assert_eq!(rule, "xss"),
            other => panic!("expected Block, got {other:?}"),
        }
    }

    #[test]
    fn blocks_traversal() {
        let engine = WafEngine::new(WafConfig::default()).unwrap();
        let req = make_req("10.0.0.1", "GET", "/static/../../etc/passwd");
        match engine.check(&req) {
            WafDecision::Block { rule, .. } => assert_eq!(rule, "path_traversal"),
            other => panic!("expected Block, got {other:?}"),
        }
    }

    #[test]
    fn blocks_scanner_ua() {
        let engine = WafEngine::new(WafConfig::default()).unwrap();
        let mut req = make_req("10.0.0.1", "GET", "/");
        req.user_agent = Some("sqlmap/1.5".into());
        match engine.check(&req) {
            WafDecision::Block { rule, .. } => assert_eq!(rule, "scanner_detection"),
            other => panic!("expected Block, got {other:?}"),
        }
    }

    #[test]
    fn ip_deny_blocks() {
        let config = WafConfig {
            ip_filter: IpFilterConfig {
                mode: IpFilterMode::Deny,
                allow_list: vec![],
                deny_list: vec!["10.0.0.0/8".into()],
            },
            ..Default::default()
        };
        let engine = WafEngine::new(config).unwrap();
        let req = make_req("10.1.2.3", "GET", "/");
        match engine.check(&req) {
            WafDecision::Block { rule, .. } => assert_eq!(rule, "ip_filter_deny"),
            other => panic!("expected Block, got {other:?}"),
        }
    }

    #[test]
    fn ip_allow_overrides() {
        let config = WafConfig {
            ip_filter: IpFilterConfig {
                mode: IpFilterMode::Deny,
                allow_list: vec!["10.0.0.1".into()],
                deny_list: vec!["10.0.0.0/8".into()],
            },
            ..Default::default()
        };
        let engine = WafEngine::new(config).unwrap();
        let req = make_req("10.0.0.1", "GET", "/");
        assert_eq!(engine.check(&req), WafDecision::Allow);
    }

    #[test]
    fn ddos_body_limit() {
        let config = WafConfig {
            ddos: DdosConfig {
                max_request_body_size: 50,
                ..Default::default()
            },
            ..Default::default()
        };
        let engine = WafEngine::new(config).unwrap();
        let mut req = make_req("10.0.0.1", "POST", "/upload");
        req.body = Some("x".repeat(100));
        match engine.check(&req) {
            WafDecision::Block { status: 413, .. } => {}
            other => panic!("expected 413 Block, got {other:?}"),
        }
    }

    #[test]
    fn stats_tracking() {
        let engine = WafEngine::new(WafConfig::default()).unwrap();

        // Clean request.
        engine.check(&make_req("10.0.0.1", "GET", "/"));
        // Attack request.
        let mut attack = make_req("10.0.0.1", "GET", "/search");
        attack.query = Some("q=1 UNION SELECT *".into());
        engine.check(&attack);

        let stats = engine.stats();
        assert_eq!(stats.total_requests, 2);
        assert_eq!(stats.allowed, 1);
        assert_eq!(stats.blocked, 1);
    }

    #[test]
    fn audit_log_populated() {
        let engine = WafEngine::new(WafConfig::default()).unwrap();
        let mut req = make_req("10.0.0.1", "GET", "/search");
        req.query = Some("q=1 UNION SELECT *".into());
        engine.check(&req);

        let events = engine.recent_events(10);
        assert_eq!(events.len(), 1);
        assert_eq!(events[0].rule, "sql_injection");
        assert_eq!(events[0].action, "block");
    }

    #[test]
    fn prometheus_metrics_format() {
        let engine = WafEngine::new(WafConfig::default()).unwrap();
        engine.check(&make_req("10.0.0.1", "GET", "/"));
        let metrics = engine.prometheus_metrics();
        assert!(metrics.contains("waf_requests_total"));
        assert!(metrics.contains("waf_requests_allowed"));
    }

    #[test]
    fn reload_ip_filter() {
        let engine = WafEngine::new(WafConfig::default()).unwrap();
        let req = make_req("10.0.0.1", "GET", "/");
        assert_eq!(engine.check(&req), WafDecision::Allow);

        // Reload to deny 10.0.0.1.
        engine.reload_ip_filter(IpFilterConfig {
            mode: IpFilterMode::Deny,
            allow_list: vec![],
            deny_list: vec!["10.0.0.1".into()],
        });
        assert!(matches!(engine.check(&req), WafDecision::Block { .. }));
    }

    #[test]
    fn custom_rule_blocks() {
        let config = WafConfig {
            custom_rules: vec![CustomRule {
                name: "block-admin".into(),
                match_config: MatchConfig {
                    path: Some("/admin/**".into()),
                    ..Default::default()
                },
                action: CustomRuleAction::Block,
                status: 403,
                reason: Some("Admin access denied".into()),
            }],
            ..Default::default()
        };
        let engine = WafEngine::new(config).unwrap();
        let req = make_req("10.0.0.1", "GET", "/admin/settings");
        match engine.check(&req) {
            WafDecision::Block { status: 403, .. } => {}
            other => panic!("expected 403 Block, got {other:?}"),
        }
    }

    #[test]
    fn check_with_headers_returns_headers_on_rate_limit() {
        let config = WafConfig {
            rate_limit_rules: vec![RateLimitRule {
                name: "strict".into(),
                pattern: "/**".into(),
                rpm: 1,
                burst: 0,
                key_source: rate_limit::KeySource::Ip,
                delay_mode: rate_limit::DelayMode::NoDelay,
            }],
            ..Default::default()
        };
        let engine = WafEngine::new(config).unwrap();
        let req = make_req("10.0.0.1", "GET", "/");

        // First request: allowed.
        let (d1, h1) = engine.check_with_headers(&req);
        // The check_with_headers for allow returns empty headers.
        assert!(matches!(d1, WafDecision::Allow) || h1.is_empty());

        // Second request: rate limited with headers.
        let (d2, h2) = engine.check_with_headers(&req);
        if matches!(d2, WafDecision::RateLimit { .. }) {
            assert!(h2.iter().any(|(k, _)| k == "Retry-After"));
        }
    }

    #[test]
    fn cleanup_does_not_panic() {
        let engine = WafEngine::new(WafConfig::default()).unwrap();
        engine.check(&make_req("10.0.0.1", "GET", "/"));
        engine.cleanup();
    }

    // Integration test: multiple attack vectors in one request.
    #[test]
    fn first_matching_rule_wins() {
        let config = WafConfig {
            ip_filter: IpFilterConfig {
                mode: IpFilterMode::Deny,
                allow_list: vec![],
                deny_list: vec!["10.0.0.1".into()],
            },
            ..Default::default()
        };
        let engine = WafEngine::new(config).unwrap();
        // Request has both a denied IP AND sqli — IP filter runs first.
        let mut req = make_req("10.0.0.1", "GET", "/search");
        req.query = Some("q=UNION SELECT *".into());
        match engine.check(&req) {
            WafDecision::Block { rule, .. } => assert_eq!(rule, "ip_filter_deny"),
            other => panic!("expected ip_filter_deny, got {other:?}"),
        }
    }
}