Skip to main content

aria2_protocol/http/
proxy.rs

1use std::env;
2use tracing::{debug, info};
3
4#[derive(Debug, Clone, PartialEq)]
5pub enum ProxyType {
6    Http,
7    Socks5,
8}
9
10#[derive(Debug, Clone)]
11pub struct ProxyConfig {
12    pub proxy_type: ProxyType,
13    pub host: String,
14    pub port: u16,
15    pub username: Option<String>,
16    pub password: Option<String>,
17}
18
19impl ProxyConfig {
20    pub fn new_http(host: &str, port: u16) -> Self {
21        Self {
22            proxy_type: ProxyType::Http,
23            host: host.to_string(),
24            port,
25            username: None,
26            password: None,
27        }
28    }
29
30    pub fn new_socks5(host: &str, port: u16) -> Self {
31        Self {
32            proxy_type: ProxyType::Socks5,
33            host: host.to_string(),
34            port,
35            username: None,
36            password: None,
37        }
38    }
39
40    pub fn with_auth(mut self, username: &str, password: &str) -> Self {
41        self.username = Some(username.to_string());
42        self.password = Some(password.to_string());
43        self
44    }
45
46    pub fn to_url(&self) -> String {
47        match (&self.username, &self.password) {
48            (Some(user), Some(pass)) => match self.proxy_type {
49                ProxyType::Http => format!("http://{}:{}@{}:{}", user, pass, self.host, self.port),
50                ProxyType::Socks5 => {
51                    format!("socks5://{}:{}@{}:{}", user, pass, self.host, self.port)
52                }
53            },
54            _ => match self.proxy_type {
55                ProxyType::Http => format!("http://{}:{}", self.host, self.port),
56                ProxyType::Socks5 => format!("socks5://{}:{}", self.host, self.port),
57            },
58        }
59    }
60
61    pub fn from_env() -> Option<Self> {
62        if let Ok(http_proxy) = env::var("HTTP_PROXY") {
63            return Self::from_url(&http_proxy);
64        }
65        if let Ok(http_proxy) = env::var("http_proxy") {
66            return Self::from_url(&http_proxy);
67        }
68        if let Ok(https_proxy) = env::var("HTTPS_PROXY") {
69            return Self::from_url(&https_proxy);
70        }
71        if let Ok(https_proxy) = env::var("https_proxy") {
72            return Self::from_url(&https_proxy);
73        }
74        if let Ok(all_proxy) = env::var("ALL_PROXY") {
75            return Self::from_url(&all_proxy);
76        }
77        if let Ok(all_proxy) = env::var("all_proxy") {
78            return Self::from_url(&all_proxy);
79        }
80        None
81    }
82
83    fn from_url(url: &str) -> Option<Self> {
84        let url = url.trim();
85
86        let (proxy_type, rest) = if let Some(r) = url.strip_prefix("socks5://") {
87            (ProxyType::Socks5, r)
88        } else if let Some(r) = url.strip_prefix("http://") {
89            (ProxyType::Http, r)
90        } else if let Some(r) = url.strip_prefix("https://") {
91            (ProxyType::Http, r)
92        } else {
93            (ProxyType::Http, url)
94        };
95
96        let auth_split_pos = rest.find('@');
97        let (auth_part, addr_part) = if let Some(pos) = auth_split_pos {
98            (Some(&rest[..pos]), &rest[pos + 1..])
99        } else {
100            (None, rest)
101        };
102
103        let (username, password) = if let Some(auth) = auth_part {
104            if let Some(colon_pos) = auth.find(':') {
105                (
106                    Some(auth[..colon_pos].to_string()),
107                    Some(auth[colon_pos + 1..].to_string()),
108                )
109            } else {
110                (Some(auth.to_string()), None)
111            }
112        } else {
113            (None, None)
114        };
115
116        let colon_pos = addr_part.rfind(':')?;
117        let host = &addr_part[..colon_pos];
118        let port: u16 = addr_part[colon_pos + 1..].parse().ok()?;
119
120        let config = match proxy_type {
121            ProxyType::Http => Self::new_http(host, port),
122            ProxyType::Socks5 => Self::new_socks5(host, port),
123        };
124
125        if let (Some(user), Some(pass)) = (username, password) {
126            Some(config.with_auth(&user, &pass))
127        } else {
128            Some(config)
129        }
130    }
131
132    pub fn should_bypass(&self, hostname: &str) -> bool {
133        if let Ok(no_proxy) = env::var("NO_PROXY") {
134            for pattern in no_proxy.split(',') {
135                let pattern = pattern.trim();
136                if pattern.is_empty() {
137                    continue;
138                }
139                if pattern == "*" {
140                    return true;
141                }
142                if pattern.starts_with('.') {
143                    if hostname.ends_with(pattern) || hostname == &pattern[1..] {
144                        return true;
145                    }
146                } else if hostname.eq_ignore_ascii_case(pattern) {
147                    return true;
148                }
149            }
150        }
151        if let Ok(no_proxy) = env::var("no_proxy") {
152            for pattern in no_proxy.split(',') {
153                let pattern = pattern.trim();
154                if pattern.is_empty() {
155                    continue;
156                }
157                if pattern == "*" {
158                    return true;
159                }
160                if pattern.starts_with('.') {
161                    if hostname.ends_with(pattern) || hostname == &pattern[1..] {
162                        return true;
163                    }
164                } else if hostname.eq_ignore_ascii_case(pattern) {
165                    return true;
166                }
167            }
168        }
169        false
170    }
171}
172
173pub struct ProxyManager {
174    config: Option<ProxyConfig>,
175}
176
177impl ProxyManager {
178    pub fn new(config: Option<ProxyConfig>) -> Self {
179        if let Some(ref cfg) = config {
180            info!(
181                "代理配置已启用: {}:{} ({:?})",
182                cfg.host, cfg.port, cfg.proxy_type
183            );
184        }
185        Self { config }
186    }
187
188    pub fn from_env() -> Self {
189        let config = ProxyConfig::from_env();
190        Self::new(config)
191    }
192
193    pub fn get_proxy_for_url(&self, url: &str) -> Option<String> {
194        let config = self.config.as_ref()?;
195        if let Some(hostname) = Self::extract_hostname(url)
196            && config.should_bypass(&hostname)
197        {
198            debug!("跳过代理 (NO_PROXY规则匹配): {}", hostname);
199            return None;
200        }
201        Some(config.to_url())
202    }
203
204    pub fn config(&self) -> Option<&ProxyConfig> {
205        self.config.as_ref()
206    }
207
208    fn extract_hostname(url: &str) -> Option<String> {
209        let url = url
210            .strip_prefix("http://")
211            .or_else(|| url.strip_prefix("https://"))
212            .or_else(|| url.strip_prefix("ftp://"))?;
213        let end_pos = url.find('/')?;
214        Some(url[..end_pos].to_string())
215    }
216}
217
218#[cfg(test)]
219mod tests {
220    use super::*;
221
222    #[test]
223    fn test_proxy_config_to_url() {
224        let config = ProxyConfig::new_http("proxy.example.com", 8080);
225        assert_eq!(config.to_url(), "http://proxy.example.com:8080");
226
227        let socks = ProxyConfig::new_socks5("localhost", 1080);
228        assert_eq!(socks.to_url(), "socks5://localhost:1080");
229
230        let with_auth = ProxyConfig::new_http("proxy.example.com", 8080).with_auth("user", "pass");
231        assert_eq!(
232            with_auth.to_url(),
233            "http://user:pass@proxy.example.com:8080"
234        );
235    }
236
237    #[test]
238    fn test_proxy_from_url() {
239        let config = ProxyConfig::from_url("http://proxy.example.com:8080").unwrap();
240        assert_eq!(config.proxy_type, ProxyType::Http);
241        assert_eq!(config.host, "proxy.example.com");
242        assert_eq!(config.port, 8080);
243
244        let with_auth = ProxyConfig::from_url("http://user:pass@proxy.example.com:3128").unwrap();
245        assert_eq!(with_auth.username.as_deref(), Some("user"));
246        assert_eq!(with_auth.password.as_deref(), Some("pass"));
247
248        let socks = ProxyConfig::from_url("socks5://127.0.0.1:1080").unwrap();
249        assert_eq!(socks.proxy_type, ProxyType::Socks5);
250    }
251
252    #[test]
253    fn test_should_bypass() {
254        let config = ProxyConfig::new_http("proxy.example.com", 8080);
255        assert!(!config.should_bypass("example.com"));
256    }
257}