capp_config/
proxy.rs

1use rand::{rng, seq::IndexedRandom};
2use regex::Regex;
3use serde::{Deserialize, Serialize};
4use std::{
5    fmt,
6    sync::{LazyLock, Mutex},
7};
8use toml::Value;
9
10// Add Debug to the trait bounds
11pub trait ProxyProvider: Send + Sync + fmt::Debug {
12    fn get_proxy(&self) -> Option<String>;
13}
14
15static PORT_RANGE_RE: LazyLock<Regex> = LazyLock::new(|| {
16    Regex::new(r"\{(\d+)-(\d+)\}").expect("Failed to compile port range regex")
17});
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct ProxyConfig {
21    pub use_proxy: bool,
22    pub uris: Vec<String>,
23}
24
25impl ProxyConfig {
26    fn expand_uri(uri: &str) -> Vec<String> {
27        if let Some(captures) = PORT_RANGE_RE.captures(uri) {
28            if let (Some(start), Some(end)) = (
29                captures.get(1).and_then(|m| m.as_str().parse::<u16>().ok()),
30                captures.get(2).and_then(|m| m.as_str().parse::<u16>().ok()),
31            ) {
32                if start <= end {
33                    return (start..=end)
34                        .map(|port| uri.replace(&captures[0], &port.to_string()))
35                        .collect();
36                }
37            }
38        }
39        vec![uri.to_string()]
40    }
41
42    pub fn from_config(config: &Value) -> Option<Self> {
43        let use_proxy = config.get("use")?.as_bool()?;
44        let uris = if use_proxy {
45            if let Some(seq) = config.get("uris").and_then(|v| v.as_array()) {
46                seq.iter()
47                    .filter_map(|v| v.as_str())
48                    .flat_map(Self::expand_uri)
49                    .collect()
50            } else if let Some(uri) = config.get("uri").and_then(|v| v.as_str()) {
51                Self::expand_uri(uri)
52            } else {
53                vec![]
54            }
55        } else {
56            vec![]
57        };
58
59        Some(ProxyConfig { use_proxy, uris })
60    }
61}
62
63#[derive(Debug)]
64pub struct RandomProxyProvider {
65    config: ProxyConfig,
66}
67
68impl RandomProxyProvider {
69    pub fn new(config: ProxyConfig) -> Self {
70        Self { config }
71    }
72
73    pub fn from_config(config: &Value) -> Option<Self> {
74        let config = ProxyConfig::from_config(config)?;
75        Some(Self::new(config))
76    }
77}
78
79impl ProxyProvider for RandomProxyProvider {
80    fn get_proxy(&self) -> Option<String> {
81        if !self.config.use_proxy || self.config.uris.is_empty() {
82            return None;
83        }
84
85        // Create a new thread_rng for each call
86        let mut rng = rng();
87        self.config.uris.choose(&mut rng).cloned()
88    }
89}
90
91// Round robin implementation
92#[derive(Debug)]
93pub struct RoundRobinProxyProvider {
94    config: ProxyConfig,
95    current: Mutex<usize>,
96}
97
98impl RoundRobinProxyProvider {
99    pub fn new(config: ProxyConfig) -> Self {
100        Self {
101            config,
102            current: Mutex::new(0),
103        }
104    }
105
106    pub fn from_config(config: &Value) -> Option<Self> {
107        let config = ProxyConfig::from_config(config)?;
108        Some(Self::new(config))
109    }
110}
111
112impl ProxyProvider for RoundRobinProxyProvider {
113    fn get_proxy(&self) -> Option<String> {
114        if !self.config.use_proxy || self.config.uris.is_empty() {
115            return None;
116        }
117
118        let mut current = self.current.lock().ok()?;
119        let proxy = self.config.uris[*current].clone();
120        *current = (*current + 1) % self.config.uris.len();
121        Some(proxy)
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    const TOML_CONF_SINGLE: &str = r#"
130    [proxy]
131    use = true
132    uri = "http://proxy1.example.com:8080"
133    timeout = 30
134    connect_timeout = 10
135    "#;
136
137    const TOML_CONF_MULTIPLE: &str = r#"
138    [proxy]
139    use = true
140    uris = [
141        "http://proxy1.example.com:8080",
142        "http://proxy2.example.com:8080",
143        "http://proxy3.example.com:8080",
144    ]
145    timeout = 30
146    connect_timeout = 10
147    "#;
148
149    #[test]
150    fn test_proxy_provider() {
151        let config: Value = toml::from_str(TOML_CONF_SINGLE).unwrap();
152        let provider = RandomProxyProvider::from_config(&config["proxy"]).unwrap();
153
154        let proxy_url = provider.get_proxy().unwrap();
155        assert!(proxy_url.as_str() == "http://proxy1.example.com:8080");
156    }
157
158    #[test]
159    fn test_random_proxy_provider() {
160        let config: Value = toml::from_str(TOML_CONF_MULTIPLE).unwrap();
161        let provider = RandomProxyProvider::from_config(&config["proxy"]).unwrap();
162
163        // Test multiple calls return different proxies
164        let mut seen = std::collections::HashSet::new();
165        for _ in 0..10 {
166            if let Some(proxy) = provider.get_proxy() {
167                seen.insert(proxy);
168            }
169        }
170        assert!(seen.len() > 1); // Should see multiple different proxies
171    }
172
173    #[test]
174    fn test_round_robin_proxy_provider() {
175        let config: Value = toml::from_str(TOML_CONF_MULTIPLE).unwrap();
176        let provider =
177            RoundRobinProxyProvider::from_config(&config["proxy"]).unwrap();
178
179        // Test round robin behavior
180        let first = provider.get_proxy();
181        let second = provider.get_proxy();
182        let third = provider.get_proxy();
183        let fourth = provider.get_proxy(); // Should wrap around to first
184
185        assert_ne!(first, second);
186        assert_ne!(second, third);
187        assert_eq!(first, fourth);
188    }
189
190    #[test]
191    fn test_proxy_provider_disabled() {
192        let config: Value = toml::from_str(
193            r#"
194        [proxy]
195        use = false
196        uri = "http://proxy1.example.com:8080"
197        "#,
198        )
199        .unwrap();
200
201        let provider = RandomProxyProvider::from_config(&config["proxy"]).unwrap();
202        assert_eq!(provider.get_proxy(), None);
203    }
204
205    #[test]
206    fn test_proxy_config_port_range() {
207        let toml = r#"
208        [proxy]
209        use = true
210        uri = "http://proxy1.example.com:{8080-8082}"
211        "#;
212
213        let config: Value = toml::from_str(toml).unwrap();
214        let proxy_config = ProxyConfig::from_config(&config["proxy"]).unwrap();
215
216        assert_eq!(proxy_config.uris.len(), 3);
217        assert!(
218            proxy_config
219                .uris
220                .contains(&"http://proxy1.example.com:8080".to_string())
221        );
222        assert!(
223            proxy_config
224                .uris
225                .contains(&"http://proxy1.example.com:8081".to_string())
226        );
227        assert!(
228            proxy_config
229                .uris
230                .contains(&"http://proxy1.example.com:8082".to_string())
231        );
232    }
233
234    #[test]
235    fn test_proxy_config_multiple_uris_with_ranges() {
236        let toml = r#"
237        [proxy]
238        use = true
239        uris = [
240            "http://proxy1.example.com:{8080-8082}",
241            "http://proxy2.example.com:8090",
242            "http://proxy3.example.com:{9000-9001}",
243        ]
244        "#;
245
246        let config: Value = toml::from_str(toml).unwrap();
247        let proxy_config = ProxyConfig::from_config(&config["proxy"]).unwrap();
248
249        assert_eq!(proxy_config.uris.len(), 6); // 3 from first range + 1 single + 2 from last range
250        assert!(
251            proxy_config
252                .uris
253                .contains(&"http://proxy2.example.com:8090".to_string())
254        );
255    }
256
257    #[test]
258    fn test_invalid_port_range() {
259        let toml = r#"
260        [proxy]
261        use = true
262        uri = "http://proxy1.example.com:{8082-8080}"
263        "#;
264
265        let config: Value = toml::from_str(toml).unwrap();
266        let proxy_config = ProxyConfig::from_config(&config["proxy"]).unwrap();
267
268        // Should treat invalid range as literal string
269        assert_eq!(proxy_config.uris.len(), 1);
270        assert_eq!(
271            proxy_config.uris[0],
272            "http://proxy1.example.com:{8082-8080}".to_string()
273        );
274    }
275
276    #[test]
277    fn test_no_port_range() {
278        let toml = r#"
279        [proxy]
280        use = true
281        uri = "http://proxy1.example.com:8080"
282        "#;
283
284        let config: Value = toml::from_str(toml).unwrap();
285        let proxy_config = ProxyConfig::from_config(&config["proxy"]).unwrap();
286
287        assert_eq!(proxy_config.uris.len(), 1);
288        assert_eq!(
289            proxy_config.uris[0],
290            "http://proxy1.example.com:8080".to_string()
291        );
292    }
293}
capp_config/proxy.rs

capp_config/
proxy.rs
