diff --git a/proxy/cats/src/main.rs b/proxy/cats/src/main.rs index a61c816..c7310d7 100644 --- a/proxy/cats/src/main.rs +++ b/proxy/cats/src/main.rs @@ -3,15 +3,20 @@ use hyper::service::{make_service_fn, service_fn}; use hyper::{Body, Request, Response, Server, StatusCode}; use std::{ convert::Infallible, + env, net::{IpAddr, SocketAddr}, }; mod hyper_reverse_proxy; -async fn handle(client_ip: IpAddr, req: Request) -> Result, Infallible> { +async fn handle( + client_ip: IpAddr, + req: Request, + upstream: String, +) -> Result, Infallible> { println!("{} {} {}", client_ip, req.method(), req.uri()); - match hyper_reverse_proxy::call(client_ip, "https://dl.kitsu.red", req).await { + match hyper_reverse_proxy::call(client_ip, upstream.as_str(), req).await { Ok(response) => Ok(response), Err(_error) => { eprintln!("Error: {:#?}", _error); @@ -32,18 +37,25 @@ async fn shutdown_signal() { #[tokio::main] async fn main() { - let addr = SocketAddr::from(([127, 0, 0, 1], 3000)); + let args: Vec = env::args().collect(); + let upstream = args.get(1).expect("Argument upstream is required").clone(); let make_svc = make_service_fn(|socket: &AddrStream| { let remote_addr = socket.remote_addr().ip(); - async move { Ok::<_, Infallible>(service_fn(move |req| handle(remote_addr, req))) } + let upstream = upstream.clone(); + async move { + Ok::<_, Infallible>(service_fn(move |req| { + handle(remote_addr, req, upstream.clone()) + })) + } }); + let addr = SocketAddr::from(([127, 0, 0, 1], 3000)); let server = Server::bind(&addr) .serve(make_svc) .with_graceful_shutdown(shutdown_signal()); - println!("Listening on http://{}", addr); + println!("Listening on {}", addr); if let Err(e) = server.await { eprintln!("server error: {}", e); }