leyline/core/
proxy_pool.rs

1mod sessions;
2mod state;
3
4use std::sync::{Arc, Mutex, PoisonError};
5use std::time::{Duration, Instant};
6
7use crate::core::block::BlockRules;
8use crate::core::config::ProxyConfig;
9use crate::core::session::Identity;
10
11use sessions::IdentitySessions;
12pub(crate) use state::Lease;
13use state::PoolState;
14
15const DEFAULT_BAN_AFTER: u32 = 3;
16const DEFAULT_BAN_FOR: Duration = Duration::from_secs(60);
17
18#[derive(Debug, Clone)]
19struct Entry {
20    config: ProxyConfig,
21    label: String,
22    identity: Option<Identity>,
23}
24
25#[derive(Debug, Clone)]
26struct Settings {
27    sticky_for: Duration,
28    ban_after: u32,
29    ban_for: Duration,
30    block_rules: Option<BlockRules>,
31}
32
33#[derive(Debug, Clone)]
34pub struct ProxyPool {
35    entries: Arc<[Entry]>,
36    settings: Settings,
37    state: Arc<Mutex<PoolState>>,
38    sessions: Arc<IdentitySessions>,
39}
40
41#[derive(Debug, Clone)]
42#[non_exhaustive]
43pub struct ProxyHealth {
44    pub proxy: String,
45    pub in_use: u32,
46    pub failures: u32,
47    pub banned_until: Option<Instant>,
48}
49
50impl ProxyPool {
51    pub fn new<P: Into<ProxyConfig>>(proxies: impl IntoIterator<Item = P>) -> Self {
52        Self::from_entries(proxies.into_iter().map(|proxy| entry(proxy.into(), None)))
53    }
54
55    pub fn identified<P: Into<ProxyConfig>>(
56        proxies: impl IntoIterator<Item = (P, Identity)>,
57    ) -> Self {
58        Self::from_entries(
59            proxies
60                .into_iter()
61                .map(|(proxy, identity)| entry(proxy.into(), Some(identity))),
62        )
63    }
64
65    pub(crate) fn needs_impersonation(&self) -> bool {
66        self.entries.iter().any(|entry| entry.identity.is_some())
67    }
68
69    pub(crate) fn configs(&self) -> impl Iterator<Item = &ProxyConfig> {
70        self.entries.iter().map(|entry| &entry.config)
71    }
72
73    fn from_entries(entries: impl Iterator<Item = Entry>) -> Self {
74        let entries: Arc<[Entry]> = entries.collect();
75        let state = PoolState::new(entries.len());
76        Self {
77            entries,
78            settings: Settings {
79                sticky_for: Duration::ZERO,
80                ban_after: DEFAULT_BAN_AFTER,
81                ban_for: DEFAULT_BAN_FOR,
82                block_rules: None,
83            },
84            state: Arc::new(Mutex::new(state)),
85            sessions: Arc::default(),
86        }
87    }
88
89    pub fn sticky_for(mut self, duration: Duration) -> Self {
90        self.settings.sticky_for = duration;
91        self
92    }
93
94    pub fn ban_after(mut self, failures: u32) -> Self {
95        self.settings.ban_after = failures.max(1);
96        self
97    }
98
99    pub fn ban_for(mut self, duration: Duration) -> Self {
100        self.settings.ban_for = duration;
101        self
102    }
103
104    pub fn rotate_on_block(mut self, rules: BlockRules) -> Self {
105        match &mut self.settings.block_rules {
106            Some(held) => held.extend(rules),
107            None => self.settings.block_rules = Some(rules),
108        }
109        self
110    }
111
112    #[must_use]
113    pub fn stats(&self) -> Vec<ProxyHealth> {
114        let state = self.lock();
115        self.entries
116            .iter()
117            .zip(state.health())
118            .map(|(entry, health)| ProxyHealth {
119                proxy: entry.label.clone(),
120                in_use: health.in_use,
121                failures: health.failures,
122                banned_until: health.banned_until,
123            })
124            .collect()
125    }
126
127    fn lock(&self) -> std::sync::MutexGuard<'_, PoolState> {
128        self.state.lock().unwrap_or_else(PoisonError::into_inner)
129    }
130}
131
132fn entry(config: ProxyConfig, identity: Option<Identity>) -> Entry {
133    let label = config
134        .primary()
135        .map(crate::util::redact)
136        .unwrap_or_default();
137    Entry {
138        config,
139        label,
140        identity,
141    }
142}
leyline/core/proxy_pool.rs

leyline/core/
proxy_pool.rs
