Rewrite in Rust

2021-03-18 18:37:44 +01:00
parent b2125dd4be
commit 1239ee3dad
9 changed files with 1352 additions and 50 deletions
--- a/src/main.rs
+++ b/src/main.rs
@ -0,0 +1,103 @@
+use compound_error::CompoundError;
+use hyper::service::{make_service_fn, service_fn};
+use hyper::{Body, Client, Request, Response, Server, Uri};
+use kv_log_macro::info;
+use log::error;
+use std::net::SocketAddr;
+use std::path::PathBuf;
+use structopt::StructOpt;
+use tokio::fs;
+
+#[derive(StructOpt)]
+struct Opt {
+    #[structopt(short, long)]
+    file: PathBuf,
+
+    #[structopt(short, long)]
+    user_agent: String,
+
+    #[structopt(short, long)]
+    proxy_pass: String,
+}
+
+#[derive(Debug, CompoundError)]
+enum Error {
+    Hyper(hyper::Error),
+    IO(std::io::Error),
+}
+
+async fn proxy_pass(mut req: Request<Body>, opt: &Opt) -> Result<Response<Body>, Error> {
+    let user_agent: Option<String> = req
+        .headers()
+        .get("User-Agent")
+        .and_then(|h| h.to_str().ok())
+        .map(|s| s.to_string());
+
+    let user_agent_matches = user_agent
+        .as_ref()
+        .map(|user_agent| user_agent.contains(&opt.user_agent))
+        .unwrap_or(false);
+
+    let req_uri = req.uri().clone();
+
+    let mut error = None;
+
+    let response = if user_agent_matches {
+        let file = fs::read_to_string(&opt.file).await?;
+        Response::new(file.into())
+    } else {
+        let proxy_uri: Uri = opt.proxy_pass.parse().expect("proxy uri");
+
+        let mut new_uri = Uri::builder().authority(proxy_uri.authority().unwrap().clone());
+
+        new_uri = new_uri.scheme(proxy_uri.scheme_str().unwrap_or("http"));
+
+        if let Some(paq) = req_uri.path_and_query().cloned() {
+            new_uri = new_uri.path_and_query(paq);
+        }
+
+        *req.uri_mut() = new_uri.build().expect("uri");
+
+        let client = Client::new();
+        match client.request(req).await {
+            Ok(response) => response,
+            Err(e) => {
+                error = Some(e);
+                Response::builder()
+                    .status(503)
+                    .body("503 Service Unavailable".into())
+                    .unwrap()
+            }
+        }
+    };
+
+    info!("request received", {
+        user_agent: format!("{:?}", user_agent).as_str(),
+        pattern: opt.user_agent.as_str(),
+        matches: user_agent_matches,
+        uri: format!("{}", req_uri).as_str(),
+        error: format!("{:?}", error).as_str(),
+    });
+
+    Ok(response)
+}
+
+#[tokio::main]
+async fn main() {
+    let opt: &'static Opt = Box::leak(Box::new(Opt::from_args()));
+
+    femme::start();
+
+    let addr = SocketAddr::from(([0, 0, 0, 0], 3000));
+
+    let make_svc =
+        make_service_fn(
+            |_conn| async move { Ok::<_, Error>(service_fn(move |r| proxy_pass(r, opt))) },
+        );
+
+    let server = Server::bind(&addr).serve(make_svc);
+
+    if let Err(e) = server.await {
+        error!("server error: {}", e);
+    }
+}