Add HTTP proxy example (#425)

2021-10-26 20:46:40 +02:00 · 2021-10-26 20:46:40 +02:00 · fc9bfb8a50
parent 302a720271
commit fc9bfb8a50
2 changed files with 108 additions and 0 deletions
--- a/examples/http-proxy/Cargo.toml
+++ 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"
--- a/examples/http-proxy/src/main.rs
+++ 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<Body>| {
+        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<Body>) -> Result<Response<Body>, 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(())
+}