captchaforge 0.2.36

//! Proxy pool — first-class HTTP/SOCKS5 proxy support for the
//! captchaforge solver chain.
//!
//! Why: every public residential-proxy provider (Bright Data,
//! Oxylabs, IPRoyal, Smartproxy, …) hands the operator a pool of
//! endpoints. Without first-class support every operator hand-rolls
//! the same wrapper. This module provides:
//!
//! - [`ProxyEntry`] — one endpoint description.
//! - [`ProxyPool`] — round-robin / sticky-per-domain selection with
//!   health tracking + auto-failover on connect-failure.
//! - [`ProxyPool::for_domain`] — sticky selection so re-solves of
//!   the same captcha on the same domain use the same exit IP
//!   (vendor anti-bot doesn't like IP-flip mid-challenge).
//!
//! The pool is wired into the solver chain via
//! `CaptchaSolverChain::with_proxy_pool` (TODO — opt-in builder
//! added when chromiumoxide's per-page-proxy support stabilises).

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::RwLock;
use std::time::{Duration, Instant};

/// One proxy endpoint. URL form: `http://user:pass@host:port` or
/// `socks5://user:pass@host:port`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProxyEntry {
    pub url: String,
    /// Optional region tag — operators can request a region-pinned
    /// solve to match the target site's expected geo.
    #[serde(default)]
    pub region: Option<String>,
    /// Operator-supplied tag for diagnostics ("ipr-rotating-us",
    /// "oxy-residential-de", …).
    #[serde(default)]
    pub label: Option<String>,
}

impl ProxyEntry {
    pub fn new(url: impl Into<String>) -> Self {
        Self {
            url: url.into(),
            region: None,
            label: None,
        }
    }

    pub fn with_region(mut self, region: impl Into<String>) -> Self {
        self.region = Some(region.into());
        self
    }

    pub fn with_label(mut self, label: impl Into<String>) -> Self {
        self.label = Some(label.into());
        self
    }
}

/// In-memory pool of proxy endpoints with health tracking.
pub struct ProxyPool {
    entries: Vec<ProxyEntry>,
    health: RwLock<HashMap<String, HealthState>>,
    /// Sticky map from domain → preferred entry index.
    sticky: RwLock<HashMap<String, usize>>,
    cursor: std::sync::atomic::AtomicUsize,
    cooldown: Duration,
}

#[derive(Debug, Clone)]
struct HealthState {
    failures: u32,
    last_failure: Instant,
}

impl ProxyPool {
    /// Construct a pool from a list of entries. Empty pool returns
    /// `None` from every selector — the chain then runs direct.
    pub fn new(entries: Vec<ProxyEntry>) -> Self {
        Self {
            entries,
            health: RwLock::new(HashMap::new()),
            sticky: RwLock::new(HashMap::new()),
            cursor: std::sync::atomic::AtomicUsize::new(0),
            cooldown: Duration::from_secs(60),
        }
    }

    pub fn with_cooldown(mut self, cooldown: Duration) -> Self {
        self.cooldown = cooldown;
        self
    }

    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    pub fn len(&self) -> usize {
        self.entries.len()
    }

    /// Pick the sticky-preferred entry for `domain`. Falls through
    /// to round-robin selection on first call OR when the previous
    /// pick has hit `cooldown` failures recently.
    pub fn for_domain(&self, domain: &str, region: Option<&str>) -> Option<ProxyEntry> {
        if self.entries.is_empty() {
            return None;
        }
        // Sticky hit?
        if let Some(&idx) = self
            .sticky
            .read()
            .unwrap_or_else(|e| e.into_inner())
            .get(domain)
        {
            if idx < self.entries.len() && !self.is_in_cooldown(&self.entries[idx].url) {
                let e = &self.entries[idx];
                if region.is_none() || e.region.as_deref() == region {
                    return Some(e.clone());
                }
            }
        }
        // Round-robin among healthy entries matching the region filter.
        let n = self.entries.len();
        for offset in 0..n {
            let idx = (self.cursor.load(std::sync::atomic::Ordering::SeqCst) + offset) % n;
            let e = &self.entries[idx];
            if region.is_some() && e.region.as_deref() != region {
                continue;
            }
            if self.is_in_cooldown(&e.url) {
                continue;
            }
            self.cursor.store(idx + 1, std::sync::atomic::Ordering::SeqCst);
            self.sticky
                .write()
                .unwrap_or_else(|e| e.into_inner())
                .insert(domain.to_string(), idx);
            return Some(e.clone());
        }
        // All entries in cooldown → fall back to the first (better
        // than nothing).
        Some(self.entries[0].clone())
    }

    fn is_in_cooldown(&self, url: &str) -> bool {
        let map = self.health.read().unwrap_or_else(|e| e.into_inner());
        match map.get(url) {
            Some(s) if s.failures >= 3 => s.last_failure.elapsed() < self.cooldown,
            _ => false,
        }
    }

    /// Record a failed solve / connect for the given entry. After 3
    /// consecutive failures the entry enters cooldown.
    pub fn record_failure(&self, url: &str) {
        let mut map = self.health.write().unwrap_or_else(|e| e.into_inner());
        let s = map.entry(url.to_string()).or_insert(HealthState {
            failures: 0,
            last_failure: Instant::now(),
        });
        s.failures = s.failures.saturating_add(1);
        s.last_failure = Instant::now();
    }

    /// Record a successful solve — resets the failure counter.
    pub fn record_success(&self, url: &str) {
        let mut map = self.health.write().unwrap_or_else(|e| e.into_inner());
        if let Some(s) = map.get_mut(url) {
            s.failures = 0;
        }
    }

    /// Snapshot every entry's health for /metrics / diagnostics.
    pub fn health_snapshot(&self) -> Vec<(String, u32)> {
        let map = self.health.read().unwrap_or_else(|e| e.into_inner());
        self.entries
            .iter()
            .map(|e| (e.url.clone(), map.get(&e.url).map(|s| s.failures).unwrap_or(0)))
            .collect()
    }
}

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

    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"),
        ]);
        // Force the cursor to land on `bad` first.
        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");
        // bad is now in cooldown; subsequent picks must avoid it.
        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());
    }
}