scatter_proxy/

proxy.rs

use dashmap::DashMap;

use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use std::time::Duration;

use crate::error::ScatterProxyError;

/// State of a proxy in the manager.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProxyState {
    Unknown,
    Active,
    /// Proxy was eliminated due to sustained failure.  Not automatically
    /// recovered — requires a restart or explicit state reset.
    Dead,
    /// Proxy reached its per-node request limit and was voluntarily retired.
    /// Unlike `Dead`, `Retired` nodes can be reactivated on source refresh.
    Retired,
}

/// Internal info stored per proxy.
struct ProxyInfo {
    state: ProxyState,
}

/// Manages a set of SOCKS5 proxies: fetching from sources, normalizing URLs,
/// tracking state, and caching `reqwest::Client` instances.
pub struct ProxyManager {
    proxies: DashMap<String, ProxyInfo>,
    clients: DashMap<String, reqwest::Client>,
    proxy_timeout: Duration,
    /// Per-proxy WAF cookie values (`proxy_url` → cookie value).
    cookies: DashMap<String, String>,
    /// Per-proxy cumulative request counters.
    request_counts: DashMap<String, Arc<AtomicU32>>,
    /// Per-proxy semaphore ensuring at most one concurrent challenge solve.
    challenge_locks: DashMap<String, Arc<tokio::sync::Semaphore>>,
}

impl ProxyManager {
    /// Create a new `ProxyManager` with the given per-proxy connection timeout.
    pub fn new(proxy_timeout: Duration) -> Self {
        let proxies = DashMap::new();
        Self {
            proxies,
            clients: DashMap::new(),
            proxy_timeout,
            cookies: DashMap::new(),
            request_counts: DashMap::new(),
            challenge_locks: DashMap::new(),
        }
    }

    // ── Per-proxy cookie management ──────────────────────────────────────────

    /// Retrieve the stored cookie value for a proxy node, if any.
    pub fn get_cookie(&self, proxy_url: &str) -> Option<String> {
        self.cookies.get(proxy_url).map(|v| v.clone())
    }

    /// Store a cookie value for a proxy node, replacing any existing value.
    pub fn set_cookie(&self, proxy_url: &str, value: String) {
        self.cookies.insert(proxy_url.to_string(), value);
    }

    // ── Per-proxy request counting ───────────────────────────────────────────

    /// Atomically increment the request count for a proxy and return the new value.
    pub fn increment_request_count(&self, proxy_url: &str) -> u32 {
        self.request_counts
            .entry(proxy_url.to_string())
            .or_insert_with(|| Arc::new(AtomicU32::new(0)))
            .fetch_add(1, Ordering::Relaxed)
            + 1
    }

    /// Return the current request count for a proxy node.
    pub fn request_count(&self, proxy_url: &str) -> u32 {
        self.request_counts
            .get(proxy_url)
            .map(|c| c.load(Ordering::Relaxed))
            .unwrap_or(0)
    }

    /// Reset the request count for a proxy node and, if the node was `Retired`,
    /// restore it to `Active`.  `Dead` nodes are not restored.
    pub fn reset_request_count(&self, proxy_url: &str) {
        if let Some(counter) = self.request_counts.get(proxy_url) {
            counter.store(0, Ordering::Relaxed);
        }
        if matches!(self.get_state(proxy_url), ProxyState::Retired) {
            self.set_state(proxy_url, ProxyState::Active);
        }
    }

    // ── Per-proxy challenge serialization ────────────────────────────────────

    /// Acquire the per-proxy challenge lock (semaphore with 1 permit).
    ///
    /// Ensures at most one concurrent challenge solve per proxy node.
    pub fn challenge_lock(&self, proxy_url: &str) -> Arc<tokio::sync::Semaphore> {
        self.challenge_locks
            .entry(proxy_url.to_string())
            .or_insert_with(|| Arc::new(tokio::sync::Semaphore::new(1)))
            .clone()
    }

    /// Fetch proxy lists from the given source URLs, normalize each line, and
    /// add any new proxies to the manager.
    ///
    /// Returns the count of newly added proxies across all sources.
    /// Errors fetching individual sources are logged but do not cause the
    /// overall call to fail.
    pub async fn fetch_and_add(
        &self,
        sources: &[String],
        prefer_remote_dns: bool,
    ) -> Result<usize, ScatterProxyError> {
        let http_client = reqwest::Client::builder()
            .timeout(Duration::from_secs(30))
            .build()
            .map_err(|e| ScatterProxyError::Init(format!("failed to build HTTP client: {e}")))?;

        let mut added = 0usize;

        for source in sources {
            match http_client.get(source).send().await {
                Ok(resp) => match resp.text().await {
                    Ok(body) => {
                        for line in body.lines() {
                            let normalized = Self::normalize_proxy_url(line, prefer_remote_dns);
                            if normalized.is_empty() {
                                continue;
                            }
                            if !self.proxies.contains_key(&normalized) {
                                self.proxies.insert(
                                    normalized,
                                    ProxyInfo {
                                        state: ProxyState::Unknown,
                                    },
                                );
                                added += 1;
                            }
                        }
                    }
                    Err(e) => {
                        tracing::error!("failed to read body from source {source}: {e}");
                    }
                },
                Err(e) => {
                    tracing::error!("failed to fetch proxy source {source}: {e}");
                }
            }
        }

        Ok(added)
    }

    /// Get all proxy URLs currently in the manager.
    pub fn all_proxy_urls(&self) -> Vec<String> {
        self.proxies.iter().map(|r| r.key().clone()).collect()
    }

    /// Get proxy URLs that are eligible for scheduling (not `Dead` or `Retired`).
    pub fn get_active_proxies(&self) -> Vec<String> {
        self.proxies
            .iter()
            .filter(|r| !matches!(r.value().state, ProxyState::Dead | ProxyState::Retired))
            .map(|r| r.key().clone())
            .collect()
    }

    /// Set the state of a proxy.
    pub fn set_state(&self, proxy: &str, state: ProxyState) {
        if let Some(mut entry) = self.proxies.get_mut(proxy) {
            entry.state = state;
        }
    }

    /// Get the state of a proxy. Returns `Unknown` if the proxy is not tracked.
    pub fn get_state(&self, proxy: &str) -> ProxyState {
        self.proxies
            .get(proxy)
            .map(|r| r.value().state)
            .unwrap_or(ProxyState::Unknown)
    }

    /// Count proxies by state.
    ///
    /// Returns `(total, active_or_unknown, cooldown_placeholder, dead, retired)`.
    /// The cooldown placeholder is always 0 — cooldown tracking lives in a separate module.
    pub fn proxy_counts(&self) -> (usize, usize, usize, usize, usize) {
        let mut total = 0usize;
        let mut active_or_unknown = 0usize;
        let mut dead = 0usize;
        let mut retired = 0usize;

        for entry in self.proxies.iter() {
            total += 1;
            match entry.value().state {
                ProxyState::Dead => dead += 1,
                ProxyState::Retired => retired += 1,
                _ => active_or_unknown += 1,
            }
        }

        (total, active_or_unknown, 0, dead, retired)
    }

    /// Get or create a `reqwest::Client` configured to route through the given
    /// proxy URL.  Clients are cached for reuse.
    pub fn get_client(&self, proxy_url: &str) -> Result<reqwest::Client, ScatterProxyError> {
        if let Some(client) = self.clients.get(proxy_url) {
            return Ok(client.value().clone());
        }

        let proxy = reqwest::Proxy::all(proxy_url).map_err(|e| {
            ScatterProxyError::Init(format!("invalid proxy URL '{proxy_url}': {e}"))
        })?;

        let client = reqwest::Client::builder()
            .proxy(proxy)
            .timeout(self.proxy_timeout)
            .build()
            .map_err(|e| {
                ScatterProxyError::Init(format!(
                    "failed to build client for proxy '{proxy_url}': {e}"
                ))
            })?;

        self.clients.insert(proxy_url.to_string(), client.clone());
        Ok(client)
    }

    /// Normalize a raw proxy string to a full URL.
    ///
    /// Rules:
    /// - Empty lines and lines starting with `#` are returned as empty string.
    /// - `"socks5h://ip:port"` → returned as-is.
    /// - `"socks5://ip:port"` → replaced with `"socks5h://"` when
    ///   `prefer_remote_dns` is true, otherwise kept as-is.
    /// - Bare `"ip:port"` → prepended with `"socks5h://"` when
    ///   `prefer_remote_dns` is true, otherwise `"socks5://"`.
    pub fn normalize_proxy_url(raw: &str, prefer_remote_dns: bool) -> String {
        let trimmed = raw.trim();

        if trimmed.is_empty() || trimmed.starts_with('#') {
            return String::new();
        }

        if trimmed.starts_with("socks5h://") {
            return trimmed.to_string();
        }

        if let Some(rest) = trimmed.strip_prefix("socks5://") {
            if prefer_remote_dns {
                return format!("socks5h://{rest}");
            }
            return trimmed.to_string();
        }

        // Also handle http/https proxies as-is
        if trimmed.starts_with("http://") || trimmed.starts_with("https://") {
            return trimmed.to_string();
        }

        // Bare ip:port
        if prefer_remote_dns {
            format!("socks5h://{trimmed}")
        } else {
            format!("socks5://{trimmed}")
        }
    }

    /// Total number of proxies being tracked.
    #[allow(dead_code)]
    pub fn total_count(&self) -> usize {
        self.proxies.len()
    }
}

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

    // -----------------------------------------------------------------------
    // normalize_proxy_url
    // -----------------------------------------------------------------------

    #[test]
    fn normalize_bare_ip_port_prefer_remote_dns() {
        let result = ProxyManager::normalize_proxy_url("1.2.3.4:1080", true);
        assert_eq!(result, "socks5h://1.2.3.4:1080");
    }

    #[test]
    fn normalize_bare_ip_port_local_dns() {
        let result = ProxyManager::normalize_proxy_url("1.2.3.4:1080", false);
        assert_eq!(result, "socks5://1.2.3.4:1080");
    }

    #[test]
    fn normalize_socks5_to_socks5h_when_prefer_remote() {
        let result = ProxyManager::normalize_proxy_url("socks5://1.2.3.4:1080", true);
        assert_eq!(result, "socks5h://1.2.3.4:1080");
    }

    #[test]
    fn normalize_socks5_stays_when_not_prefer_remote() {
        let result = ProxyManager::normalize_proxy_url("socks5://1.2.3.4:1080", false);
        assert_eq!(result, "socks5://1.2.3.4:1080");
    }

    #[test]
    fn normalize_socks5h_unchanged() {
        let result = ProxyManager::normalize_proxy_url("socks5h://1.2.3.4:1080", true);
        assert_eq!(result, "socks5h://1.2.3.4:1080");
    }

    #[test]
    fn normalize_socks5h_unchanged_no_prefer() {
        let result = ProxyManager::normalize_proxy_url("socks5h://1.2.3.4:1080", false);
        assert_eq!(result, "socks5h://1.2.3.4:1080");
    }

    #[test]
    fn normalize_empty_string() {
        let result = ProxyManager::normalize_proxy_url("", true);
        assert!(result.is_empty());
    }

    #[test]
    fn normalize_whitespace_only() {
        let result = ProxyManager::normalize_proxy_url("   ", true);
        assert!(result.is_empty());
    }

    #[test]
    fn normalize_comment_line() {
        let result = ProxyManager::normalize_proxy_url("# this is a comment", true);
        assert!(result.is_empty());
    }

    #[test]
    fn normalize_trims_whitespace() {
        let result = ProxyManager::normalize_proxy_url("  1.2.3.4:1080  ", true);
        assert_eq!(result, "socks5h://1.2.3.4:1080");
    }

    #[test]
    fn normalize_http_proxy_passed_through() {
        let result = ProxyManager::normalize_proxy_url("http://proxy.example.com:8080", true);
        assert_eq!(result, "http://proxy.example.com:8080");
    }

    #[test]
    fn normalize_https_proxy_passed_through() {
        let result = ProxyManager::normalize_proxy_url("https://proxy.example.com:8080", false);
        assert_eq!(result, "https://proxy.example.com:8080");
    }

    #[test]
    fn normalize_socks5_with_auth() {
        let result = ProxyManager::normalize_proxy_url("socks5://user:pass@1.2.3.4:1080", true);
        assert_eq!(result, "socks5h://user:pass@1.2.3.4:1080");
    }

    // -----------------------------------------------------------------------
    // ProxyManager::new
    // -----------------------------------------------------------------------

    // -----------------------------------------------------------------------
    // State management
    // -----------------------------------------------------------------------

    #[test]
    fn get_state_returns_unknown_for_missing_proxy() {
        let mgr = ProxyManager::new(Duration::from_secs(5));
        assert_eq!(mgr.get_state("socks5h://9.9.9.9:1080"), ProxyState::Unknown);
    }

    #[test]
    fn set_and_get_state() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        pm.proxies.insert(
            "socks5h://1.2.3.4:1080".to_string(),
            ProxyInfo {
                state: ProxyState::Unknown,
            },
        );

        assert_eq!(pm.get_state("socks5h://1.2.3.4:1080"), ProxyState::Unknown);

        pm.set_state("socks5h://1.2.3.4:1080", ProxyState::Active);
        assert_eq!(pm.get_state("socks5h://1.2.3.4:1080"), ProxyState::Active);

        pm.set_state("socks5h://1.2.3.4:1080", ProxyState::Dead);
        assert_eq!(pm.get_state("socks5h://1.2.3.4:1080"), ProxyState::Dead);
    }

    #[test]
    fn set_state_on_missing_proxy_is_noop() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        pm.set_state("socks5h://nonexistent:1080", ProxyState::Dead);
        assert_eq!(pm.total_count(), 0);
    }

    // -----------------------------------------------------------------------
    // get_active_proxies
    // -----------------------------------------------------------------------

    #[test]
    fn get_active_proxies_excludes_dead() {
        let mgr = ProxyManager::new(Duration::from_secs(5));

        // Manually insert some proxies.
        mgr.proxies.insert(
            "socks5h://1.1.1.1:1080".into(),
            ProxyInfo {
                state: ProxyState::Active,
            },
        );
        mgr.proxies.insert(
            "socks5h://2.2.2.2:1080".into(),
            ProxyInfo {
                state: ProxyState::Dead,
            },
        );
        mgr.proxies.insert(
            "socks5h://3.3.3.3:1080".into(),
            ProxyInfo {
                state: ProxyState::Unknown,
            },
        );

        let active = mgr.get_active_proxies();
        assert_eq!(active.len(), 2);
        assert!(active.contains(&"socks5h://1.1.1.1:1080".to_string()));
        assert!(!active.contains(&"socks5h://2.2.2.2:1080".to_string()));
        assert!(active.contains(&"socks5h://3.3.3.3:1080".to_string()));
    }

    // -----------------------------------------------------------------------
    // proxy_counts
    // -----------------------------------------------------------------------

    #[test]
    fn proxy_counts_correct() {
        let pm = ProxyManager::new(Duration::from_secs(8));

        pm.proxies.insert(
            "a".to_string(),
            ProxyInfo {
                state: ProxyState::Active,
            },
        );
        pm.proxies.insert(
            "b".to_string(),
            ProxyInfo {
                state: ProxyState::Unknown,
            },
        );
        pm.proxies.insert(
            "c".to_string(),
            ProxyInfo {
                state: ProxyState::Dead,
            },
        );
        pm.proxies.insert(
            "d".to_string(),
            ProxyInfo {
                state: ProxyState::Dead,
            },
        );

        let (total, active_or_unknown, cooldown, dead, retired) = pm.proxy_counts();
        assert_eq!(total, 4);
        assert_eq!(active_or_unknown, 2);
        assert_eq!(cooldown, 0);
        assert_eq!(dead, 2);
        assert_eq!(retired, 0);
    }

    #[test]
    fn proxy_counts_empty() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        let (total, active, cooldown, dead, retired) = pm.proxy_counts();
        assert_eq!(total, 0);
        assert_eq!(active, 0);
        assert_eq!(cooldown, 0);
        assert_eq!(dead, 0);
        assert_eq!(retired, 0);
    }

    // -----------------------------------------------------------------------
    // get_client
    // -----------------------------------------------------------------------

    #[test]
    fn get_client_creates_and_caches() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        let client1 = pm.get_client("socks5h://127.0.0.1:1080").unwrap();
        let client2 = pm.get_client("socks5h://127.0.0.1:1080").unwrap();

        // Both should be the same cached client (we can't compare reqwest::Client
        // directly, but we can verify the cache contains one entry).
        assert_eq!(pm.clients.len(), 1);

        // Ensure we actually got clients back.
        drop(client1);
        drop(client2);
    }

    #[test]
    fn get_client_different_proxies_get_different_clients() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        let _c1 = pm.get_client("socks5h://127.0.0.1:1080").unwrap();
        let _c2 = pm.get_client("socks5h://127.0.0.1:1081").unwrap();
        assert_eq!(pm.clients.len(), 2);
    }

    #[test]
    fn get_client_invalid_url_returns_error() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        let result = pm.get_client("not a valid url at all");
        // reqwest may or may not reject this; if it does, we expect an Init error.
        // If reqwest is lenient, the test still passes.
        if let Err(e) = result {
            match e {
                ScatterProxyError::Init(msg) => {
                    assert!(msg.contains("proxy") || msg.contains("invalid"));
                }
                other => panic!("expected Init error, got: {other:?}"),
            }
        }
    }

    // -----------------------------------------------------------------------
    // all_proxy_urls
    // -----------------------------------------------------------------------

    #[test]
    fn all_proxy_urls_returns_all() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        pm.proxies.insert(
            "a".to_string(),
            ProxyInfo {
                state: ProxyState::Active,
            },
        );
        pm.proxies.insert(
            "b".to_string(),
            ProxyInfo {
                state: ProxyState::Dead,
            },
        );

        let urls = pm.all_proxy_urls();
        assert_eq!(urls.len(), 2);
        assert!(urls.contains(&"a".to_string()));
        assert!(urls.contains(&"b".to_string()));
    }

    // -----------------------------------------------------------------------
    // total_count
    // -----------------------------------------------------------------------

    #[test]
    fn total_count_tracks_insertions() {
        let mgr = ProxyManager::new(Duration::from_secs(5));
        assert_eq!(mgr.total_count(), 0);

        mgr.proxies.insert(
            "socks5h://1.1.1.1:1080".into(),
            ProxyInfo {
                state: ProxyState::Active,
            },
        );
        assert_eq!(mgr.total_count(), 1);

        mgr.proxies.insert(
            "socks5h://2.2.2.2:1080".into(),
            ProxyInfo {
                state: ProxyState::Dead,
            },
        );
        assert_eq!(mgr.total_count(), 2);
    }

    // -----------------------------------------------------------------------
    // ProxyState enum
    // -----------------------------------------------------------------------

    #[test]
    fn proxy_state_debug() {
        assert_eq!(format!("{:?}", ProxyState::Unknown), "Unknown");
        assert_eq!(format!("{:?}", ProxyState::Active), "Active");
        assert_eq!(format!("{:?}", ProxyState::Dead), "Dead");
        assert_eq!(format!("{:?}", ProxyState::Retired), "Retired");
    }

    #[test]
    fn proxy_state_clone_and_eq() {
        let s = ProxyState::Active;
        let s2 = s;
        assert_eq!(s, s2);
    }

    // -----------------------------------------------------------------------
    // ProxyState::Retired
    // -----------------------------------------------------------------------

    #[test]
    fn set_state_to_retired() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        pm.proxies.insert(
            "socks5h://1.2.3.4:1080".to_string(),
            ProxyInfo {
                state: ProxyState::Active,
            },
        );
        pm.set_state("socks5h://1.2.3.4:1080", ProxyState::Retired);
        assert_eq!(pm.get_state("socks5h://1.2.3.4:1080"), ProxyState::Retired);
    }

    #[test]
    fn get_active_proxies_excludes_retired() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        pm.proxies.insert(
            "a".to_string(),
            ProxyInfo {
                state: ProxyState::Active,
            },
        );
        pm.proxies.insert(
            "b".to_string(),
            ProxyInfo {
                state: ProxyState::Retired,
            },
        );
        pm.proxies.insert(
            "c".to_string(),
            ProxyInfo {
                state: ProxyState::Unknown,
            },
        );
        let active = pm.get_active_proxies();
        assert_eq!(active.len(), 2);
        assert!(active.contains(&"a".to_string()));
        assert!(!active.contains(&"b".to_string()));
        assert!(active.contains(&"c".to_string()));
    }

    #[test]
    fn proxy_counts_includes_retired() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        pm.proxies.insert(
            "a".to_string(),
            ProxyInfo {
                state: ProxyState::Active,
            },
        );
        pm.proxies.insert(
            "b".to_string(),
            ProxyInfo {
                state: ProxyState::Retired,
            },
        );
        pm.proxies.insert(
            "c".to_string(),
            ProxyInfo {
                state: ProxyState::Dead,
            },
        );

        let (total, active, cooldown, dead, retired) = pm.proxy_counts();
        assert_eq!(total, 3);
        assert_eq!(active, 1);
        assert_eq!(cooldown, 0);
        assert_eq!(dead, 1);
        assert_eq!(retired, 1);
    }

    // -----------------------------------------------------------------------
    // Per-proxy cookie management
    // -----------------------------------------------------------------------

    #[test]
    fn get_cookie_returns_none_when_not_set() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        assert!(pm.get_cookie("socks5h://1.2.3.4:1080").is_none());
    }

    #[test]
    fn set_and_get_cookie() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        pm.set_cookie("socks5h://1.2.3.4:1080", "abc123".to_string());
        assert_eq!(
            pm.get_cookie("socks5h://1.2.3.4:1080"),
            Some("abc123".to_string())
        );
    }

    #[test]
    fn set_cookie_overwrites_existing() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        pm.set_cookie("socks5h://1.2.3.4:1080", "first".to_string());
        pm.set_cookie("socks5h://1.2.3.4:1080", "second".to_string());
        assert_eq!(
            pm.get_cookie("socks5h://1.2.3.4:1080"),
            Some("second".to_string())
        );
    }

    #[test]
    fn cookies_are_per_proxy() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        pm.set_cookie("socks5h://1.1.1.1:1080", "cookie_a".to_string());
        pm.set_cookie("socks5h://2.2.2.2:1080", "cookie_b".to_string());
        assert_eq!(
            pm.get_cookie("socks5h://1.1.1.1:1080"),
            Some("cookie_a".to_string())
        );
        assert_eq!(
            pm.get_cookie("socks5h://2.2.2.2:1080"),
            Some("cookie_b".to_string())
        );
    }

    // -----------------------------------------------------------------------
    // Per-proxy request counting
    // -----------------------------------------------------------------------

    #[test]
    fn request_count_starts_at_zero() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        assert_eq!(pm.request_count("socks5h://1.2.3.4:1080"), 0);
    }

    #[test]
    fn increment_request_count_increments_correctly() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        assert_eq!(pm.increment_request_count("socks5h://1.2.3.4:1080"), 1);
        assert_eq!(pm.increment_request_count("socks5h://1.2.3.4:1080"), 2);
        assert_eq!(pm.increment_request_count("socks5h://1.2.3.4:1080"), 3);
        assert_eq!(pm.request_count("socks5h://1.2.3.4:1080"), 3);
    }

    #[test]
    fn request_counts_are_per_proxy() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        pm.increment_request_count("socks5h://1.1.1.1:1080");
        pm.increment_request_count("socks5h://1.1.1.1:1080");
        pm.increment_request_count("socks5h://2.2.2.2:1080");
        assert_eq!(pm.request_count("socks5h://1.1.1.1:1080"), 2);
        assert_eq!(pm.request_count("socks5h://2.2.2.2:1080"), 1);
    }

    #[test]
    fn reset_request_count_clears_count() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        pm.increment_request_count("socks5h://1.2.3.4:1080");
        pm.increment_request_count("socks5h://1.2.3.4:1080");
        pm.reset_request_count("socks5h://1.2.3.4:1080");
        assert_eq!(pm.request_count("socks5h://1.2.3.4:1080"), 0);
    }

    #[test]
    fn reset_request_count_restores_retired_to_active() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        pm.proxies.insert(
            "socks5h://1.2.3.4:1080".to_string(),
            ProxyInfo {
                state: ProxyState::Retired,
            },
        );
        pm.increment_request_count("socks5h://1.2.3.4:1080");
        pm.reset_request_count("socks5h://1.2.3.4:1080");
        assert_eq!(pm.get_state("socks5h://1.2.3.4:1080"), ProxyState::Active);
        assert_eq!(pm.request_count("socks5h://1.2.3.4:1080"), 0);
    }

    #[test]
    fn reset_request_count_does_not_restore_dead() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        pm.proxies.insert(
            "socks5h://1.2.3.4:1080".to_string(),
            ProxyInfo {
                state: ProxyState::Dead,
            },
        );
        pm.reset_request_count("socks5h://1.2.3.4:1080");
        assert_eq!(pm.get_state("socks5h://1.2.3.4:1080"), ProxyState::Dead);
    }

    // -----------------------------------------------------------------------
    // Challenge lock
    // -----------------------------------------------------------------------

    #[test]
    fn challenge_lock_returns_semaphore_with_one_permit() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        let sem = pm.challenge_lock("socks5h://1.2.3.4:1080");
        assert_eq!(sem.available_permits(), 1);
    }

    #[test]
    fn challenge_lock_same_proxy_returns_same_semaphore() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        let sem1 = pm.challenge_lock("socks5h://1.2.3.4:1080");
        let sem2 = pm.challenge_lock("socks5h://1.2.3.4:1080");
        // Both Arc references point to the same semaphore instance.
        assert!(std::ptr::eq(sem1.as_ref(), sem2.as_ref()));
    }

    // -----------------------------------------------------------------------
    // fetch_and_add (unit tests with no real HTTP server)
    // -----------------------------------------------------------------------

    #[tokio::test]
    async fn fetch_and_add_with_empty_sources() {
        let mgr = ProxyManager::new(Duration::from_secs(5));
        let added = mgr.fetch_and_add(&[], true).await.unwrap();
        assert_eq!(added, 0);
        assert_eq!(mgr.total_count(), 0);
    }

    #[tokio::test]
    async fn fetch_and_add_with_unreachable_source_logs_error() {
        let pm = ProxyManager::new(Duration::from_secs(8));
        // This source won't be reachable but should not cause a hard failure.
        let added = pm
            .fetch_and_add(
                &["http://127.0.0.1:1/nonexistent-proxy-list".to_string()],
                true,
            )
            .await
            .unwrap();
        assert_eq!(added, 0);
        assert_eq!(pm.total_count(), 0);
    }

    // -----------------------------------------------------------------------
    // normalize edge cases
    // -----------------------------------------------------------------------

    #[test]
    fn normalize_comment_with_leading_whitespace() {
        // After trimming, starts with '#'
        let result = ProxyManager::normalize_proxy_url("  # comment", true);
        assert!(result.is_empty());
    }

    #[test]
    fn normalize_ipv6_bare() {
        let result = ProxyManager::normalize_proxy_url("[::1]:1080", true);
        assert_eq!(result, "socks5h://[::1]:1080");
    }

    #[test]
    fn normalize_socks5_uppercase_not_matched_becomes_bare() {
        // "SOCKS5://..." doesn't match the lowercase prefix, so it's treated as bare.
        let result = ProxyManager::normalize_proxy_url("SOCKS5://1.2.3.4:1080", true);
        assert_eq!(result, "socks5h://SOCKS5://1.2.3.4:1080");
    }
}