use hyper::service::{make_service_fn, service_fn}; use hyper::{header::HeaderValue, Body, Client, Request, Response, Server}; use std::env; use std::net::SocketAddr; use toml; #[macro_use] extern crate log; use es_public_proxy::{filter_request, ProxyConfig, ProxyError}; const CARGO_VERSION: Option<&'static str> = option_env!("CARGO_PKG_VERSION"); async fn upstream_req( req: Request, config: ProxyConfig, ) -> Result, hyper::Error> { info!("request: {} {}", req.method(), req.uri()); let upstream_resp: Result, ProxyError> = match filter_request(req, &config).await { Ok(parsed_req) => { debug!("forwarding request upstream"); Client::new().request(parsed_req).await.map_err(|e| { warn!("upstream error: {}", e); ProxyError::UpstreamError(e.to_string()) }) } Err(e) => { warn!("blocked: {:?}", e); Err(e) } }; let mut resp = match upstream_resp { Ok(resp) => resp, Err(other) => Response::builder() .status(other.http_status_code()) .header("Content-Type", "application/json; charset=UTF-8") .body( serde_json::to_string(&other.to_json_value()) .expect("trivial JSON value serialization failed") .into(), ) .expect("trivial HTTP response construction failed"), }; // regardless of the response type (error or upstream response), add HTTP headers resp.headers_mut() .insert("Via", HeaderValue::from_static("HTTP/1.1 es-public-proxy")); if config.enable_cors == Some(true) { resp.headers_mut() .insert("Access-Control-Allow-Origin", HeaderValue::from_static("*")); resp.headers_mut().insert( "Access-Control-Allow-Methods", HeaderValue::from_static("GET, POST, DELETE, HEAD, OPTIONS"), ); resp.headers_mut().insert( "Access-Control-Allow-Headers", HeaderValue::from_static("DNT,User-Agent,Content-Type"), ); } Ok(resp) } async fn shutdown_signal() { // Wait for the CTRL+C signal tokio::signal::ctrl_c() .await .expect("failed to install CTRL+C signal handler"); } 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(), }; warn!("Listening on http://{}", addr); // 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 { error!("server error: {}", e); } } 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; let mut unsafe_all_indices = false; // first parse CLI arg match args.as_slice() { [_] | [] => {} [_, "-h"] | [_, "--help"] => { println!("{}", usage()); std::process::exit(0); } [_, "--version"] => { println!("es-public-proxy v{}", CARGO_VERSION.unwrap_or("UNKNOWN")); std::process::exit(0); } [_, "--example-config"] => { println!("{}", include_str!("../extra/example_config.toml")); std::process::exit(0); } [_, "--config", p] => config_path = Some(p.to_string()), [_, "--unsafe-all-indices"] => unsafe_all_indices = true, _ => { 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 let mut 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() }; if unsafe_all_indices { config.unsafe_all_indices = Some(true); } config } #[tokio::main] async fn main() { let config = load_config(); env_logger::init_from_env( env_logger::Env::default().filter_or(env_logger::DEFAULT_FILTER_ENV, "info"), ); run_server(config).await; }