#![deny(warnings)]
use std::convert::Infallible;
use std::net::SocketAddr;
use fluxio::service::{make_service_fn, service_fn};
use fluxio::upgrade::Upgraded;
use fluxio::{Body, Client, Method, Request, Response, Server};
use tokio::net::TcpStream;
type HttpClient = Client<fluxio::client::HttpConnector>;
#[tokio::main]
async fn main() {
let addr = SocketAddr::from(([127, 0, 0, 1], 8100));
let client = Client::builder()
.http1_title_case_headers(true)
.http1_preserve_header_case(true)
.build_http();
let make_service = make_service_fn(move |_| {
let client = client.clone();
async move { Ok::<_, Infallible>(service_fn(move |req| proxy(client.clone(), req))) }
});
let server = Server::bind(&addr)
.http1_preserve_header_case(true)
.http1_title_case_headers(true)
.serve(make_service);
println!("Listening on http://{}", addr);
if let Err(e) = server.await {
eprintln!("server error: {}", e);
}
}
async fn proxy(client: HttpClient, req: Request<Body>) -> Result<Response<Body>, fluxio::Error> {
println!("req: {:?}", req);
if Method::CONNECT == req.method() {
if let Some(addr) = host_addr(req.uri()) {
tokio::task::spawn(async move {
match fluxio::upgrade::on(req).await {
Ok(upgraded) => {
if let Err(e) = tunnel(upgraded, addr).await {
eprintln!("server io error: {}", e);
};
}
Err(e) => eprintln!("upgrade error: {}", e),
}
});
Ok(Response::new(Body::empty()))
} else {
eprintln!("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() = http::StatusCode::BAD_REQUEST;
Ok(resp)
}
} else {
client.request(req).await
}
}
fn host_addr(uri: &http::Uri) -> Option<String> {
uri.authority().and_then(|auth| Some(auth.to_string()))
}
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?;
println!(
"client wrote {} bytes and received {} bytes",
from_client, from_server
);
Ok(())
}