1+ use base64:: engine:: general_purpose;
2+ use base64:: Engine ;
3+ use hyper:: client:: HttpConnector ;
4+ use hyper:: service:: Service ;
5+ use hyper:: Uri ;
6+ use log:: debug;
7+ use std:: env;
8+ use std:: error:: Error ;
9+ use std:: fmt;
10+ use std:: future:: Future ;
11+ use std:: pin:: Pin ;
12+ use std:: task:: { Context , Poll } ;
13+ use tokio:: net:: TcpStream ;
14+ use tokio_socks:: tcp:: Socks5Stream ;
15+
16+ type BoxError = Box < dyn Error + Send + Sync > ;
17+ type BoxFuture < T > = Pin < Box < dyn Future < Output = Result < T , BoxError > > + Send > > ;
18+ type Credentials = ( String , String ) ;
19+
20+ #[ derive( Clone ) ]
21+ pub enum ProxyConnector {
22+ NoProxy ( HttpConnector ) ,
23+ Socks ( String ) ,
24+ Http ( String ) ,
25+ }
26+
27+ #[ derive( Debug ) ]
28+ pub struct ProxyError ( String ) ;
29+
30+ impl fmt:: Display for ProxyError {
31+ fn fmt ( & self , f : & mut fmt:: Formatter < ' _ > ) -> fmt:: Result {
32+ write ! ( f, "Proxy error: {}" , self . 0 )
33+ }
34+ }
35+
36+ impl Error for ProxyError { }
37+
38+ impl Service < Uri > for ProxyConnector {
39+ type Response = TcpStream ;
40+ type Error = BoxError ;
41+ type Future = BoxFuture < Self :: Response > ;
42+
43+ fn poll_ready ( & mut self , cx : & mut Context < ' _ > ) -> Poll < Result < ( ) , Self :: Error > > {
44+ match self {
45+ ProxyConnector :: NoProxy ( connector) => connector. poll_ready ( cx) . map_err ( Into :: into) ,
46+ _ => Poll :: Ready ( Ok ( ( ) ) ) ,
47+ }
48+ }
49+
50+ fn call ( & mut self , uri : Uri ) -> Self :: Future {
51+ let this = self . clone ( ) ;
52+ Box :: pin ( async move {
53+ match this {
54+ ProxyConnector :: NoProxy ( mut connector) => {
55+ connector. call ( uri) . await . map_err ( Into :: into)
56+ }
57+ ProxyConnector :: Socks ( proxy_addr) => handle_socks_connection ( & proxy_addr, & uri) . await ,
58+ ProxyConnector :: Http ( proxy_addr) => handle_http_connection ( & proxy_addr, & uri) . await ,
59+ }
60+ } )
61+ }
62+ }
63+
64+ impl ProxyConnector {
65+ pub fn new ( ) -> Self {
66+ if let Ok ( socks_proxy) = env:: var ( "SOCKS_PROXY" ) {
67+ debug ! ( "Using SOCKS proxy: {}" , socks_proxy) ;
68+ return ProxyConnector :: Socks ( socks_proxy) ;
69+ }
70+
71+ if let Ok ( http_proxy) = env:: var ( "HTTP_PROXY" ) . or_else ( |_| env:: var ( "HTTPS_PROXY" ) ) {
72+ debug ! ( "Using HTTP proxy: {}" , http_proxy) ;
73+ return ProxyConnector :: Http ( http_proxy) ;
74+ }
75+
76+ let mut connector = HttpConnector :: new ( ) ;
77+ connector. enforce_http ( false ) ;
78+ ProxyConnector :: NoProxy ( connector)
79+ }
80+ }
81+
82+ async fn handle_socks_connection ( proxy_addr : & str , uri : & Uri ) -> Result < TcpStream , BoxError > {
83+ let ( host, port, credentials) = parse_proxy_addr ( proxy_addr) ?;
84+ let target_addr = get_target_addr ( uri) ?;
85+
86+ let stream = match credentials {
87+ Some ( ( username, password) ) => {
88+ Socks5Stream :: connect_with_password ( ( host. as_str ( ) , port) , target_addr, & username, & password) . await
89+ }
90+ None => Socks5Stream :: connect ( ( host. as_str ( ) , port) , target_addr) . await ,
91+ } ?;
92+
93+ Ok ( stream. into_inner ( ) )
94+ }
95+
96+ async fn handle_http_connection ( proxy_addr : & str , uri : & Uri ) -> Result < TcpStream , BoxError > {
97+ let ( host, port, credentials) = parse_proxy_addr ( proxy_addr) ?;
98+ let proxy_stream = TcpStream :: connect ( ( host. as_str ( ) , port) ) . await ?;
99+ let target_addr = get_target_addr ( uri) ?;
100+
101+ let connect_req = build_connect_request ( & target_addr, credentials) ?;
102+ write_and_verify_connection ( & proxy_stream, & connect_req) . await ?;
103+
104+ Ok ( proxy_stream)
105+ }
106+
107+ fn build_connect_request ( target_addr : & str , credentials : Option < Credentials > ) -> Result < String , BoxError > {
108+ let mut req = format ! (
109+ "CONNECT {target_addr} HTTP/1.1\r \n \
110+ \r \n \
111+ \r \n "
112+ ) ;
113+
114+ if let Some ( ( username, password) ) = credentials {
115+ let auth = general_purpose:: STANDARD . encode ( format ! ( "{}:{}" , username, password) ) ;
116+ req. push_str ( & format ! ( "Proxy-Authorization: Basic {}\r \n " , auth) ) ;
117+ }
118+
119+ req. push_str ( "\r \n " ) ;
120+ Ok ( req)
121+ }
122+
123+ async fn write_and_verify_connection ( proxy_stream : & TcpStream , connect_req : & str ) -> Result < ( ) , BoxError > {
124+ proxy_stream. writable ( ) . await ?;
125+ proxy_stream. try_write ( connect_req. as_bytes ( ) ) ?;
126+
127+ let mut response = [ 0u8 ; 1024 ] ;
128+ proxy_stream. readable ( ) . await ?;
129+ let n = proxy_stream. try_read ( & mut response) ?;
130+
131+ let response = String :: from_utf8_lossy ( & response[ ..n] ) ;
132+ if !response. starts_with ( "HTTP/1.1 200" ) {
133+ return Err ( Box :: new ( ProxyError ( format ! ( "Proxy CONNECT failed: {}" , response) ) ) ) ;
134+ }
135+
136+ Ok ( ( ) )
137+ }
138+
139+ fn parse_proxy_addr ( addr : & str ) -> Result < ( String , u16 , Option < Credentials > ) , BoxError > {
140+ let uri: Uri = addr. parse ( ) ?;
141+ let host = uri. host ( ) . ok_or ( "Missing proxy host" ) ?. to_string ( ) ;
142+ let port = uri. port_u16 ( ) . unwrap_or_else ( || {
143+ if uri. scheme_str ( ) == Some ( "https" ) { 443 } else { 80 }
144+ } ) ;
145+
146+ let credentials = extract_credentials ( uri. authority ( ) ) ?;
147+ Ok ( ( host, port, credentials) )
148+ }
149+
150+ fn extract_credentials ( authority : Option < & hyper:: http:: uri:: Authority > ) -> Result < Option < Credentials > , BoxError > {
151+ let Some ( authority) = authority else {
152+ return Ok ( None ) ;
153+ } ;
154+
155+ let Some ( credentials) = authority. as_str ( ) . split ( '@' ) . next ( ) else {
156+ return Ok ( None ) ;
157+ } ;
158+
159+ if credentials == authority. as_str ( ) {
160+ return Ok ( None ) ;
161+ }
162+
163+ let creds: Vec < & str > = credentials. split ( ':' ) . collect ( ) ;
164+ if creds. len ( ) == 2 {
165+ Ok ( Some ( ( creds[ 0 ] . to_string ( ) , creds[ 1 ] . to_string ( ) ) ) )
166+ } else {
167+ Ok ( None )
168+ }
169+ }
170+
171+ fn get_target_addr ( uri : & Uri ) -> Result < String , BoxError > {
172+ let host = uri. host ( ) . ok_or ( "Missing target host" ) ?;
173+ let port = uri. port_u16 ( ) . unwrap_or_else ( || {
174+ if uri. scheme_str ( ) == Some ( "https" ) { 443 } else { 80 }
175+ } ) ;
176+ Ok ( format ! ( "{}:{}" , host, port) )
177+ }
0 commit comments