Skip to content

Instantly share code, notes, and snippets.

@carlosm27

carlosm27/main.rs

Created April 18, 2023 23:50
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save carlosm27/f57c4931dff31e22d07950353eb5b076 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save carlosm27/f57c4931dff31e22d07950353eb5b076 to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. carlosm27 created this gist Apr 18, 2023.
    86 changes: 86 additions & 0 deletions main.rs
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,86 @@
    use axum::{
    body::{self, Body},
    http::{Method, Request, StatusCode},
    response::{IntoResponse, Response},
    routing::get,
    Router,
    };
    use hyper::upgrade::Upgraded;
    use std::net::SocketAddr;
    use tokio::net::TcpStream;
    use tower::{make::Shared, ServiceExt};
    use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

    #[tokio::main]
    async fn main() {
    tracing_subscriber::registry()
    .with(
    tracing_subscriber::EnvFilter::try_from_default_env()
    .unwrap_or_else(|_| "example_http_proxy=trace,tower_http=debug".into()),
    )
    .with(tracing_subscriber::fmt::layer())
    .init();

    let router_svc = Router::new().route("/", get(|| async { "Hello, World!" }));

    let service = tower::service_fn(move |req: Request<Body>| {
    let router_svc = router_svc.clone();
    async move {
    if req.method() == Method::CONNECT {
    proxy(req).await
    } else {
    router_svc.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, 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::boxed(body::Empty::new())))
    } else {
    tracing::warn!("CONNECT host is not socket addr: {:?}", req.uri());
    Ok((
    StatusCode::BAD_REQUEST,
    "CONNECT must be to a socket address",
    )
    .into_response())
    }
    }

    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(())
    }