captchaforge 0.2.38

//! CaptchaForge proxy-pool compatibility surface backed by proxywire.
//!
//! CaptchaForge originally exposed `ProxyEntry { url, region, label }` and a
//! sticky `ProxyPool::for_domain` selector. The canonical proxy routing and
//! health implementation now lives in `proxywire`; this module keeps the
//! CaptchaForge import path and method names stable while delegating selection,
//! cooldown, and provider routing to that shared library.

use std::{collections::HashMap, time::Duration};

use serde::{Deserialize, Serialize};

pub use proxywire::{PoolRotation, ProxyEndpoint, ProxyProtocol};

/// One proxy endpoint in CaptchaForge's original URL-oriented shape.
///
/// URL form: `http://user:pass@host:port`, `socks5://host:port`, or another
/// scheme accepted by [`proxywire::ProxyEndpoint::from_url`].
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ProxyEntry {
    /// Proxy URL.
    pub url: String,
    /// Optional region tag used by region-pinned selection.
    #[serde(default)]
    pub region: Option<String>,
    /// Operator-supplied diagnostic label.
    #[serde(default)]
    pub label: Option<String>,
}

impl ProxyEntry {
    /// Creates an entry from a proxy URL.
    #[must_use]
    pub fn new(url: impl Into<String>) -> Self {
        Self {
            url: url.into(),
            region: None,
            label: None,
        }
    }

    /// Adds a region tag to the entry.
    #[must_use]
    pub fn with_region(mut self, region: impl Into<String>) -> Self {
        self.region = Some(region.into());
        self
    }

    /// Adds a diagnostic label to the entry.
    #[must_use]
    pub fn with_label(mut self, label: impl Into<String>) -> Self {
        self.label = Some(label.into());
        self
    }

    fn to_endpoint(&self) -> Option<ProxyEndpoint> {
        let mut endpoint = ProxyEndpoint::from_url(&self.url).ok()?;
        endpoint.region = self.region.clone();
        endpoint.label = self.label.clone();
        Some(endpoint)
    }
}

/// In-memory sticky proxy pool with health tracking.
#[derive(Debug)]
pub struct ProxyPool {
    entries: Vec<ProxyEntry>,
    endpoints: Vec<ProxyEndpoint>,
    endpoint_entry_indices: Vec<usize>,
    rotation: PoolRotation,
    pool: proxywire::PoolProxy,
}

impl ProxyPool {
    /// Constructs a sticky per-domain proxy pool from CaptchaForge entries.
    #[must_use]
    pub fn new(entries: Vec<ProxyEntry>) -> Self {
        Self::with_rotation(entries, PoolRotation::PerDomain)
    }

    /// Constructs a pool with an explicit shared proxywire rotation strategy.
    #[must_use]
    pub fn with_rotation(entries: Vec<ProxyEntry>, rotation: PoolRotation) -> Self {
        let mut endpoints = Vec::new();
        let mut endpoint_entry_indices = Vec::new();

        for (index, entry) in entries.iter().enumerate() {
            if let Some(endpoint) = entry.to_endpoint() {
                endpoints.push(endpoint);
                endpoint_entry_indices.push(index);
            }
        }

        Self {
            entries,
            endpoints: endpoints.clone(),
            endpoint_entry_indices,
            rotation,
            pool: proxywire::PoolProxy::new(endpoints, rotation),
        }
    }

    /// Overrides the cooldown duration before failed entries re-enter rotation.
    #[must_use]
    pub fn with_cooldown(mut self, cooldown: Duration) -> Self {
        self.pool = proxywire::PoolProxy::new(self.endpoints.clone(), self.rotation)
            .with_cooldown(cooldown);
        self
    }

    /// Number of entries supplied to the pool.
    #[must_use]
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    /// Whether the pool has no configured entries.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// Selects the sticky-preferred entry for a domain and optional region.
    #[must_use]
    pub fn for_domain(&self, domain: &str, region: Option<&str>) -> Option<ProxyEntry> {
        let endpoint = self.pool.select_with_region(domain, region)?;
        self.entry_for_endpoint(&endpoint)
    }

    /// Exposes the consolidated proxywire endpoint selector for callers that
    /// need the shared endpoint type directly.
    #[must_use]
    pub fn select_with_region(&self, key: &str, region: Option<&str>) -> Option<ProxyEndpoint> {
        self.pool.select_with_region(key, region)
    }

    /// Records a failed solve or connect for a URL returned by this pool.
    pub fn record_failure(&self, url: &str) {
        if let Some(endpoint) = self.endpoint_for_url(url) {
            self.pool.record_failure(&endpoint);
        }
    }

    /// Records a successful solve or connect for a URL returned by this pool.
    pub fn record_success(&self, url: &str) {
        if let Some(endpoint) = self.endpoint_for_url(url) {
            self.pool.record_success(&endpoint);
        }
    }

    /// Snapshot every configured entry's current failure count.
    #[must_use]
    pub fn health_snapshot(&self) -> Vec<(String, u32)> {
        let failures: HashMap<_, _> = self.pool.health_snapshot().into_iter().collect();
        self.entries
            .iter()
            .map(|entry| {
                let count = entry
                    .to_endpoint()
                    .and_then(|endpoint| failures.get(&endpoint.to_url()).copied())
                    .unwrap_or(0);
                (entry.url.clone(), count)
            })
            .collect()
    }

    fn entry_for_endpoint(&self, endpoint: &ProxyEndpoint) -> Option<ProxyEntry> {
        let endpoint_index = self
            .endpoints
            .iter()
            .position(|candidate| candidate == endpoint)?;
        let entry_index = self.endpoint_entry_indices.get(endpoint_index)?;
        self.entries.get(*entry_index).cloned()
    }

    fn endpoint_for_url(&self, url: &str) -> Option<ProxyEndpoint> {
        self.endpoint_entry_indices
            .iter()
            .enumerate()
            .find_map(|(endpoint_index, entry_index)| {
                let entry = self.entries.get(*entry_index)?;
                (entry.url == url).then(|| self.endpoints[endpoint_index].clone())
            })
            .or_else(|| ProxyEndpoint::from_url(url).ok())
    }
}

impl proxywire::ProxyProvider for ProxyPool {
    fn route_for(
        &self,
        context: &proxywire::ProxyRequestContext,
    ) -> proxywire::Result<proxywire::ProxyRoute> {
        proxywire::ProxyProvider::route_for(&self.pool, context)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use proxywire::{ProxyProvider, ProxyRequestContext, Scheme};

    fn pool() -> ProxyPool {
        ProxyPool::new(vec![
            ProxyEntry::new("http://p1.example:8080").with_region("us"),
            ProxyEntry::new("http://p2.example:8080").with_region("us"),
            ProxyEntry::new("http://p3.example:8080").with_region("de"),
        ])
    }

    #[test]
    fn for_domain_returns_sticky_pick() {
        let p = pool();
        let first = p.for_domain("example.com", None).expect("non-empty");
        let second = p.for_domain("example.com", None).expect("sticky");
        assert_eq!(first.url, second.url);
    }

    #[test]
    fn region_filter_excludes_other_regions() {
        let p = pool();
        let de = p.for_domain("example.com", Some("de")).expect("de");
        assert_eq!(de.region.as_deref(), Some("de"));
    }

    #[test]
    fn empty_pool_returns_none() {
        let p = ProxyPool::new(Vec::new());
        assert!(p.for_domain("any", None).is_none());
    }

    #[test]
    fn cooldown_after_three_failures_excludes_entry() {
        let p = ProxyPool::new(vec![
            ProxyEntry::new("http://bad.example:8080").with_region("x"),
            ProxyEntry::new("http://good.example:8080").with_region("x"),
        ]);
        let _ = p.for_domain("d1", Some("x"));
        p.record_failure("http://bad.example:8080");
        p.record_failure("http://bad.example:8080");
        p.record_failure("http://bad.example:8080");
        let pick = p.for_domain("d2", Some("x")).unwrap();
        assert_ne!(pick.url, "http://bad.example:8080");
    }

    #[test]
    fn record_success_resets_failures() {
        let p = ProxyPool::new(vec![ProxyEntry::new("http://x.example:8080")]);
        p.record_failure("http://x.example:8080");
        p.record_failure("http://x.example:8080");
        p.record_success("http://x.example:8080");
        let snap = p.health_snapshot();
        assert_eq!(snap[0].1, 0);
    }

    #[test]
    fn health_snapshot_contains_every_entry() {
        let p = pool();
        let snap = p.health_snapshot();
        assert_eq!(snap.len(), p.len());
    }

    #[test]
    fn proxywire_provider_routes_through_selected_entry() {
        let p = ProxyPool::new(vec![ProxyEntry::new("http://p1.example:8080")]);
        let route = p
            .route_for(&ProxyRequestContext {
                scheme: Scheme::Https,
                host: "example.com".to_string(),
                port: 443,
            })
            .unwrap();
        assert_eq!(route.hops()[0].host, "p1.example");
    }
}