mirror of
https://github.com/tokio-rs/axum.git
synced 2025-01-06 10:44:32 +01:00
96 lines
2.7 KiB
Rust
96 lines
2.7 KiB
Rust
|
//! 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(())
|
||
|
|
}
|