rullst-core 12.2.0

//! WebAssembly-compatible WAF (Web Application Firewall) middleware.

use axum::{
    body::Body,
    extract::Request,
    http::{HeaderMap, HeaderValue, StatusCode, header},
    middleware::Next,
    response::Response,
};

const MAX_INSPECTED_REQUEST_BYTES: usize = 1024 * 1024;

static MALICIOUS_PATTERNS: &[&str] = &[
    "select ",
    "union ",
    "insert ",
    "delete ",
    "drop table",
    "alter table", // SQLi
    "<script",
    "javascript:",
    "onload=",
    "onerror=",
    "document.cookie", // XSS
    "../",
    "..\\",
    "/etc/passwd",
    "win.ini", // Path Traversal
    "; ls",
    "&& cat",
    "| bash",
    "| sh",
    "wget ",
    "curl ",
    "ping -c", // Command Injection
];

fn plain_response(status: StatusCode, message: &'static str) -> Response {
    let mut response = Response::new(Body::from(message));
    *response.status_mut() = status;
    response.headers_mut().insert(
        header::CONTENT_TYPE,
        HeaderValue::from_static("text/plain; charset=utf-8"),
    );
    response
        .headers_mut()
        .insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store"));
    response
}

fn forbidden_response() -> Response {
    plain_response(
        StatusCode::FORBIDDEN,
        "Access Denied: Malicious pattern detected by Rullst Shield WAF.",
    )
}

fn should_inspect_body(headers: &HeaderMap) -> bool {
    headers
        .get(header::CONTENT_TYPE)
        .and_then(|value| value.to_str().ok())
        .map(super::media_type::essence)
        .is_some_and(super::media_type::is_inspected_request_body)
}

fn has_identity_encoding(headers: &HeaderMap) -> bool {
    headers
        .get(header::CONTENT_ENCODING)
        .is_none_or(|value| value.as_bytes().eq_ignore_ascii_case(b"identity"))
}

fn declared_body_is_too_large(headers: &HeaderMap) -> bool {
    headers
        .get(header::CONTENT_LENGTH)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.parse::<usize>().ok())
        .is_some_and(|length| length > MAX_INSPECTED_REQUEST_BYTES)
}

/// Helper to decode a hex char pair to a single byte.
#[cfg_attr(mutants, mutants::skip)]
fn hex_decode_char(c1: u8, c2: u8) -> Option<u8> {
    let b1 = (c1 as char).to_digit(16)?;
    let b2 = (c2 as char).to_digit(16)?;
    Some(((b1 << 4) | b2) as u8)
}

/// WebAssembly-compatible URL decoding helper.
#[cfg_attr(mutants, mutants::skip)]
fn url_decode(s: &str) -> String {
    let mut decoded_bytes = Vec::with_capacity(s.len());
    let bytes = s.as_bytes();
    let mut i = 0;
    while i < bytes.len() {
        let b = bytes[i];
        if b == b'+' {
            decoded_bytes.push(b' ');
            i += 1;
            continue;
        }
        if b == b'%' && i + 2 < bytes.len() {
            let h1 = bytes[i + 1];
            let h2 = bytes[i + 2];
            if let Some(d) = hex_decode_char(h1, h2) {
                decoded_bytes.push(d);
                i += 3;
                continue;
            }
        }
        decoded_bytes.push(b);
        i += 1;
    }
    String::from_utf8_lossy(&decoded_bytes).into_owned()
}

fn contains_malicious_pattern(payload: &str) -> bool {
    let payload_decoded = url_decode(payload);
    let payload_lower = payload_decoded.to_lowercase();
    MALICIOUS_PATTERNS
        .iter()
        .any(|pattern| payload_lower.contains(pattern))
}

async fn inspect_and_restore_body(req: Request) -> Result<Request, Box<Response>> {
    if !should_inspect_body(req.headers()) {
        return Ok(req);
    }

    if !has_identity_encoding(req.headers()) {
        return Err(Box::new(plain_response(
            StatusCode::UNSUPPORTED_MEDIA_TYPE,
            "Encoded request bodies cannot be inspected by the WAF.",
        )));
    }

    if declared_body_is_too_large(req.headers()) {
        return Err(Box::new(plain_response(
            StatusCode::PAYLOAD_TOO_LARGE,
            "Request body exceeds the WAF inspection limit.",
        )));
    }

    let (parts, body) = req.into_parts();
    let bytes = axum::body::to_bytes(body, MAX_INSPECTED_REQUEST_BYTES)
        .await
        .map_err(|_| {
            Box::new(plain_response(
                StatusCode::PAYLOAD_TOO_LARGE,
                "Request body could not be inspected within the WAF limit.",
            ))
        })?;

    let payload = std::str::from_utf8(&bytes).map_err(|_| {
        Box::new(plain_response(
            StatusCode::BAD_REQUEST,
            "Declared textual request body is not valid UTF-8.",
        ))
    })?;
    if contains_malicious_pattern(payload) {
        return Err(Box::new(forbidden_response()));
    }

    Ok(Request::from_parts(parts, Body::from(bytes)))
}

/// WebAssembly-compatible WAF middleware for traffic control and malicious bot protection.
pub async fn waf_middleware(mut req: Request, next: Next) -> Response {
    // 1. Inspect User-Agent for known bots or scrapers
    if let Some(ua) = req
        .headers()
        .get(header::USER_AGENT)
        .and_then(|v| v.to_str().ok())
    {
        let ua_lower = ua.to_lowercase();
        let suspicious_agents = req
            .extensions()
            .get::<crate::config::SecurityConfig>()
            .map(|cfg| cfg.user_agent_blocklist.clone())
            .unwrap_or_else(|| {
                crate::config::RullstConfig::global()
                    .security
                    .user_agent_blocklist
                    .clone()
            });

        for agent in suspicious_agents {
            if ua_lower.contains(&agent.to_lowercase()) {
                return plain_response(
                    StatusCode::FORBIDDEN,
                    "Access Denied: Suspicious User-Agent blocked by Rullst Shield WAF.",
                );
            }
        }
    }

    // 2. Inspect query parameters and selected headers for common attack vectors.
    if let Some(query) = req.uri().query() {
        if contains_malicious_pattern(query) {
            return forbidden_response();
        }
    }

    for header_name in [header::REFERER, header::COOKIE] {
        if let Some(payload) = req
            .headers()
            .get(header_name)
            .and_then(|value| value.to_str().ok())
            && contains_malicious_pattern(payload)
        {
            return forbidden_response();
        }
    }

    // 3. Inspect bounded textual/JSON/form bodies, then reconstruct the exact request for
    // downstream extractors. Unsupported encodings and over-limit text fail closed.
    req = match inspect_and_restore_body(req).await {
        Ok(req) => req,
        Err(response) => return *response,
    };

    next.run(req).await
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;
    use axum::{Router, body::Bytes, http::Request, routing::post};
    use tower::ServiceExt;

    #[tokio::test]
    async fn case_and_suffix_variants_cannot_skip_body_inspection() {
        let app = Router::new()
            .route("/items", post(|body: Bytes| async move { body }))
            .route_layer(axum::middleware::from_fn(waf_middleware));
        for (media_type, body) in [
            (
                "application/vnd.api+JSON",
                r#"{"bio":"<script>document.cookie</script>"}"#,
            ),
            (
                "Application/problem+json",
                r#"{"q":"x' union select pw from users--"}"#,
            ),
            ("application/json+patch", r#"{"path":"../../etc/passwd"}"#),
            (
                "application/x-www-form-urlencodedX",
                "q=x%27%20union%20select%20pw%20from%20users--",
            ),
        ] {
            let request = Request::post("/items")
                .header(header::CONTENT_TYPE, media_type)
                .body(Body::from(body))
                .unwrap();
            assert_eq!(
                app.clone().oneshot(request).await.unwrap().status(),
                StatusCode::FORBIDDEN,
                "{media_type} body was not inspected"
            );
        }
    }
}