From 1309571f5dcb4f24701ca913ef8160205d6f6a39 Mon Sep 17 00:00:00 2001 From: Alex Chi Z Date: Fri, 14 Jul 2023 12:11:01 -0400 Subject: [PATCH] proxy: switch to structopt for clap parsing (#4714) Using `#[clap]` for parsing cli opts, which is easier to maintain. --------- Signed-off-by: Alex Chi Z --- proxy/src/bin/proxy.rs | 219 ++++++++++++++++------------------------- 1 file changed, 84 insertions(+), 135 deletions(-) diff --git a/proxy/src/bin/proxy.rs b/proxy/src/bin/proxy.rs index fc8bc39742..6b46eaddfa 100644 --- a/proxy/src/bin/proxy.rs +++ b/proxy/src/bin/proxy.rs @@ -5,7 +5,6 @@ use proxy::http; use proxy::metrics; use anyhow::bail; -use clap::{self, Arg}; use proxy::config::{self, ProxyConfig}; use std::pin::pin; use std::{borrow::Cow, net::SocketAddr}; @@ -18,6 +17,70 @@ use utils::{project_git_version, sentry_init::init_sentry}; project_git_version!(GIT_VERSION); +use clap::{Parser, ValueEnum}; + +#[derive(Clone, Debug, ValueEnum)] +enum AuthBackend { + Console, + Postgres, + Link, +} + +/// Neon proxy/router +#[derive(Parser)] +#[command(version = GIT_VERSION, about)] +struct ProxyCliArgs { + /// listen for incoming client connections on ip:port + #[clap(short, long, default_value = "127.0.0.1:4432")] + proxy: String, + #[clap(value_enum, long, default_value_t = AuthBackend::Link)] + auth_backend: AuthBackend, + /// listen for management callback connection on ip:port + #[clap(short, long, default_value = "127.0.0.1:7000")] + mgmt: String, + /// listen for incoming http connections (metrics, etc) on ip:port + #[clap(long, default_value = "127.0.0.1:7001")] + http: String, + /// listen for incoming wss connections on ip:port + #[clap(long)] + wss: Option, + /// redirect unauthenticated users to the given uri in case of link auth + #[clap(short, long, default_value = "http://localhost:3000/psql_session/")] + uri: String, + /// cloud API endpoint for authenticating users + #[clap( + short, + long, + default_value = "http://localhost:3000/authenticate_proxy_request/" + )] + auth_endpoint: String, + /// path to TLS key for client postgres connections + /// + /// tls-key and tls-cert are for backwards compatibility, we can put all certs in one dir + #[clap(short = 'k', long, alias = "ssl-key")] + tls_key: Option, + /// path to TLS cert for client postgres connections + /// + /// tls-key and tls-cert are for backwards compatibility, we can put all certs in one dir + #[clap(short = 'c', long, alias = "ssl-cert")] + tls_cert: Option, + /// path to directory with TLS certificates for client postgres connections + #[clap(long)] + certs_dir: Option, + /// http endpoint to receive periodic metric updates + #[clap(long)] + metric_collection_endpoint: Option, + /// how often metrics should be sent to a collection endpoint + #[clap(long)] + metric_collection_interval: Option, + /// cache for `wake_compute` api method (use `size=0` to disable) + #[clap(long, default_value = config::CacheOptions::DEFAULT_OPTIONS_NODE_INFO)] + wake_compute_cache: String, + /// Allow self-signed certificates for compute nodes (for testing) + #[clap(long, default_value_t = false, value_parser = clap::builder::BoolishValueParser::new(), action = clap::ArgAction::Set)] + allow_self_signed_compute: bool, +} + #[tokio::main] async fn main() -> anyhow::Result<()> { let _logging_guard = proxy::logging::init().await?; @@ -27,21 +90,21 @@ async fn main() -> anyhow::Result<()> { info!("Version: {GIT_VERSION}"); ::metrics::set_build_info_metric(GIT_VERSION); - let args = cli().get_matches(); + let args = ProxyCliArgs::parse(); let config = build_config(&args)?; info!("Authentication backend: {}", config.auth_backend); // Check that we can bind to address before further initialization - let http_address: SocketAddr = args.get_one::("http").unwrap().parse()?; + let http_address: SocketAddr = args.http.parse()?; info!("Starting http on {http_address}"); let http_listener = TcpListener::bind(http_address).await?.into_std()?; - let mgmt_address: SocketAddr = args.get_one::("mgmt").unwrap().parse()?; + let mgmt_address: SocketAddr = args.mgmt.parse()?; info!("Starting mgmt on {mgmt_address}"); let mgmt_listener = TcpListener::bind(mgmt_address).await?; - let proxy_address: SocketAddr = args.get_one::("proxy").unwrap().parse()?; + let proxy_address: SocketAddr = args.proxy.parse()?; info!("Starting proxy on {proxy_address}"); let proxy_listener = TcpListener::bind(proxy_address).await?; let cancellation_token = CancellationToken::new(); @@ -55,7 +118,7 @@ async fn main() -> anyhow::Result<()> { cancellation_token.clone(), )); - if let Some(wss_address) = args.get_one::("wss") { + if let Some(wss_address) = args.wss { let wss_address: SocketAddr = wss_address.parse()?; info!("Starting wss on {wss_address}"); let wss_listener = TcpListener::bind(wss_address).await?; @@ -102,31 +165,24 @@ async fn main() -> anyhow::Result<()> { } /// ProxyConfig is created at proxy startup, and lives forever. -fn build_config(args: &clap::ArgMatches) -> anyhow::Result<&'static ProxyConfig> { - let tls_config = match ( - args.get_one::("tls-key"), - args.get_one::("tls-cert"), - ) { +fn build_config(args: &ProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> { + let tls_config = match (&args.tls_key, &args.tls_cert) { (Some(key_path), Some(cert_path)) => Some(config::configure_tls( key_path, cert_path, - args.get_one::("certs-dir"), + args.certs_dir.as_ref(), )?), (None, None) => None, _ => bail!("either both or neither tls-key and tls-cert must be specified"), }; - let allow_self_signed_compute: bool = args - .get_one::("allow-self-signed-compute") - .unwrap() - .parse()?; - if allow_self_signed_compute { + if args.allow_self_signed_compute { warn!("allowing self-signed compute certificates"); } let metric_collection = match ( - args.get_one::("metric-collection-endpoint"), - args.get_one::("metric-collection-interval"), + &args.metric_collection_endpoint, + &args.metric_collection_interval, ) { (Some(endpoint), Some(interval)) => Some(config::MetricCollectionConfig { endpoint: endpoint.parse()?, @@ -139,145 +195,38 @@ fn build_config(args: &clap::ArgMatches) -> anyhow::Result<&'static ProxyConfig> ), }; - let auth_backend = match args.get_one::("auth-backend").unwrap().as_str() { - "console" => { - let config::CacheOptions { size, ttl } = args - .get_one::("wake-compute-cache") - .unwrap() - .parse()?; + let auth_backend = match &args.auth_backend { + AuthBackend::Console => { + let config::CacheOptions { size, ttl } = args.wake_compute_cache.parse()?; info!("Using NodeInfoCache (wake_compute) with size={size} ttl={ttl:?}"); let caches = Box::leak(Box::new(console::caches::ApiCaches { node_info: console::caches::NodeInfoCache::new("node_info_cache", size, ttl), })); - let url = args.get_one::("auth-endpoint").unwrap().parse()?; + let url = args.auth_endpoint.parse()?; let endpoint = http::Endpoint::new(url, http::new_client()); let api = console::provider::neon::Api::new(endpoint, caches); auth::BackendType::Console(Cow::Owned(api), ()) } - "postgres" => { - let url = args.get_one::("auth-endpoint").unwrap().parse()?; + AuthBackend::Postgres => { + let url = args.auth_endpoint.parse()?; let api = console::provider::mock::Api::new(url); auth::BackendType::Postgres(Cow::Owned(api), ()) } - "link" => { - let url = args.get_one::("uri").unwrap().parse()?; + AuthBackend::Link => { + let url = args.uri.parse()?; auth::BackendType::Link(Cow::Owned(url)) } - other => bail!("unsupported auth backend: {other}"), }; let config = Box::leak(Box::new(ProxyConfig { tls_config, auth_backend, metric_collection, - allow_self_signed_compute, + allow_self_signed_compute: args.allow_self_signed_compute, })); Ok(config) } - -fn cli() -> clap::Command { - clap::Command::new("Neon proxy/router") - .disable_help_flag(true) - .version(GIT_VERSION) - .arg( - Arg::new("proxy") - .short('p') - .long("proxy") - .help("listen for incoming client connections on ip:port") - .default_value("127.0.0.1:4432"), - ) - .arg( - Arg::new("auth-backend") - .long("auth-backend") - .value_parser(["console", "postgres", "link"]) - .default_value("link"), - ) - .arg( - Arg::new("mgmt") - .short('m') - .long("mgmt") - .help("listen for management callback connection on ip:port") - .default_value("127.0.0.1:7000"), - ) - .arg( - Arg::new("http") - .long("http") - .help("listen for incoming http connections (metrics, etc) on ip:port") - .default_value("127.0.0.1:7001"), - ) - .arg( - Arg::new("wss") - .long("wss") - .help("listen for incoming wss connections on ip:port"), - ) - .arg( - Arg::new("uri") - .short('u') - .long("uri") - .help("redirect unauthenticated users to the given uri in case of link auth") - .default_value("http://localhost:3000/psql_session/"), - ) - .arg( - Arg::new("auth-endpoint") - .short('a') - .long("auth-endpoint") - .help("cloud API endpoint for authenticating users") - .default_value("http://localhost:3000/authenticate_proxy_request/"), - ) - .arg( - Arg::new("tls-key") - .short('k') - .long("tls-key") - .alias("ssl-key") // backwards compatibility - .help("path to TLS key for client postgres connections"), - ) - .arg( - Arg::new("tls-cert") - .short('c') - .long("tls-cert") - .alias("ssl-cert") // backwards compatibility - .help("path to TLS cert for client postgres connections"), - ) - // tls-key and tls-cert are for backwards compatibility, we can put all certs in one dir - .arg( - Arg::new("certs-dir") - .long("certs-dir") - .help("path to directory with TLS certificates for client postgres connections"), - ) - .arg( - Arg::new("metric-collection-endpoint") - .long("metric-collection-endpoint") - .help("http endpoint to receive periodic metric updates"), - ) - .arg( - Arg::new("metric-collection-interval") - .long("metric-collection-interval") - .help("how often metrics should be sent to a collection endpoint"), - ) - .arg( - Arg::new("wake-compute-cache") - .long("wake-compute-cache") - .help("cache for `wake_compute` api method (use `size=0` to disable)") - .default_value(config::CacheOptions::DEFAULT_OPTIONS_NODE_INFO), - ) - .arg( - Arg::new("allow-self-signed-compute") - .long("allow-self-signed-compute") - .help("Allow self-signed certificates for compute nodes (for testing)") - .default_value("false"), - ) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn verify_cli() { - cli().debug_assert(); - } -}