use hyper::body::Incoming;
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper::{Method, Request, Response, StatusCode};
use hyper_util::rt::TokioIo;
use std::io;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use base64::Engine;
use tracing::{debug, error, info};
#[derive(Debug, Clone)]
struct ProxyConfig {
bind_addr: SocketAddr,
upstream_proxy: Option<String>,
upstream_username: Option<String>,
upstream_password: Option<String>,
}
impl Default for ProxyConfig {
fn default() -> Self {
Self {
bind_addr: "127.0.0.1:0".parse().unwrap(),
upstream_proxy: None,
upstream_username: None,
upstream_password: None,
}
}
}
impl ProxyConfig {
fn from_proxy_url(url: &str) -> Result<Self, String> {
let rest = url
.strip_prefix("https://")
.or_else(|| url.strip_prefix("http://"))
.ok_or_else(|| format!("Unsupported or missing scheme in proxy URL: {url}"))?;
let (creds, hostport) = if let Some(at) = rest.rfind('@') {
(Some(&rest[..at]), &rest[at + 1..])
} else {
(None, rest)
};
let (username, password) = match creds {
Some(c) => {
if let Some((u, p)) = c.split_once(':') {
(Some(u.to_string()), Some(p.to_string()))
} else {
(Some(c.to_string()), None)
}
}
None => (None, None),
};
if hostport.is_empty() {
return Err("Missing host in proxy URL".into());
}
Ok(Self {
bind_addr: "127.0.0.1:0".parse().unwrap(),
upstream_proxy: Some(hostport.to_string()),
upstream_username: username,
upstream_password: password,
})
}
}
struct ProxyServer {
config: Arc<ProxyConfig>,
}
impl ProxyServer {
fn new(config: ProxyConfig) -> Self {
Self { config: Arc::new(config) }
}
async fn start_background(self) -> io::Result<SocketAddr> {
let listener = TcpListener::bind(self.config.bind_addr).await?;
let addr = listener.local_addr()?;
info!("Local proxy overlay listening on {}", addr);
tokio::spawn(async move {
if let Err(e) = self.run(listener).await {
error!("Proxy server exited: {}", e);
}
});
Ok(addr)
}
async fn run(&self, listener: TcpListener) -> Result<(), Box<dyn std::error::Error>> {
loop {
let (stream, client_addr) = listener.accept().await?;
let config = self.config.clone();
tokio::spawn(async move {
let io = TokioIo::new(stream);
let service = service_fn(move |req| {
let config = config.clone();
async move { handle_request(req, config, client_addr).await }
});
if let Err(e) = http1::Builder::new()
.preserve_header_case(true)
.title_case_headers(true)
.serve_connection(io, service)
.with_upgrades()
.await
{
error!("Connection error from {}: {}", client_addr, e);
}
});
}
}
}
async fn handle_request(
req: Request<Incoming>,
config: Arc<ProxyConfig>,
client_addr: SocketAddr,
) -> Result<Response<String>, hyper::Error> {
let method = req.method().clone();
let uri = req.uri().clone();
debug!("[{}] {} {}", client_addr, method, uri);
if method != Method::CONNECT {
return Ok(Response::builder()
.status(StatusCode::METHOD_NOT_ALLOWED)
.body("Only CONNECT is supported".into())
.unwrap());
}
let destination = uri
.authority()
.map(|a| a.as_str().to_string())
.unwrap_or_default();
if destination.is_empty() {
return Ok(Response::builder()
.status(StatusCode::BAD_REQUEST)
.body("Missing CONNECT destination".into())
.unwrap());
}
debug!("[{}] CONNECT → {}", client_addr, destination);
tokio::spawn(async move {
match hyper::upgrade::on(req).await {
Ok(upgraded) => {
let io = TokioIo::new(upgraded);
if let Err(e) = handle_tunnel(io, destination, config).await {
error!("[{}] Tunnel error: {}", client_addr, e);
}
}
Err(e) => error!("[{}] Upgrade error: {}", client_addr, e),
}
});
Ok(Response::builder()
.status(StatusCode::OK)
.body(String::new())
.unwrap())
}
async fn handle_tunnel(
client: TokioIo<hyper::upgrade::Upgraded>,
destination: String,
config: Arc<ProxyConfig>,
) -> io::Result<()> {
let mut upstream = match &config.upstream_proxy {
Some(proxy_addr) => {
connect_via_upstream(
proxy_addr,
&destination,
config.upstream_username.as_deref(),
config.upstream_password.as_deref(),
)
.await?
}
None => TcpStream::connect(&destination).await?,
};
let (mut cr, mut cw) = tokio::io::split(client);
let (mut ur, mut uw) = upstream.split();
match tokio::try_join!(
tokio::io::copy(&mut cr, &mut uw),
tokio::io::copy(&mut ur, &mut cw),
) {
Ok((up, down)) => {
debug!("Tunnel closed {}: ↑{}B ↓{}B", destination, up, down);
Ok(())
}
Err(e) => Err(e),
}
}
async fn connect_via_upstream(
proxy_addr: &str,
destination: &str,
username: Option<&str>,
password: Option<&str>,
) -> io::Result<TcpStream> {
let mut stream = TcpStream::connect(proxy_addr).await?;
let auth_header = match (username, password) {
(Some(u), Some(p)) => {
let encoded = base64::engine::general_purpose::STANDARD
.encode(format!("{u}:{p}"));
format!("Proxy-Authorization: Basic {encoded}\r\n")
}
_ => String::new(),
};
let req = format!(
"CONNECT {destination} HTTP/1.1\r\nHost: {destination}\r\n{auth_header}Connection: keep-alive\r\n\r\n"
);
stream.write_all(req.as_bytes()).await?;
let mut buf = vec![0u8; 4096];
let n = stream.read(&mut buf).await?;
let resp = String::from_utf8_lossy(&buf[..n]);
if !resp.contains("200") {
return Err(io::Error::new(
io::ErrorKind::ConnectionRefused,
format!(
"Upstream proxy rejected CONNECT: {}",
resp.lines().next().unwrap_or("unknown")
),
));
}
Ok(stream)
}
pub async fn start_overlay(proxy_url: &str) -> Result<SocketAddr, Box<dyn std::error::Error>> {
let config = ProxyConfig::from_proxy_url(proxy_url)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
let addr = ProxyServer::new(config).start_background().await?;
Ok(addr)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_http_url_with_creds() {
let c = ProxyConfig::from_proxy_url("http://alice:secret@proxy.example.com:8080").unwrap();
assert_eq!(c.upstream_proxy.unwrap(), "proxy.example.com:8080");
assert_eq!(c.upstream_username.unwrap(), "alice");
assert_eq!(c.upstream_password.unwrap(), "secret");
}
#[test]
fn parse_https_url_no_creds() {
let c = ProxyConfig::from_proxy_url("https://10.0.0.1:3128").unwrap();
assert_eq!(c.upstream_proxy.unwrap(), "10.0.0.1:3128");
assert!(c.upstream_username.is_none());
}
#[test]
fn parse_invalid_scheme() {
assert!(ProxyConfig::from_proxy_url("socks5://host:1080").is_err());
}
#[test]
fn default_config() {
let c = ProxyConfig::default();
assert_eq!(c.bind_addr.to_string(), "127.0.0.1:0");
assert!(c.upstream_proxy.is_none());
}
}