1use base64::{prelude::BASE64_STANDARD, Engine};
2
3use crate::{
4 error::{Error, ErrorKind},
5 Response,
6};
7
8/// Proxy protocol
9#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
10pub enum Proto {
11 HTTP,
12 SOCKS4,
13 SOCKS4A,
14 SOCKS5,
15}
16
17/// Proxy server definition
18#[derive(Clone, Debug, Eq, Hash, PartialEq)]
19pub struct Proxy {
20 pub(crate) server: String,
21 pub(crate) port: u32,
22 pub(crate) user: Option<String>,
23 pub(crate) password: Option<String>,
24 pub(crate) proto: Proto,
25}
26
27impl Proxy {
28 fn parse_creds<S: AsRef<str>>(
29 creds: &Option<S>,
30 ) -> Result<(Option<String>, Option<String>), Error> {
31 match creds {
32 Some(creds) => {
33 let mut parts = creds
34 .as_ref()
35 .splitn(2, ':')
36 .collect::<Vec<&str>>()
37 .into_iter();
38
39 if parts.len() != 2 {
40 Err(ErrorKind::InvalidProxyUrl.new())
41 } else {
42 Ok((
43 parts.next().map(String::from),
44 parts.next().map(String::from),
45 ))
46 }
47 }
48 None => Ok((None, None)),
49 }
50 }
51
52 fn parse_address<S: AsRef<str>>(host: &Option<S>) -> Result<(String, Option<u32>), Error> {
53 match host {
54 Some(host) => {
55 let mut parts = host.as_ref().split(':').collect::<Vec<&str>>().into_iter();
56 let host = parts
57 .next()
58 .ok_or_else(|| ErrorKind::InvalidProxyUrl.new())?;
59 let port = parts.next();
60 Ok((
61 String::from(host),
62 port.and_then(|port| port.parse::<u32>().ok()),
63 ))
64 }
65 None => Err(ErrorKind::InvalidProxyUrl.new()),
66 }
67 }
68
69 pub(crate) fn use_authorization(&self) -> bool {
70 self.user.is_some() && self.password.is_some()
71 }
72
73 pub(crate) fn try_from_system() -> Option<Self> {
74 macro_rules! try_env {
75 ($($env:literal),+) => {
76 $(
77 if let Ok(env) = std::env::var($env) {
78 if let Ok(proxy) = Self::new(env) {
79 return Some(proxy);
80 }
81 }
82 )+
83 };
84 }
85
86 try_env!(
87 "ALL_PROXY",
88 "all_proxy",
89 "HTTPS_PROXY",
90 "https_proxy",
91 "HTTP_PROXY",
92 "http_proxy"
93 );
94 None
95 }
96
97 /// Create a proxy from a format string.
98 /// # Arguments:
99 /// * `proxy` - a str of format `<protocol>://<user>:<password>@<host>:port` . All parts except host are optional.
100 /// # Protocols
101 /// * `http`: HTTP
102 /// * `socks4`: SOCKS4 (requires socks feature)
103 /// * `socks4a`: SOCKS4A (requires socks feature)
104 /// * `socks5` and `socks`: SOCKS5 (requires socks feature)
105 /// # Examples
106 /// * `http://127.0.0.1:8080`
107 /// * `socks5://john:smith@socks.google.com`
108 /// * `john:smith@socks.google.com:8000`
109 /// * `localhost`
110 pub fn new<S: AsRef<str>>(proxy: S) -> Result<Self, Error> {
111 let mut proxy = proxy.as_ref();
112
113 while proxy.ends_with('/') {
114 proxy = &proxy[..(proxy.len() - 1)];
115 }
116
117 let mut proxy_parts = proxy.splitn(2, "://").collect::<Vec<&str>>().into_iter();
118
119 let proto = if proxy_parts.len() == 2 {
120 match proxy_parts.next() {
121 Some("http") => Proto::HTTP,
122 Some("socks4") => Proto::SOCKS4,
123 Some("socks4a") => Proto::SOCKS4A,
124 Some("socks") => Proto::SOCKS5,
125 Some("socks5") => Proto::SOCKS5,
126 _ => return Err(ErrorKind::InvalidProxyUrl.new()),
127 }
128 } else {
129 Proto::HTTP
130 };
131
132 let remaining_parts = proxy_parts.next();
133 if remaining_parts.is_none() {
134 return Err(ErrorKind::InvalidProxyUrl.new());
135 }
136
137 let mut creds_server_port_parts = remaining_parts
138 .unwrap()
139 .rsplitn(2, '@')
140 .collect::<Vec<&str>>()
141 .into_iter()
142 .rev();
143
144 let (user, password) = if creds_server_port_parts.len() == 2 {
145 Proxy::parse_creds(&creds_server_port_parts.next())?
146 } else {
147 (None, None)
148 };
149
150 let (server, port) = Proxy::parse_address(&creds_server_port_parts.next())?;
151
152 Ok(Self {
153 server,
154 user,
155 password,
156 port: port.unwrap_or(match proto {
157 Proto::HTTP => 80,
158 Proto::SOCKS4 | Proto::SOCKS4A | Proto::SOCKS5 => 1080,
159 }),
160 proto,
161 })
162 }
163
164 pub(crate) fn connect<S: AsRef<str>>(&self, host: S, port: u16, user_agent: &str) -> String {
165 let authorization = if self.use_authorization() {
166 let creds = BASE64_STANDARD.encode(format!(
167 "{}:{}",
168 self.user.clone().unwrap_or_default(),
169 self.password.clone().unwrap_or_default()
170 ));
171
172 match self.proto {
173 Proto::HTTP => format!("Proxy-Authorization: basic {}\r\n", creds),
174 _ => String::new(),
175 }
176 } else {
177 String::new()
178 };
179
180 format!(
181 "CONNECT {}:{} HTTP/1.1\r\n\
182Host: {}:{}\r\n\
183User-Agent: {}\r\n\
184Proxy-Connection: Keep-Alive\r\n\
185{}\
186\r\n",
187 host.as_ref(),
188 port,
189 host.as_ref(),
190 port,
191 user_agent,
192 authorization
193 )
194 }
195
196 pub(crate) fn verify_response(response: &Response) -> Result<(), Error> {
197 match response.status() {
198 200 => Ok(()),
199 401 | 407 => Err(ErrorKind::ProxyUnauthorized.new()),
200 _ => Err(ErrorKind::ProxyConnect.new()),
201 }
202 }
203}
204
205#[cfg(test)]
206mod tests {
207 use super::*;
208
209 #[test]
210 fn parse_proxy_fakeproto() {
211 assert!(Proxy::new("fakeproto://localhost").is_err());
212 }
213
214 #[test]
215 fn parse_proxy_http_user_pass_server_port() {
216 let proxy = Proxy::new("http://user:p@ssw0rd@localhost:9999").unwrap();
217 assert_eq!(proxy.user, Some(String::from("user")));
218 assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
219 assert_eq!(proxy.server, String::from("localhost"));
220 assert_eq!(proxy.port, 9999);
221 assert_eq!(proxy.proto, Proto::HTTP);
222 }
223
224 #[test]
225 fn parse_proxy_http_user_pass_server_port_trailing_slash() {
226 let proxy = Proxy::new("http://user:p@ssw0rd@localhost:9999/").unwrap();
227 assert_eq!(proxy.user, Some(String::from("user")));
228 assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
229 assert_eq!(proxy.server, String::from("localhost"));
230 assert_eq!(proxy.port, 9999);
231 assert_eq!(proxy.proto, Proto::HTTP);
232 }
233
234 #[cfg(feature = "socks-proxy")]
235 #[test]
236 fn parse_proxy_socks4_user_pass_server_port() {
237 let proxy = Proxy::new("socks4://user:p@ssw0rd@localhost:9999").unwrap();
238 assert_eq!(proxy.user, Some(String::from("user")));
239 assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
240 assert_eq!(proxy.server, String::from("localhost"));
241 assert_eq!(proxy.port, 9999);
242 assert_eq!(proxy.proto, Proto::SOCKS4);
243 }
244
245 #[cfg(feature = "socks-proxy")]
246 #[test]
247 fn parse_proxy_socks4a_user_pass_server_port() {
248 let proxy = Proxy::new("socks4a://user:p@ssw0rd@localhost:9999").unwrap();
249 assert_eq!(proxy.user, Some(String::from("user")));
250 assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
251 assert_eq!(proxy.server, String::from("localhost"));
252 assert_eq!(proxy.port, 9999);
253 assert_eq!(proxy.proto, Proto::SOCKS4A);
254 }
255
256 #[cfg(feature = "socks-proxy")]
257 #[test]
258 fn parse_proxy_socks_user_pass_server_port() {
259 let proxy = Proxy::new("socks://user:p@ssw0rd@localhost:9999").unwrap();
260 assert_eq!(proxy.user, Some(String::from("user")));
261 assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
262 assert_eq!(proxy.server, String::from("localhost"));
263 assert_eq!(proxy.port, 9999);
264 assert_eq!(proxy.proto, Proto::SOCKS5);
265 }
266
267 #[cfg(feature = "socks-proxy")]
268 #[test]
269 fn parse_proxy_socks5_user_pass_server_port() {
270 let proxy = Proxy::new("socks5://user:p@ssw0rd@localhost:9999").unwrap();
271 assert_eq!(proxy.user, Some(String::from("user")));
272 assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
273 assert_eq!(proxy.server, String::from("localhost"));
274 assert_eq!(proxy.port, 9999);
275 assert_eq!(proxy.proto, Proto::SOCKS5);
276 }
277
278 #[test]
279 fn parse_proxy_user_pass_server_port() {
280 let proxy = Proxy::new("user:p@ssw0rd@localhost:9999").unwrap();
281 assert_eq!(proxy.user, Some(String::from("user")));
282 assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
283 assert_eq!(proxy.server, String::from("localhost"));
284 assert_eq!(proxy.port, 9999);
285 }
286
287 #[test]
288 fn parse_proxy_server_port() {
289 let proxy = Proxy::new("localhost:9999").unwrap();
290 assert_eq!(proxy.user, None);
291 assert_eq!(proxy.password, None);
292 assert_eq!(proxy.server, String::from("localhost"));
293 assert_eq!(proxy.port, 9999);
294 }
295
296 #[test]
297 fn parse_proxy_server() {
298 let proxy = Proxy::new("localhost").unwrap();
299 assert_eq!(proxy.user, None);
300 assert_eq!(proxy.password, None);
301 assert_eq!(proxy.server, String::from("localhost"));
302 }
303}
304