Skip to main content

mocra_proxy/
proxy_impl.rs

1use crate::{IpProvider, IpProxy, IpProxyLoader};
2use async_trait::async_trait;
3use log::info;
4use std::time::Duration;
5
6async fn fetch_text_proxies(config: &IpProvider) -> crate::error::Result<Vec<IpProxy>> {
7    info!("Fetching proxies from {}: {}", config.name, config.url);
8    let client = reqwest::Client::builder()
9        .timeout(Duration::from_secs(config.timeout))
10        .build()
11        .map_err(|e| crate::error::ProxyError::GetProxy(Box::new(e)))?;
12
13    let resp = client
14        .get(&config.url)
15        .send()
16        .await
17        .map_err(|e| crate::error::ProxyError::GetProxy(Box::new(e)))?
18        .text()
19        .await
20        .map_err(|e| crate::error::ProxyError::GetProxy(Box::new(e)))?;
21
22    let mut proxies = Vec::new();
23    for line in resp.lines() {
24        let line = line.trim();
25        if line.is_empty() {
26            continue;
27        }
28        // Parse IP:PORT or IP:PORT:USER:PASS
29        let parts: Vec<&str> = line.split(':').collect();
30        if parts.len() >= 2
31            && let Ok(port) = parts[1].parse::<u16>()
32        {
33            proxies.push(IpProxy {
34                ip: parts[0].to_string(),
35                port,
36                username: if parts.len() >= 4 {
37                    Some(parts[2].to_string())
38                } else {
39                    None
40                },
41                password: if parts.len() >= 4 {
42                    Some(parts[3].to_string())
43                } else {
44                    None
45                },
46                proxy_type: Some("http".to_string()),
47                rate_limit: config.rate_limit,
48            });
49        }
50    }
51    info!("Fetched {} proxies from {}", proxies.len(), config.name);
52    Ok(proxies)
53}
54
55/// Kuaidaili proxy loader.
56pub(crate) struct KuaiDaiLiLoader {
57    pub(crate) config: IpProvider,
58}
59
60#[async_trait]
61impl IpProxyLoader for KuaiDaiLiLoader {
62    async fn get_ip_proxies(&self) -> crate::error::Result<Vec<IpProxy>> {
63        fetch_text_proxies(&self.config).await
64    }
65    fn is_retry_code(&self, code: &u16) -> bool {
66        self.config.retry_codes.contains(code)
67    }
68
69    fn get_name(&self) -> String {
70        self.config.name.clone()
71    }
72
73    fn get_weight(&self) -> u32 {
74        self.config.weight.unwrap_or(1)
75    }
76
77    fn get_config(&self) -> &IpProvider {
78        &self.config
79    }
80
81    async fn health_check(&self, proxy: &IpProxy) -> bool {
82        check_proxy_health(proxy).await
83    }
84}
85
86/// Zhima proxy loader.
87pub(crate) struct ZmLoader {
88    pub(crate) config: IpProvider,
89}
90
91#[async_trait]
92impl IpProxyLoader for ZmLoader {
93    async fn get_ip_proxies(&self) -> crate::error::Result<Vec<IpProxy>> {
94        fetch_text_proxies(&self.config).await
95    }
96
97    fn is_retry_code(&self, code: &u16) -> bool {
98        self.config.retry_codes.contains(code)
99    }
100
101    fn get_name(&self) -> String {
102        self.config.name.clone()
103    }
104
105    fn get_weight(&self) -> u32 {
106        self.config.weight.unwrap_or(1)
107    }
108
109    fn get_config(&self) -> &IpProvider {
110        &self.config
111    }
112
113    async fn health_check(&self, proxy: &IpProxy) -> bool {
114        check_proxy_health(proxy).await
115    }
116}
117
118/// Generic text-list proxy loader (IP:PORT[:USER:PASS])
119pub(crate) struct TextListLoader {
120    pub(crate) config: IpProvider,
121}
122
123#[async_trait]
124impl IpProxyLoader for TextListLoader {
125    async fn get_ip_proxies(&self) -> crate::error::Result<Vec<IpProxy>> {
126        fetch_text_proxies(&self.config).await
127    }
128
129    fn is_retry_code(&self, code: &u16) -> bool {
130        self.config.retry_codes.contains(code)
131    }
132
133    fn get_name(&self) -> String {
134        self.config.name.clone()
135    }
136
137    fn get_weight(&self) -> u32 {
138        self.config.weight.unwrap_or(1)
139    }
140
141    fn get_config(&self) -> &IpProvider {
142        &self.config
143    }
144
145    async fn health_check(&self, proxy: &IpProxy) -> bool {
146        check_proxy_health(proxy).await
147    }
148}
149
150pub(crate) fn build_ip_proxy_loader(config: IpProvider) -> Box<dyn IpProxyLoader> {
151    let name = config.name.to_lowercase();
152    match name.as_str() {
153        "kuaidaili" => Box::new(KuaiDaiLiLoader { config }),
154        "zm" | "zhima" => Box::new(ZmLoader { config }),
155        _ => Box::new(TextListLoader { config }),
156    }
157}
158
159async fn check_proxy_health(proxy: &IpProxy) -> bool {
160    let proxy_url = proxy.to_string();
161    if let Ok(p) = reqwest::Proxy::all(&proxy_url)
162        && let Ok(client) = reqwest::Client::builder()
163            .proxy(p)
164            .timeout(Duration::from_secs(5))
165            .build()
166    {
167        // Use a lightweight check
168        return client.head("http://www.baidu.com").send().await.is_ok();
169    }
170    false
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176    use crate::IpProvider;
177
178    #[test]
179    fn test_kuaidaili_config() {
180        let config = IpProvider {
181            name: "kuaidaili".to_string(),
182            url: "http://example.com".to_string(),
183            retry_codes: vec![429, 503],
184            timeout: 10,
185            rate_limit: 10.0,
186            provider_expire_time: None,
187            proxy_expire_time: 300,
188            weight: Some(10),
189        };
190        let loader = KuaiDaiLiLoader {
191            config: config.clone(),
192        };
193        assert_eq!(loader.get_name(), "kuaidaili");
194        assert_eq!(loader.get_weight(), 10);
195        assert!(loader.is_retry_code(&429));
196        assert!(!loader.is_retry_code(&200));
197    }
198
199    #[test]
200    fn test_zm_config() {
201        let config = IpProvider {
202            name: "zm".to_string(),
203            url: "http://example.com".to_string(),
204            retry_codes: vec![],
205            timeout: 10,
206            rate_limit: 5.0,
207            provider_expire_time: None,
208            proxy_expire_time: 300,
209            weight: None,
210        };
211        let loader = ZmLoader { config };
212        assert_eq!(loader.get_name(), "zm");
213        assert_eq!(loader.get_weight(), 1); // Default
214    }
215}