1 files changed, 64 insertions, 9 deletions
diff --git a/src/main.rs b/src/main.rs
index 92d030a..62cfd37 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,8 +2,12 @@
 use hyper::service::{make_service_fn, service_fn};
 use hyper::{Body, Client, Request, Response, Server, Uri};
 use std::net::SocketAddr;
+use std::env;
+use toml;
 
-async fn upstream_req(req: Request<Body>) -> Result<Response<Body>, hyper::Error> {
+use es_public_proxy::ProxyConfig;
+
+async fn upstream_req(req: Request<Body>, _config: ProxyConfig) -> Result<Response<Body>, hyper::Error> {
     println!("hit: {}", req.uri());
     let req_uri = req.uri();
     let upstream_uri = Uri::builder()
@@ -25,13 +29,25 @@ async fn shutdown_signal() {
         .expect("failed to install CTRL+C signal handler");
 }
 
-async fn run_server(addr: SocketAddr) {
+async fn run_server(config: ProxyConfig) {
+
+    let addr = match &config.bind_addr {
+        None => SocketAddr::from(([127, 0, 0, 1], 9292)),
+        Some(addr) => addr.parse().unwrap(),
+    };
+
     println!("Listening on http://{}", addr);
 
-    let serve_future = Server::bind(&addr)
-        .serve(make_service_fn(|_| async {
-            Ok::<_, hyper::Error>(service_fn(upstream_req))
-        }));
+    // TODO: possible to avoid cloning config on every connection?
+    let make_svc = make_service_fn(move |_| {
+        let inner = config.clone();
+        async move {
+            Ok::<_, hyper::Error>(service_fn(move |req| {
+                upstream_req(req, inner.clone())
+            }))
+        }
+    });
+    let serve_future = Server::bind(&addr).serve(make_svc);
     let graceful = serve_future.with_graceful_shutdown(shutdown_signal());
 
     if let Err(e) = graceful.await {
@@ -39,9 +55,48 @@ async fn run_server(addr: SocketAddr) {
     }
 }
 
+fn usage() -> String {
+    "es-public-proxy [--config CONFIG_FILE] [--help]".to_string()
+}
+
+fn load_config() -> ProxyConfig {
+
+    let args: Vec<String> = env::args().collect();
+    let args: Vec<&str> = args.iter().map(|x| x.as_str()).collect();
+    let mut config_path: Option<String> = None;
+
+    // first parse CLI arg
+    match args.as_slice() {
+        [_] | [] => {},
+        [_, "-h"] | [_, "--help"] => {
+            println!("{}", usage());
+            std::process::exit(0);
+        },
+        [_, "--config", p] => { config_path = Some(p.to_string()) },
+        _ => {
+            eprintln!("{}", usage());
+            eprintln!("couldn't parse arguments");
+            std::process::exit(1);
+        }
+    }
+
+    // then try environment variables
+    if let None = config_path {
+        config_path = std::env::var("ES_PUBLIC_PROXY_CONFIG_PATH").ok();
+    }
+
+    // then either load config file (TOML), or use default config
+    if let Some(config_path) = config_path {
+        let config_toml = std::fs::read_to_string(config_path).unwrap();
+        let config: ProxyConfig = toml::from_str(&config_toml).unwrap();
+        config
+    } else {
+        ProxyConfig::default()
+    }
+}
+
 #[tokio::main]
 async fn main() {
-  let addr = SocketAddr::from(([127, 0, 0, 1], 3030));
-
-  run_server(addr).await;
+    let config = load_config();
+    run_server(config).await;
 }