wmhttp/
proxy.rs

1// Copyright 2022 - 2023 Wenmeng See the COPYRIGHT
2// file at the top-level directory of this distribution.
3//
4// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
5// http://www.apache.org/licenses/LICENSE-2.0>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8//
9// Author: tickbh
10// -----
11// Created Date: 2023/12/07 03:05:04
12
13use algorithm::buf::{BinaryMut, BtMut};
14use lazy_static::lazy_static;
15use std::{net::SocketAddr, env, collections::HashSet, fmt::Display};
16
17use tokio::{net::TcpStream, io::{AsyncRead, AsyncWrite}};
18use webparse::{Url, HeaderValue, Scheme};
19
20use crate::{ProtError, ProtResult, RecvRequest};
21
22
23
24/// 客户端代理类
25#[derive(Debug, Clone)]
26pub enum ProxyScheme {
27    Http {
28        addr: SocketAddr,
29        auth: Option<HeaderValue>,
30    },
31    Https {
32        addr: SocketAddr,
33        auth: Option<HeaderValue>,
34    },
35    Socks5 {
36        addr: SocketAddr,
37        auth: Option<(String, String)>,
38    },
39}
40
41async fn tunnel<T>(
42    mut conn: T,
43    host: String,
44    port: u16,
45    user_agent: &Option<HeaderValue>,
46    auth: &Option<HeaderValue>,
47) -> ProtResult<T>
48where
49    T: AsyncRead + AsyncWrite + Unpin,
50{
51    use tokio::io::{AsyncReadExt, AsyncWriteExt};
52
53    let mut buf = format!(
54        "\
55         CONNECT {0}:{1} HTTP/1.1\r\n\
56         Host: {0}:{1}\r\n\
57         ",
58        host, port
59    )
60    .into_bytes();
61
62    // user-agent
63    if let Some(user_agent) = user_agent {
64        buf.extend_from_slice(b"User-Agent: ");
65        buf.extend_from_slice(user_agent.as_bytes());
66        buf.extend_from_slice(b"\r\n");
67    }
68
69    // proxy-authorization
70    if let Some(value) = auth {
71        log::debug!("建立连接 {}:{} 使用基础加密", host, port);
72        buf.extend_from_slice(b"Proxy-Authorization: ");
73        buf.extend_from_slice(value.as_bytes());
74        buf.extend_from_slice(b"\r\n");
75    }
76
77    // headers end
78    buf.extend_from_slice(b"\r\n");
79
80    conn.write_all(&buf).await?;
81
82    let mut buf = [0; 8192];
83    let mut pos = 0;
84
85    loop {
86        let n = conn.read(&mut buf[pos..]).await?;
87
88        if n == 0 {
89            return Err(ProtError::Extension("eof error"));
90        }
91        pos += n;
92
93        let recvd = &buf[..pos];
94        if recvd.starts_with(b"HTTP/1.1 200") || recvd.starts_with(b"HTTP/1.0 200") {
95            if recvd.ends_with(b"\r\n\r\n") {
96                return Ok(conn);
97            }
98            if pos == buf.len() {
99                return Err(ProtError::Extension("proxy headers too long for tunnel"));
100            }
101        // else read more
102        } else if recvd.starts_with(b"HTTP/1.1 407") {
103            return Err(ProtError::Extension("proxy authentication required"));
104        } else {
105            return Err(ProtError::Extension("unsuccessful tunnel"));
106        }
107    }
108}
109
110pub fn basic_auth(auth: &Option<(String, String)>) -> Option<HeaderValue>
111{
112    use base64::prelude::BASE64_STANDARD;
113    use base64::write::EncoderWriter;
114    use std::io::Write;
115    if auth.is_none() {
116        return  None;
117    }
118
119    let mut buf = b"Basic ".to_vec();
120    {
121        let mut encoder = EncoderWriter::new(&mut buf, &BASE64_STANDARD);
122        let _ = write!(encoder, "{}:{}", auth.as_ref().unwrap().0, auth.as_ref().unwrap().1);
123    }
124    let header = HeaderValue::from_bytes(&buf);
125    Some(header)
126}
127
128fn insert_from_env(proxies: &mut Vec<ProxyScheme>, scheme: Scheme, key: &str) -> bool {
129    if let Ok(val) = env::var(key) {
130        if let Ok(proxy) = ProxyScheme::try_from(&*val) {
131            if scheme.is_http() {
132                if let Ok(proxy) = proxy.trans_http() {
133                    proxies.push(proxy);
134                    return true;
135                }
136            } else {
137                if let Ok(proxy) = proxy.trans_https() {
138                    proxies.push(proxy);
139                    return true;
140                }
141            }
142        }
143    }
144    false
145}
146
147fn get_from_environment() -> Vec<ProxyScheme> {
148    let mut proxies = vec![];
149
150    if !insert_from_env(&mut proxies, Scheme::Http, "HTTP_PROXY") {
151        insert_from_env(&mut proxies, Scheme::Http, "http_proxy");
152    }
153
154    if !insert_from_env(&mut proxies, Scheme::Https, "HTTPS_PROXY") {
155        insert_from_env(&mut proxies, Scheme::Https, "https_proxy");
156    }
157
158    if !(insert_from_env(&mut proxies, Scheme::Http, "ALL_PROXY")
159        && insert_from_env(&mut proxies, Scheme::Https, "ALL_PROXY"))
160    {
161        insert_from_env(&mut proxies, Scheme::Http, "all_proxy");
162        insert_from_env(&mut proxies, Scheme::Https, "all_proxy");
163    }
164
165    proxies
166}
167
168impl ProxyScheme {
169
170    pub fn get_env_proxies() -> &'static Vec<ProxyScheme> {
171        lazy_static! {
172            static ref ENV_PROXIES: Vec<ProxyScheme> = get_from_environment();
173        }
174        &ENV_PROXIES
175    }
176
177    pub fn get_env_no_proxy() -> &'static HashSet<String> {
178        lazy_static! {
179            static ref ENV_NO_PROXY: HashSet<String> = {
180                let mut hash = HashSet::new();
181                hash.insert("localhost".to_string());
182                hash.insert("127.0.0.1".to_string());
183                hash.insert("::1".to_string());
184                fn insert_no_proxy(all_hash: &mut HashSet<String>, key: &str) -> bool {
185                    if let Ok(val) = env::var(key) {
186                        let all = val.split(",").collect::<Vec<&str>>();
187                        for one in all {
188                            all_hash.insert(one.trim().to_string());
189                        }
190                        return true
191                    }
192                    false
193                }
194                if !insert_no_proxy(&mut hash, "NO_PROXY") {
195                    insert_no_proxy(&mut hash, "no_proxy");
196                }
197                hash
198            };
199        }
200        &ENV_NO_PROXY
201    }
202
203    pub fn is_no_proxy(host: &String) -> bool {
204        let hash = Self::get_env_no_proxy();
205        hash.contains(host)
206    }
207
208    pub fn fix_request(&self, req: &mut RecvRequest) -> ProtResult<()> {
209        match self {
210            ProxyScheme::Http {addr: _, auth} => {
211                if auth.is_some() {
212                    req.headers_mut().insert("Proxy-Authorization", auth.clone().unwrap());
213                }
214            },
215            _ => {}
216
217        }
218        Ok(())
219    }
220
221    pub async fn connect(&self, url:&Url) -> ProtResult<Option<TcpStream>> {
222        log::trace!("客户端访问\"{}\", 尝试通过代理\"{}\"", url, self);
223        match self {
224            ProxyScheme::Http {addr, auth} => {
225                let tcp = TcpStream::connect(addr).await?;
226                if url.scheme.is_https() {
227                    let tcp = tunnel(tcp, url.domain.clone().unwrap_or_default(), url.port.unwrap_or(443), &None, &auth).await?;
228                    return Ok(Some(tcp));
229                } else {
230                    return Ok(Some(tcp));
231                }
232            },
233            ProxyScheme::Https {addr, auth }  => {
234                if !url.scheme.is_https() {
235                    return Ok(None);
236                }
237                let tcp = TcpStream::connect(addr).await?;
238                let tcp = tunnel(tcp, url.domain.clone().unwrap_or_default(), url.port.unwrap_or(443), &None, &auth).await?;
239                return Ok(Some(tcp));
240            },
241            ProxyScheme::Socks5 { addr, auth } => {
242                let tcp = TcpStream::connect(addr).await?;
243                let tcp = Self::socks5_connect(tcp, &url, auth).await?;
244                return Ok(Some(tcp))
245            },
246        }
247    }
248
249    fn trans_http(self) -> ProtResult<Self> {
250        match self {
251            ProxyScheme::Http { addr, auth } => {
252                Ok(ProxyScheme::Http { addr, auth })
253            }
254            ProxyScheme::Https { addr, auth } => {
255                Ok(ProxyScheme::Http { addr, auth })
256            }
257            _ => {
258                Err(ProtError::Extension("unknow type"))
259            }
260        }
261    }
262
263
264    fn trans_https(self) -> ProtResult<Self> {
265        match self {
266            ProxyScheme::Http { addr, auth } => {
267                Ok(ProxyScheme::Https { addr, auth })
268            }
269            ProxyScheme::Https { addr, auth } => {
270                Ok(ProxyScheme::Https { addr, auth })
271            }
272            _ => {
273                Err(ProtError::Extension("unknow type"))
274            }
275        }
276    }
277
278    async fn socks5_connect<T>(
279        mut conn: T,
280        url: &Url,
281        auth: &Option<(String, String)>,
282    ) -> ProtResult<T>
283    where
284        T: AsyncRead + AsyncWrite + Unpin,
285    {
286        use tokio::io::{AsyncReadExt, AsyncWriteExt};
287        let mut binary = BinaryMut::new();
288        let mut data = vec![0;1024];
289        if let Some(_auth) = auth {
290            conn.write_all(&[5, 1, 2]).await?;
291        } else {
292            conn.write_all(&[5, 0]).await?;
293        }
294
295        conn.read_exact(&mut data[..2]).await?;
296        if data[0] != 5 {
297            return Err(ProtError::Extension("socks5 error"));
298        }
299        match data[1] {
300            2 => {
301                let (user, pass) = auth.as_ref().unwrap();
302                binary.put_u8(1);
303                binary.put_u8(user.as_bytes().len() as u8);
304                binary.put_slice(user.as_bytes());
305                binary.put_u8(pass.as_bytes().len() as u8);
306                binary.put_slice(pass.as_bytes());
307                conn.write_all(binary.as_slice()).await?;
308
309                conn.read_exact(&mut data[..2]).await?;
310                if data[0] != 1 || data[1] != 0 {
311                    return Err(ProtError::Extension("user password error"));
312                }
313
314                binary.clear();
315            }
316            0 => {},
317            _ => {
318                return Err(ProtError::Extension("no method for auth"));
319            }
320        }
321
322        binary.put_slice(&[5, 1, 0, 3]);
323        let domain = url.domain.as_ref().unwrap();
324        let port = url.port.unwrap_or(80);
325        binary.put_u8(domain.as_bytes().len() as u8);
326        binary.put_slice(domain.as_bytes());
327        binary.put_u16(port);
328        conn.write_all(&binary.as_slice()).await?;
329        conn.read_exact(&mut data[..10]).await?;
330        if data[0] != 5 {
331            return Err(ProtError::Extension("socks5 error"));
332        }
333        if data[1] != 0 {
334            return Err(ProtError::Extension("network error"));
335        }
336        Ok(conn)
337    }
338}
339
340impl TryFrom<&str> for ProxyScheme {
341    type Error = ProtError;
342
343    fn try_from(value: &str) -> Result<Self, Self::Error> {
344        let url = Url::try_from(value)?;
345        let (addr, auth) = if let Some(connect) = url.get_connect_url() {
346            let addr = connect
347                .parse::<SocketAddr>()
348                .map_err(|_| ProtError::Extension("unknow parse"))?;
349            let auth = if url.username.is_some() && url.password.is_some() {
350                Some((url.username.unwrap(), url.password.unwrap()))
351            } else {
352                None
353            };
354            (addr, auth)
355        } else {
356            return Err(ProtError::Extension("unknow addr"))
357        };
358        match &url.scheme {
359            webparse::Scheme::Http => Ok(ProxyScheme::Http {
360                addr, auth: basic_auth(&auth)
361            }),
362            webparse::Scheme::Https => Ok(ProxyScheme::Https {
363                addr, auth: basic_auth(&auth)
364            }),
365            webparse::Scheme::Extension(s) if s == "socks5" => {
366                Ok(ProxyScheme::Socks5 { addr, auth })
367            }
368            _ => Err(ProtError::Extension("unknow scheme")),
369        }
370    }
371}
372
373impl Display for ProxyScheme {
374    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
375        match self {
376            ProxyScheme::Http {addr, auth : _} => {
377                f.write_fmt(format_args!("HTTP {}", addr))
378            },
379            ProxyScheme::Https {addr, auth }  => {
380                if auth.is_none() {
381                    f.write_fmt(format_args!("HTTPS {}", addr))
382                } else {
383                    f.write_fmt(format_args!("HTTPS {}, Auth: {}", addr, auth.as_ref().unwrap()))
384                }
385            },
386            ProxyScheme::Socks5 { addr, auth } => {
387                if auth.is_none() {
388                    f.write_fmt(format_args!("SOCKS5 {}", addr))
389                } else {
390                    f.write_fmt(format_args!("SOCKS5 {}, Auth: {}, {}", addr, auth.as_ref().unwrap().0, auth.as_ref().unwrap().1))
391                }
392            },
393        }
394    }
395}