From ac53f1b078559fe5ba806eae6f7a768ee009b4f7 Mon Sep 17 00:00:00 2001 From: Bryan Newbold Date: Mon, 24 Aug 2020 18:02:11 -0700 Subject: config parsing; pass config object through to endpoint function --- src/main.rs | 73 +++++++++++++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 64 insertions(+), 9 deletions(-) (limited to 'src/main.rs') 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) -> Result, hyper::Error> { +use es_public_proxy::ProxyConfig; + +async fn upstream_req(req: Request, _config: ProxyConfig) -> Result, 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 = env::args().collect(); + let args: Vec<&str> = args.iter().map(|x| x.as_str()).collect(); + let mut config_path: Option = 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; } -- cgit v1.2.3