scatter_proxy/

classifier.rs

use http::{HeaderMap, StatusCode};

/// Four-level verdict returned by a [`BodyClassifier`] for every HTTP response.
///
/// The scheduler uses this to decide what to do next:
/// - `Success` — task is done, proxy gets a positive health mark.
/// - `ProxyBlocked` — this proxy failed for the target; counts as a proxy failure.
/// - `TargetError` — the target itself is unhealthy; does **not** penalise the proxy.
/// - `Challenge` — a WAF challenge page was returned; the scheduler will attempt to
///   resolve it on the *same* proxy node via [`ChallengeResolver`], then retry.
///   Without a resolver configured, `Challenge` degrades to `ProxyBlocked`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BodyVerdict {
    Success,
    ProxyBlocked,
    TargetError,
    /// WAF challenge page — resolve on the same proxy, then retry.
    Challenge,
}

/// Trait that inspects a completed HTTP response and returns a [`BodyVerdict`].
///
/// Implement this to customise how your target site's responses are interpreted.
pub trait BodyClassifier: Send + Sync + 'static {
    fn classify(&self, status: StatusCode, headers: &HeaderMap, body: &[u8]) -> BodyVerdict;
}

/// Built-in classifier that covers the common case:
///
/// | Condition | Verdict |
/// |-----------|---------|
/// | 2xx **and** non-empty body | `Success` |
/// | 2xx **and** empty body | `ProxyBlocked` |
/// | 403 or 429 | `ProxyBlocked` |
/// | 5xx | `TargetError` |
/// | everything else | `ProxyBlocked` |
pub struct DefaultClassifier;

impl BodyClassifier for DefaultClassifier {
    fn classify(&self, status: StatusCode, _headers: &HeaderMap, body: &[u8]) -> BodyVerdict {
        if status.is_success() {
            if body.is_empty() {
                BodyVerdict::ProxyBlocked
            } else {
                BodyVerdict::Success
            }
        } else if status == StatusCode::FORBIDDEN || status == StatusCode::TOO_MANY_REQUESTS {
            BodyVerdict::ProxyBlocked
        } else if status.is_server_error() {
            BodyVerdict::TargetError
        } else {
            BodyVerdict::ProxyBlocked
        }
    }
}

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

    /// Helper that classifies with the default classifier, empty headers, and the
    /// supplied status + body.
    fn classify(status: u16, body: &[u8]) -> BodyVerdict {
        let classifier = DefaultClassifier;
        let headers = HeaderMap::new();
        classifier.classify(StatusCode::from_u16(status).unwrap(), &headers, body)
    }

    // ── 2xx ──────────────────────────────────────────────────────────────

    #[test]
    fn success_200_with_body() {
        assert_eq!(classify(200, b"hello"), BodyVerdict::Success);
    }

    #[test]
    fn success_201_with_body() {
        assert_eq!(classify(201, b"{\"id\":1}"), BodyVerdict::Success);
    }

    #[test]
    fn success_204_no_content_empty_body() {
        // 204 is 2xx but body is empty → ProxyBlocked per the default rules.
        assert_eq!(classify(204, b""), BodyVerdict::ProxyBlocked);
    }

    #[test]
    fn proxy_blocked_200_empty_body() {
        assert_eq!(classify(200, b""), BodyVerdict::ProxyBlocked);
    }

    // ── 403 / 429 ───────────────────────────────────────────────────────

    #[test]
    fn proxy_blocked_403() {
        assert_eq!(classify(403, b"Forbidden"), BodyVerdict::ProxyBlocked);
    }

    #[test]
    fn proxy_blocked_429() {
        assert_eq!(classify(429, b"Rate limited"), BodyVerdict::ProxyBlocked);
    }

    // ── 5xx ─────────────────────────────────────────────────────────────

    #[test]
    fn target_error_500() {
        assert_eq!(
            classify(500, b"Internal Server Error"),
            BodyVerdict::TargetError
        );
    }

    #[test]
    fn target_error_502() {
        assert_eq!(classify(502, b"Bad Gateway"), BodyVerdict::TargetError);
    }

    #[test]
    fn target_error_503() {
        assert_eq!(
            classify(503, b"Service Unavailable"),
            BodyVerdict::TargetError
        );
    }

    // ── Other status codes ──────────────────────────────────────────────

    #[test]
    fn proxy_blocked_301_redirect() {
        assert_eq!(classify(301, b"Moved"), BodyVerdict::ProxyBlocked);
    }

    #[test]
    fn proxy_blocked_400_bad_request() {
        assert_eq!(classify(400, b"Bad Request"), BodyVerdict::ProxyBlocked);
    }

    #[test]
    fn proxy_blocked_401_unauthorised() {
        assert_eq!(classify(401, b"Unauthorised"), BodyVerdict::ProxyBlocked);
    }

    #[test]
    fn proxy_blocked_404_not_found() {
        assert_eq!(classify(404, b"Not Found"), BodyVerdict::ProxyBlocked);
    }

    #[test]
    fn proxy_blocked_407_proxy_auth_required() {
        assert_eq!(
            classify(407, b"Proxy Authentication Required"),
            BodyVerdict::ProxyBlocked
        );
    }

    // ── Headers are forwarded (even if DefaultClassifier ignores them) ──

    #[test]
    fn headers_are_available_to_classifier() {
        let classifier = DefaultClassifier;
        let mut headers = HeaderMap::new();
        headers.insert("x-custom", "value".parse().unwrap());
        // Should still produce Success for 200 + body regardless of headers.
        assert_eq!(
            classifier.classify(StatusCode::OK, &headers, b"data"),
            BodyVerdict::Success,
        );
    }

    // ── Trait object safety ─────────────────────────────────────────────

    #[test]
    fn can_be_used_as_trait_object() {
        let classifier: Box<dyn BodyClassifier> = Box::new(DefaultClassifier);
        let headers = HeaderMap::new();
        assert_eq!(
            classifier.classify(StatusCode::OK, &headers, b"body"),
            BodyVerdict::Success,
        );
    }

    // ── Custom classifier ───────────────────────────────────────────────

    struct AlwaysSuccess;

    impl BodyClassifier for AlwaysSuccess {
        fn classify(&self, _status: StatusCode, _headers: &HeaderMap, _body: &[u8]) -> BodyVerdict {
            BodyVerdict::Success
        }
    }

    #[test]
    fn custom_classifier_overrides_defaults() {
        let classifier = AlwaysSuccess;
        let headers = HeaderMap::new();
        // Even a 500 is considered success by AlwaysSuccess.
        assert_eq!(
            classifier.classify(StatusCode::INTERNAL_SERVER_ERROR, &headers, b""),
            BodyVerdict::Success,
        );
    }

    // ── BodyVerdict derives ─────────────────────────────────────────────

    #[test]
    fn verdict_is_copy_and_clone() {
        let v = BodyVerdict::Success;
        let v2 = v; // Copy
        let v3 = v; // Clone (Copy)
        assert_eq!(v, v2);
        assert_eq!(v2, v3);
    }

    #[test]
    fn verdict_debug_format() {
        assert_eq!(format!("{:?}", BodyVerdict::Success), "Success");
        assert_eq!(format!("{:?}", BodyVerdict::ProxyBlocked), "ProxyBlocked");
        assert_eq!(format!("{:?}", BodyVerdict::TargetError), "TargetError");
        assert_eq!(format!("{:?}", BodyVerdict::Challenge), "Challenge");
    }

    // ── Challenge variant ───────────────────────────────────────────────

    struct WafChallengeClassifier;

    impl BodyClassifier for WafChallengeClassifier {
        fn classify(&self, status: StatusCode, _headers: &HeaderMap, body: &[u8]) -> BodyVerdict {
            if status.is_success() {
                if body.starts_with(b"<script>") {
                    BodyVerdict::Challenge
                } else if body.is_empty() {
                    BodyVerdict::ProxyBlocked
                } else {
                    BodyVerdict::Success
                }
            } else {
                BodyVerdict::ProxyBlocked
            }
        }
    }

    #[test]
    fn custom_classifier_returns_challenge() {
        let classifier = WafChallengeClassifier;
        let headers = HeaderMap::new();
        assert_eq!(
            classifier.classify(StatusCode::OK, &headers, b"<script>var x=1</script>"),
            BodyVerdict::Challenge,
        );
    }

    #[test]
    fn challenge_variant_is_eq() {
        assert_eq!(BodyVerdict::Challenge, BodyVerdict::Challenge);
        assert_ne!(BodyVerdict::Challenge, BodyVerdict::ProxyBlocked);
    }
}