rustenium_identity/
local_proxy_server.rs

1use hyper::body::Incoming;
2use hyper::server::conn::http1;
3use hyper::service::service_fn;
4use hyper::{Method, Request, Response, StatusCode};
5use hyper_util::rt::TokioIo;
6use std::io;
7use std::net::SocketAddr;
8use std::sync::Arc;
9use tokio::io::{AsyncReadExt, AsyncWriteExt};
10use tokio::net::{TcpListener, TcpStream};
11use base64::Engine;
12use tracing::{debug, error, info};
13
14// ---------------------------------------------------------------------------
15// Config
16// ---------------------------------------------------------------------------
17
18#[derive(Debug, Clone)]
19struct ProxyConfig {
20    bind_addr: SocketAddr,
21    upstream_proxy: Option<String>,
22    upstream_username: Option<String>,
23    upstream_password: Option<String>,
24}
25
26impl Default for ProxyConfig {
27    fn default() -> Self {
28        Self {
29            bind_addr: "127.0.0.1:0".parse().unwrap(),
30            upstream_proxy: None,
31            upstream_username: None,
32            upstream_password: None,
33        }
34    }
35}
36
37impl ProxyConfig {
38    fn from_proxy_url(url: &str) -> Result<Self, String> {
39        let rest = url
40            .strip_prefix("https://")
41            .or_else(|| url.strip_prefix("http://"))
42            .ok_or_else(|| format!("Unsupported or missing scheme in proxy URL: {url}"))?;
43
44        let (creds, hostport) = if let Some(at) = rest.rfind('@') {
45            (Some(&rest[..at]), &rest[at + 1..])
46        } else {
47            (None, rest)
48        };
49
50        let (username, password) = match creds {
51            Some(c) => {
52                if let Some((u, p)) = c.split_once(':') {
53                    (Some(u.to_string()), Some(p.to_string()))
54                } else {
55                    (Some(c.to_string()), None)
56                }
57            }
58            None => (None, None),
59        };
60
61        if hostport.is_empty() {
62            return Err("Missing host in proxy URL".into());
63        }
64
65        Ok(Self {
66            bind_addr: "127.0.0.1:0".parse().unwrap(),
67            upstream_proxy: Some(hostport.to_string()),
68            upstream_username: username,
69            upstream_password: password,
70        })
71    }
72
73}
74
75// ---------------------------------------------------------------------------
76// Server
77// ---------------------------------------------------------------------------
78
79struct ProxyServer {
80    config: Arc<ProxyConfig>,
81}
82
83impl ProxyServer {
84    fn new(config: ProxyConfig) -> Self {
85        Self { config: Arc::new(config) }
86    }
87
88    async fn start_background(self) -> io::Result<SocketAddr> {
89        let listener = TcpListener::bind(self.config.bind_addr).await?;
90        let addr = listener.local_addr()?;
91        info!("Local proxy overlay listening on {}", addr);
92        tokio::spawn(async move {
93            if let Err(e) = self.run(listener).await {
94                error!("Proxy server exited: {}", e);
95            }
96        });
97        Ok(addr)
98    }
99
100    async fn run(&self, listener: TcpListener) -> Result<(), Box<dyn std::error::Error>> {
101        loop {
102            let (stream, client_addr) = listener.accept().await?;
103            let config = self.config.clone();
104            tokio::spawn(async move {
105                let io = TokioIo::new(stream);
106                let service = service_fn(move |req| {
107                    let config = config.clone();
108                    async move { handle_request(req, config, client_addr).await }
109                });
110                if let Err(e) = http1::Builder::new()
111                    .preserve_header_case(true)
112                    .title_case_headers(true)
113                    .serve_connection(io, service)
114                    .with_upgrades()
115                    .await
116                {
117                    error!("Connection error from {}: {}", client_addr, e);
118                }
119            });
120        }
121    }
122}
123
124// ---------------------------------------------------------------------------
125// Request handler
126// ---------------------------------------------------------------------------
127
128async fn handle_request(
129    req: Request<Incoming>,
130    config: Arc<ProxyConfig>,
131    client_addr: SocketAddr,
132) -> Result<Response<String>, hyper::Error> {
133    let method = req.method().clone();
134    let uri = req.uri().clone();
135    debug!("[{}] {} {}", client_addr, method, uri);
136
137    if method != Method::CONNECT {
138        return Ok(Response::builder()
139            .status(StatusCode::METHOD_NOT_ALLOWED)
140            .body("Only CONNECT is supported".into())
141            .unwrap());
142    }
143
144    let destination = uri
145        .authority()
146        .map(|a| a.as_str().to_string())
147        .unwrap_or_default();
148
149    if destination.is_empty() {
150        return Ok(Response::builder()
151            .status(StatusCode::BAD_REQUEST)
152            .body("Missing CONNECT destination".into())
153            .unwrap());
154    }
155
156    debug!("[{}] CONNECT → {}", client_addr, destination);
157
158    tokio::spawn(async move {
159        match hyper::upgrade::on(req).await {
160            Ok(upgraded) => {
161                let io = TokioIo::new(upgraded);
162                if let Err(e) = handle_tunnel(io, destination, config).await {
163                    error!("[{}] Tunnel error: {}", client_addr, e);
164                }
165            }
166            Err(e) => error!("[{}] Upgrade error: {}", client_addr, e),
167        }
168    });
169
170    Ok(Response::builder()
171        .status(StatusCode::OK)
172        .body(String::new())
173        .unwrap())
174}
175
176// ---------------------------------------------------------------------------
177// Tunnel
178// ---------------------------------------------------------------------------
179
180async fn handle_tunnel(
181    client: TokioIo<hyper::upgrade::Upgraded>,
182    destination: String,
183    config: Arc<ProxyConfig>,
184) -> io::Result<()> {
185    let mut upstream = match &config.upstream_proxy {
186        Some(proxy_addr) => {
187            connect_via_upstream(
188                proxy_addr,
189                &destination,
190                config.upstream_username.as_deref(),
191                config.upstream_password.as_deref(),
192            )
193            .await?
194        }
195        None => TcpStream::connect(&destination).await?,
196    };
197
198    let (mut cr, mut cw) = tokio::io::split(client);
199    let (mut ur, mut uw) = upstream.split();
200
201    match tokio::try_join!(
202        tokio::io::copy(&mut cr, &mut uw),
203        tokio::io::copy(&mut ur, &mut cw),
204    ) {
205        Ok((up, down)) => {
206            debug!("Tunnel closed {}: ↑{}B ↓{}B", destination, up, down);
207            Ok(())
208        }
209        Err(e) => Err(e),
210    }
211}
212
213async fn connect_via_upstream(
214    proxy_addr: &str,
215    destination: &str,
216    username: Option<&str>,
217    password: Option<&str>,
218) -> io::Result<TcpStream> {
219    let mut stream = TcpStream::connect(proxy_addr).await?;
220
221    let auth_header = match (username, password) {
222        (Some(u), Some(p)) => {
223            let encoded = base64::engine::general_purpose::STANDARD
224                .encode(format!("{u}:{p}"));
225            format!("Proxy-Authorization: Basic {encoded}\r\n")
226        }
227        _ => String::new(),
228    };
229
230    let req = format!(
231        "CONNECT {destination} HTTP/1.1\r\nHost: {destination}\r\n{auth_header}Connection: keep-alive\r\n\r\n"
232    );
233    stream.write_all(req.as_bytes()).await?;
234
235    let mut buf = vec![0u8; 4096];
236    let n = stream.read(&mut buf).await?;
237    let resp = String::from_utf8_lossy(&buf[..n]);
238    if !resp.contains("200") {
239        return Err(io::Error::new(
240            io::ErrorKind::ConnectionRefused,
241            format!(
242                "Upstream proxy rejected CONNECT: {}",
243                resp.lines().next().unwrap_or("unknown")
244            ),
245        ));
246    }
247
248    Ok(stream)
249}
250
251// ---------------------------------------------------------------------------
252// Convenience
253// ---------------------------------------------------------------------------
254
255/// Start a local overlay proxy in the background from a full proxy URL
256/// (e.g. `http://user:pass@host:port`). Returns the local `SocketAddr`
257/// Chrome should be pointed at.
258pub async fn start_overlay(proxy_url: &str) -> Result<SocketAddr, Box<dyn std::error::Error>> {
259    let config = ProxyConfig::from_proxy_url(proxy_url)
260        .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
261    let addr = ProxyServer::new(config).start_background().await?;
262    Ok(addr)
263}
264
265// ---------------------------------------------------------------------------
266// Tests
267// ---------------------------------------------------------------------------
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272
273    #[test]
274    fn parse_http_url_with_creds() {
275        let c = ProxyConfig::from_proxy_url("http://alice:secret@proxy.example.com:8080").unwrap();
276        assert_eq!(c.upstream_proxy.unwrap(), "proxy.example.com:8080");
277        assert_eq!(c.upstream_username.unwrap(), "alice");
278        assert_eq!(c.upstream_password.unwrap(), "secret");
279    }
280
281    #[test]
282    fn parse_https_url_no_creds() {
283        let c = ProxyConfig::from_proxy_url("https://10.0.0.1:3128").unwrap();
284        assert_eq!(c.upstream_proxy.unwrap(), "10.0.0.1:3128");
285        assert!(c.upstream_username.is_none());
286    }
287
288    #[test]
289    fn parse_invalid_scheme() {
290        assert!(ProxyConfig::from_proxy_url("socks5://host:1080").is_err());
291    }
292
293    #[test]
294    fn default_config() {
295        let c = ProxyConfig::default();
296        assert_eq!(c.bind_addr.to_string(), "127.0.0.1:0");
297        assert!(c.upstream_proxy.is_none());
298    }
299}
rustenium_identity/local_proxy_server.rs

rustenium_identity/
local_proxy_server.rs
