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(8080),
157 proto,
158 })
159 }
160
161 pub(crate) fn connect<S: AsRef<str>>(&self, host: S, port: u16, user_agent: &str) -> String {
162 let authorization = if self.use_authorization() {
163 let creds = BASE64_STANDARD.encode(format!(
164 "{}:{}",
165 self.user.clone().unwrap_or_default(),
166 self.password.clone().unwrap_or_default()
167 ));
168
169 match self.proto {
170 Proto::HTTP => format!("Proxy-Authorization: basic {}\r\n", creds),
171 _ => String::new(),
172 }
173 } else {
174 String::new()
175 };
176
177 format!(
178 "CONNECT {}:{} HTTP/1.1\r\n\
179Host: {}:{}\r\n\
180User-Agent: {}\r\n\
181Proxy-Connection: Keep-Alive\r\n\
182{}\
183\r\n",
184 host.as_ref(),
185 port,
186 host.as_ref(),
187 port,
188 user_agent,
189 authorization
190 )
191 }
192
193 pub(crate) fn verify_response(response: &Response) -> Result<(), Error> {
194 match response.status() {
195 200 => Ok(()),
196 401 | 407 => Err(ErrorKind::ProxyUnauthorized.new()),
197 _ => Err(ErrorKind::ProxyConnect.new()),
198 }
199 }
200}
201
202#[cfg(test)]
203mod tests {
204 use super::*;
205
206 #[test]
207 fn parse_proxy_fakeproto() {
208 assert!(Proxy::new("fakeproto://localhost").is_err());
209 }
210
211 #[test]
212 fn parse_proxy_http_user_pass_server_port() {
213 let proxy = Proxy::new("http://user:p@ssw0rd@localhost:9999").unwrap();
214 assert_eq!(proxy.user, Some(String::from("user")));
215 assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
216 assert_eq!(proxy.server, String::from("localhost"));
217 assert_eq!(proxy.port, 9999);
218 assert_eq!(proxy.proto, Proto::HTTP);
219 }
220
221 #[test]
222 fn parse_proxy_http_user_pass_server_port_trailing_slash() {
223 let proxy = Proxy::new("http://user:p@ssw0rd@localhost:9999/").unwrap();
224 assert_eq!(proxy.user, Some(String::from("user")));
225 assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
226 assert_eq!(proxy.server, String::from("localhost"));
227 assert_eq!(proxy.port, 9999);
228 assert_eq!(proxy.proto, Proto::HTTP);
229 }
230
231 #[cfg(feature = "socks-proxy")]
232 #[test]
233 fn parse_proxy_socks4_user_pass_server_port() {
234 let proxy = Proxy::new("socks4://user:p@ssw0rd@localhost:9999").unwrap();
235 assert_eq!(proxy.user, Some(String::from("user")));
236 assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
237 assert_eq!(proxy.server, String::from("localhost"));
238 assert_eq!(proxy.port, 9999);
239 assert_eq!(proxy.proto, Proto::SOCKS4);
240 }
241
242 #[cfg(feature = "socks-proxy")]
243 #[test]
244 fn parse_proxy_socks4a_user_pass_server_port() {
245 let proxy = Proxy::new("socks4a://user:p@ssw0rd@localhost:9999").unwrap();
246 assert_eq!(proxy.user, Some(String::from("user")));
247 assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
248 assert_eq!(proxy.server, String::from("localhost"));
249 assert_eq!(proxy.port, 9999);
250 assert_eq!(proxy.proto, Proto::SOCKS4A);
251 }
252
253 #[cfg(feature = "socks-proxy")]
254 #[test]
255 fn parse_proxy_socks_user_pass_server_port() {
256 let proxy = Proxy::new("socks://user:p@ssw0rd@localhost:9999").unwrap();
257 assert_eq!(proxy.user, Some(String::from("user")));
258 assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
259 assert_eq!(proxy.server, String::from("localhost"));
260 assert_eq!(proxy.port, 9999);
261 assert_eq!(proxy.proto, Proto::SOCKS5);
262 }
263
264 #[cfg(feature = "socks-proxy")]
265 #[test]
266 fn parse_proxy_socks5_user_pass_server_port() {
267 let proxy = Proxy::new("socks5://user:p@ssw0rd@localhost:9999").unwrap();
268 assert_eq!(proxy.user, Some(String::from("user")));
269 assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
270 assert_eq!(proxy.server, String::from("localhost"));
271 assert_eq!(proxy.port, 9999);
272 assert_eq!(proxy.proto, Proto::SOCKS5);
273 }
274
275 #[test]
276 fn parse_proxy_user_pass_server_port() {
277 let proxy = Proxy::new("user:p@ssw0rd@localhost:9999").unwrap();
278 assert_eq!(proxy.user, Some(String::from("user")));
279 assert_eq!(proxy.password, Some(String::from("p@ssw0rd")));
280 assert_eq!(proxy.server, String::from("localhost"));
281 assert_eq!(proxy.port, 9999);
282 }
283
284 #[test]
285 fn parse_proxy_server_port() {
286 let proxy = Proxy::new("localhost:9999").unwrap();
287 assert_eq!(proxy.user, None);
288 assert_eq!(proxy.password, None);
289 assert_eq!(proxy.server, String::from("localhost"));
290 assert_eq!(proxy.port, 9999);
291 }
292
293 #[test]
294 fn parse_proxy_server() {
295 let proxy = Proxy::new("localhost").unwrap();
296 assert_eq!(proxy.user, None);
297 assert_eq!(proxy.password, None);
298 assert_eq!(proxy.server, String::from("localhost"));
299 }
300}
301