cloudscraper_rs/modules/proxy/
mod.rs

1//! Proxy rotation and health tracking utilities.
2//!
3//! Tracks proxy performance, bans unhealthy endpoints, and selects the next
4//! candidate based on the chosen rotation strategy.
5
6use rand::Rng;
7use rand::seq::SliceRandom;
8use std::cmp::Ordering;
9use std::collections::HashMap;
10use std::time::{Duration, Instant};
11
12use crate::challenges::solvers::access_denied::ProxyPool;
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum RotationStrategy {
16    Sequential,
17    Random,
18    Smart,
19    Weighted,
20    RoundRobinSmart,
21}
22
23#[derive(Debug, Clone)]
24pub struct ProxyConfig {
25    pub rotation_strategy: RotationStrategy,
26    pub ban_time: Duration,
27    pub failure_threshold: u32,
28    pub cooldown: Duration,
29}
30
31impl Default for ProxyConfig {
32    fn default() -> Self {
33        Self {
34            rotation_strategy: RotationStrategy::Sequential,
35            ban_time: Duration::from_secs(300),
36            failure_threshold: 3,
37            cooldown: Duration::from_secs(60),
38        }
39    }
40}
41
42#[derive(Debug, Clone)]
43pub struct ProxyHealthReport {
44    pub total_proxies: usize,
45    pub available_proxies: usize,
46    pub banned_proxies: usize,
47    pub details: HashMap<String, ProxyStats>,
48}
49
50#[derive(Debug, Clone, Default)]
51pub struct ProxyStats {
52    pub successes: u64,
53    pub failures: u64,
54    pub last_used: Option<Instant>,
55    pub last_failure: Option<Instant>,
56}
57
58#[derive(Debug, Clone)]
59struct ProxyEntry {
60    endpoint: String,
61    stats: ProxyStats,
62    banned_until: Option<Instant>,
63}
64
65impl ProxyEntry {
66    fn is_available(&self) -> bool {
67        match self.banned_until {
68            Some(until) => Instant::now() >= until,
69            None => true,
70        }
71    }
72
73    fn score(&self) -> f64 {
74        let total = self.stats.successes + self.stats.failures;
75        let success_rate = if total == 0 {
76            1.0
77        } else {
78            self.stats.successes as f64 / total as f64
79        };
80        let recency = self
81            .stats
82            .last_used
83            .map(|ts| (Instant::now() - ts).as_secs_f64())
84            .unwrap_or(300.0)
85            / 300.0;
86        (success_rate * 0.7) + (recency.clamp(0.0, 1.0) * 0.3)
87    }
88}
89
90/// Proxy manager with rotation policies.
91#[derive(Debug)]
92pub struct ProxyManager {
93    config: ProxyConfig,
94    proxies: Vec<ProxyEntry>,
95    current_index: usize,
96    rng: rand::rngs::ThreadRng,
97}
98
99impl ProxyManager {
100    pub fn new(config: ProxyConfig) -> Self {
101        Self {
102            config,
103            proxies: Vec::new(),
104            current_index: 0,
105            rng: rand::thread_rng(),
106        }
107    }
108
109    pub fn load<I>(&mut self, proxies: I)
110    where
111        I: IntoIterator,
112        I::Item: Into<String>,
113    {
114        self.proxies.clear();
115        for proxy in proxies {
116            self.add_proxy(proxy);
117        }
118    }
119
120    pub fn add_proxy(&mut self, proxy: impl Into<String>) {
121        let endpoint = proxy.into();
122        if self.proxies.iter().any(|entry| entry.endpoint == endpoint) {
123            return;
124        }
125        self.proxies.push(ProxyEntry {
126            endpoint,
127            stats: ProxyStats::default(),
128            banned_until: None,
129        });
130    }
131
132    pub fn remove_proxy(&mut self, proxy: &str) {
133        self.proxies.retain(|entry| entry.endpoint != proxy);
134    }
135
136    pub fn next_proxy(&mut self) -> Option<String> {
137        if self.proxies.is_empty() {
138            return None;
139        }
140
141        let now = Instant::now();
142        let mut available_indices = Vec::new();
143        for idx in 0..self.proxies.len() {
144            let entry = &mut self.proxies[idx];
145            if let Some(until) = entry.banned_until {
146                if until <= now {
147                    entry.banned_until = None;
148                    available_indices.push(idx);
149                }
150            } else {
151                available_indices.push(idx);
152            }
153        }
154
155        let selected_index = if available_indices.is_empty() {
156            let index = self
157                .proxies
158                .iter()
159                .enumerate()
160                .min_by_key(|(_, entry)| entry.banned_until.unwrap_or(now))
161                .map(|(idx, _)| idx)?;
162            let entry = &mut self.proxies[index];
163            entry.banned_until = None;
164            index
165        } else {
166            match self.config.rotation_strategy {
167                RotationStrategy::Sequential => {
168                    let idx_in_pool = self.current_index % available_indices.len();
169                    self.current_index = (self.current_index + 1) % available_indices.len();
170                    available_indices[idx_in_pool]
171                }
172                RotationStrategy::Random => {
173                    available_indices.choose(&mut self.rng).copied().unwrap()
174                }
175                RotationStrategy::Smart => *available_indices
176                    .iter()
177                    .max_by(|&&a, &&b| {
178                        let lhs = self.proxies[a].score();
179                        let rhs = self.proxies[b].score();
180                        lhs.partial_cmp(&rhs).unwrap_or(Ordering::Equal)
181                    })
182                    .unwrap(),
183                RotationStrategy::Weighted => {
184                    weighted_choice_index(&mut self.rng, &self.proxies, &available_indices)
185                        .unwrap_or(available_indices[0])
186                }
187                RotationStrategy::RoundRobinSmart => {
188                    let filtered: Vec<usize> = available_indices
189                        .iter()
190                        .copied()
191                        .filter(|&idx| {
192                            if let Some(last_failure) = self.proxies[idx].stats.last_failure {
193                                now.duration_since(last_failure) > self.config.cooldown
194                            } else {
195                                true
196                            }
197                        })
198                        .collect();
199                    let pool = if filtered.is_empty() {
200                        &available_indices
201                    } else {
202                        &filtered
203                    };
204                    let idx_in_pool = self.current_index % pool.len();
205                    self.current_index = (self.current_index + 1) % pool.len();
206                    pool[idx_in_pool]
207                }
208            }
209        };
210
211        let entry = &mut self.proxies[selected_index];
212        entry.stats.last_used = Some(Instant::now());
213        Some(entry.endpoint.clone())
214    }
215
216    pub fn report_success(&mut self, proxy: &str) {
217        if let Some(entry) = self
218            .proxies
219            .iter_mut()
220            .find(|entry| entry.endpoint == proxy)
221        {
222            entry.stats.successes += 1;
223            entry.banned_until = None;
224        }
225    }
226
227    pub fn report_failure(&mut self, proxy: &str) {
228        if let Some(entry) = self
229            .proxies
230            .iter_mut()
231            .find(|entry| entry.endpoint == proxy)
232        {
233            entry.stats.failures += 1;
234            entry.stats.last_failure = Some(Instant::now());
235            if entry.stats.failures % self.config.failure_threshold as u64 == 0 {
236                entry.banned_until = Some(Instant::now() + self.config.ban_time);
237            }
238        }
239    }
240
241    pub fn health_report(&self) -> ProxyHealthReport {
242        let mut details = HashMap::new();
243        let mut available = 0;
244        let mut banned = 0;
245        for entry in &self.proxies {
246            if entry.is_available() {
247                available += 1;
248            } else {
249                banned += 1;
250            }
251            details.insert(entry.endpoint.clone(), entry.stats.clone());
252        }
253
254        ProxyHealthReport {
255            total_proxies: self.proxies.len(),
256            available_proxies: available,
257            banned_proxies: banned,
258            details,
259        }
260    }
261}
262
263impl Default for ProxyManager {
264    fn default() -> Self {
265        Self::new(ProxyConfig::default())
266    }
267}
268
269impl ProxyPool for ProxyManager {
270    fn report_failure(&mut self, proxy: &str) {
271        ProxyManager::report_failure(self, proxy);
272    }
273
274    fn next_proxy(&mut self) -> Option<String> {
275        ProxyManager::next_proxy(self)
276    }
277}
278
279fn weighted_choice_index(
280    rng: &mut rand::rngs::ThreadRng,
281    proxies: &[ProxyEntry],
282    indices: &[usize],
283) -> Option<usize> {
284    if indices.is_empty() {
285        return None;
286    }
287
288    let weights: Vec<f64> = indices
289        .iter()
290        .map(|&idx| proxies[idx].score().max(0.1))
291        .collect();
292    let total: f64 = weights.iter().sum();
293    if total <= f64::EPSILON {
294        return indices.choose(rng).copied();
295    }
296
297    let mut target = rng.gen_range(0.0..total);
298    for (index, weight) in indices.iter().zip(weights.iter()) {
299        if target <= *weight {
300            return Some(*index);
301        }
302        target -= *weight;
303    }
304
305    indices.last().copied()
306}
307
308#[cfg(test)]
309mod tests {
310    use super::*;
311
312    #[test]
313    fn rotates_proxies() {
314        let mut manager = ProxyManager::default();
315        manager.load(["http://1.1.1.1:8080", "http://2.2.2.2:8080"]);
316        let first = manager.next_proxy().unwrap();
317        let second = manager.next_proxy().unwrap();
318        assert!(!first.is_empty());
319        assert!(!second.is_empty());
320    }
321
322    #[test]
323    fn bans_after_failures() {
324        let mut manager = ProxyManager::new(ProxyConfig {
325            failure_threshold: 1,
326            ban_time: Duration::from_secs(60),
327            ..Default::default()
328        });
329        manager.add_proxy("http://1.1.1.1:8080");
330        let proxy = manager.next_proxy().unwrap();
331        manager.report_failure(&proxy);
332        let report = manager.health_report();
333        assert_eq!(report.banned_proxies, 1);
334    }
335}