scatter-proxy 0.9.0

/// Mutates a cloned request just before it is dispatched through a specific proxy.
///
/// The scheduler calls `mutate` after `try_clone()` and before `client.execute()`,
/// giving implementations a chance to inject per-proxy request headers such as
/// cookies obtained from a prior WAF challenge solve.
///
/// # Example
///
/// ```rust,ignore
/// use scatter_proxy::RequestMutator;
/// use std::sync::Arc;
///
/// struct CookieMutator {
///     proxy_manager: Arc<scatter_proxy::ProxyManager>,
/// }
///
/// impl RequestMutator for CookieMutator {
///     fn mutate(&self, proxy_url: &str, mut req: reqwest::Request) -> reqwest::Request {
///         if let Some(cookie_val) = self.proxy_manager.get_cookie(proxy_url) {
///             let header_val = format!("acw_sc__v2={cookie_val}");
///             if let Ok(v) = header_val.parse() {
///                 req.headers_mut().insert(reqwest::header::COOKIE, v);
///             }
///         }
///         req
///     }
/// }
/// ```
pub trait RequestMutator: Send + Sync + 'static {
    /// Inject per-proxy modifications into `req` before it is sent.
    ///
    /// `proxy_url` — the proxy node URL that will execute the request.
    /// `req`       — an already `try_clone()`d request ready for execution.
    ///
    /// Return the (possibly modified) request.
    fn mutate(&self, proxy_url: &str, req: reqwest::Request) -> reqwest::Request;
}

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

    struct NoOpMutator;

    impl RequestMutator for NoOpMutator {
        fn mutate(&self, _proxy_url: &str, req: reqwest::Request) -> reqwest::Request {
            req
        }
    }

    struct HeaderInjectMutator {
        key: &'static str,
        value: &'static str,
    }

    impl RequestMutator for HeaderInjectMutator {
        fn mutate(&self, _proxy_url: &str, mut req: reqwest::Request) -> reqwest::Request {
            if let (Ok(name), Ok(val)) = (
                reqwest::header::HeaderName::from_bytes(self.key.as_bytes()),
                self.value.parse::<reqwest::header::HeaderValue>(),
            ) {
                req.headers_mut().insert(name, val);
            }
            req
        }
    }

    fn build_request() -> reqwest::Request {
        reqwest::Client::new()
            .get("http://example.com/")
            .build()
            .unwrap()
    }

    #[test]
    fn no_op_mutator_leaves_request_unchanged() {
        let m = NoOpMutator;
        let req = build_request();
        let mutated = m.mutate("socks5h://1.2.3.4:1080", req);
        assert_eq!(mutated.url().as_str(), "http://example.com/");
    }

    #[test]
    fn header_inject_mutator_adds_header() {
        let m = HeaderInjectMutator {
            key: "x-proxy-node",
            value: "node-42",
        };
        let req = build_request();
        let mutated = m.mutate("socks5h://1.2.3.4:1080", req);
        assert_eq!(
            mutated
                .headers()
                .get("x-proxy-node")
                .and_then(|v| v.to_str().ok()),
            Some("node-42"),
        );
    }

    #[test]
    fn mutator_can_be_used_as_trait_object() {
        let m: Arc<dyn RequestMutator> = Arc::new(NoOpMutator);
        let req = build_request();
        let _mutated = m.mutate("socks5h://1.2.3.4:1080", req);
    }

    #[test]
    fn different_proxy_urls_reach_mutator() {
        struct UrlCapture(std::sync::Mutex<String>);
        impl RequestMutator for UrlCapture {
            fn mutate(&self, proxy_url: &str, req: reqwest::Request) -> reqwest::Request {
                *self.0.lock().unwrap() = proxy_url.to_string();
                req
            }
        }

        let m = UrlCapture(std::sync::Mutex::new(String::new()));
        let req = build_request();
        m.mutate("socks5h://9.9.9.9:1080", req);
        assert_eq!(*m.0.lock().unwrap(), "socks5h://9.9.9.9:1080");
    }
}