From fc9bfb8a506488440f8abc970fc5884ec8458e59 Mon Sep 17 00:00:00 2001 From: David Pedersen Date: Tue, 26 Oct 2021 20:46:40 +0200 Subject: [PATCH] Add HTTP proxy example (#425) --- examples/http-proxy/Cargo.toml | 13 +++++ examples/http-proxy/src/main.rs | 95 +++++++++++++++++++++++++++++++++ 2 files changed, 108 insertions(+) create mode 100644 examples/http-proxy/Cargo.toml create mode 100644 examples/http-proxy/src/main.rs diff --git a/examples/http-proxy/Cargo.toml b/examples/http-proxy/Cargo.toml new file mode 100644 index 00000000..833068df --- /dev/null +++ b/examples/http-proxy/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "example-http-proxy" +version = "0.1.0" +edition = "2018" +publish = false + +[dependencies] +axum = { path = "../.." } +tokio = { version = "1.0", features = ["full"] } +hyper = { version = "0.14", features = ["full"] } +tower = { version = "0.4", features = ["make"] } +tracing = "0.1" +tracing-subscriber = "0.2" diff --git a/examples/http-proxy/src/main.rs b/examples/http-proxy/src/main.rs new file mode 100644 index 00000000..95a883b3 --- /dev/null +++ b/examples/http-proxy/src/main.rs @@ -0,0 +1,95 @@ +//! Run with +//! +//! ```not_rust +//! $ cargo run -p example-http-proxy +//! ``` +//! +//! In another terminal: +//! +//! ```not_rust +//! $ curl -v -x "127.0.0.1:3000" https://tokio.rs +//! ``` +//! +//! Example is based on https://github.com/hyperium/hyper/blob/master/examples/http_proxy.rs + +use axum::{ + body::{box_body, Body}, + http::{Method, Request, Response, StatusCode}, + routing::get, + Router, +}; +use hyper::upgrade::Upgraded; +use std::net::SocketAddr; +use tokio::net::TcpStream; +use tower::{make::Shared, ServiceExt}; + +#[tokio::main] +async fn main() { + // Set the RUST_LOG, if it hasn't been explicitly defined + if std::env::var_os("RUST_LOG").is_none() { + std::env::set_var("RUST_LOG", "example_http_proxy=trace,tower_http=debug") + } + tracing_subscriber::fmt::init(); + + let router = Router::new().route("/", get(|| async { "Hello, World!" })); + + let service = tower::service_fn(move |req: Request| { + let router = router.clone(); + async move { + if req.method() == Method::CONNECT { + proxy(req).await.map(|res| res.map(box_body)) + } else { + router.oneshot(req).await.map_err(|err| match err {}) + } + } + }); + + let addr = SocketAddr::from(([127, 0, 0, 1], 3000)); + tracing::debug!("listening on {}", addr); + axum::Server::bind(&addr) + .http1_preserve_header_case(true) + .http1_title_case_headers(true) + .serve(Shared::new(service)) + .await + .unwrap(); +} + +async fn proxy(req: Request) -> Result, hyper::Error> { + tracing::trace!(?req); + + if let Some(host_addr) = req.uri().authority().map(|auth| auth.to_string()) { + tokio::task::spawn(async move { + match hyper::upgrade::on(req).await { + Ok(upgraded) => { + if let Err(e) = tunnel(upgraded, host_addr).await { + tracing::warn!("server io error: {}", e); + }; + } + Err(e) => tracing::warn!("upgrade error: {}", e), + } + }); + + Ok(Response::new(Body::empty())) + } else { + tracing::warn!("CONNECT host is not socket addr: {:?}", req.uri()); + let mut resp = Response::new(Body::from("CONNECT must be to a socket address")); + *resp.status_mut() = StatusCode::BAD_REQUEST; + + Ok(resp) + } +} + +async fn tunnel(mut upgraded: Upgraded, addr: String) -> std::io::Result<()> { + let mut server = TcpStream::connect(addr).await?; + + let (from_client, from_server) = + tokio::io::copy_bidirectional(&mut upgraded, &mut server).await?; + + tracing::debug!( + "client wrote {} bytes and received {} bytes", + from_client, + from_server + ); + + Ok(()) +}