LCOV - code coverage report
Current view: top level - proxy/src/binary - local_proxy.rs (source / functions) Coverage Total Hit
Test: 07bee600374ccd486c69370d0972d9035964fe68.info Lines: 0.0 % 251 0
Test Date: 2025-02-20 13:11:02 Functions: 0.0 % 70 0

            Line data    Source code
       1              : use std::net::SocketAddr;
       2              : use std::pin::pin;
       3              : use std::str::FromStr;
       4              : use std::sync::Arc;
       5              : use std::time::Duration;
       6              : 
       7              : use anyhow::{bail, ensure, Context};
       8              : use camino::{Utf8Path, Utf8PathBuf};
       9              : use clap::Parser;
      10              : use compute_api::spec::LocalProxySpec;
      11              : use futures::future::Either;
      12              : use thiserror::Error;
      13              : use tokio::net::TcpListener;
      14              : use tokio::sync::Notify;
      15              : use tokio::task::JoinSet;
      16              : use tokio_util::sync::CancellationToken;
      17              : use tracing::{debug, error, info, warn};
      18              : use utils::sentry_init::init_sentry;
      19              : use utils::{pid_file, project_build_tag, project_git_version};
      20              : 
      21              : use crate::auth::backend::jwt::JwkCache;
      22              : use crate::auth::backend::local::{LocalBackend, JWKS_ROLE_MAP};
      23              : use crate::auth::{self};
      24              : use crate::cancellation::CancellationHandler;
      25              : use crate::config::{
      26              :     self, AuthenticationConfig, ComputeConfig, HttpConfig, ProxyConfig, RetryConfig,
      27              : };
      28              : use crate::control_plane::locks::ApiLocks;
      29              : use crate::control_plane::messages::{EndpointJwksResponse, JwksSettings};
      30              : use crate::http::health_server::AppMetrics;
      31              : use crate::intern::RoleNameInt;
      32              : use crate::metrics::{Metrics, ThreadPoolMetrics};
      33              : use crate::rate_limiter::{
      34              :     BucketRateLimiter, EndpointRateLimiter, LeakyBucketConfig, RateBucketInfo,
      35              : };
      36              : use crate::scram::threadpool::ThreadPool;
      37              : use crate::serverless::cancel_set::CancelSet;
      38              : use crate::serverless::{self, GlobalConnPoolOptions};
      39              : use crate::tls::client_config::compute_client_config_with_root_certs;
      40              : use crate::types::RoleName;
      41              : use crate::url::ApiUrl;
      42              : 
      43              : project_git_version!(GIT_VERSION);
      44              : project_build_tag!(BUILD_TAG);
      45              : 
      46              : /// Neon proxy/router
      47              : #[derive(Parser)]
      48              : #[command(version = GIT_VERSION, about)]
      49              : struct LocalProxyCliArgs {
      50              :     /// listen for incoming metrics connections on ip:port
      51              :     #[clap(long, default_value = "127.0.0.1:7001")]
      52            0 :     metrics: String,
      53              :     /// listen for incoming http connections on ip:port
      54              :     #[clap(long)]
      55            0 :     http: String,
      56              :     /// timeout for the TLS handshake
      57              :     #[clap(long, default_value = "15s", value_parser = humantime::parse_duration)]
      58            0 :     handshake_timeout: tokio::time::Duration,
      59              :     /// lock for `connect_compute` api method. example: "shards=32,permits=4,epoch=10m,timeout=1s". (use `permits=0` to disable).
      60              :     #[clap(long, default_value = config::ConcurrencyLockOptions::DEFAULT_OPTIONS_CONNECT_COMPUTE_LOCK)]
      61            0 :     connect_compute_lock: String,
      62              :     #[clap(flatten)]
      63              :     sql_over_http: SqlOverHttpArgs,
      64              :     /// User rate limiter max number of requests per second.
      65              :     ///
      66              :     /// Provided in the form `<Requests Per Second>@<Bucket Duration Size>`.
      67              :     /// Can be given multiple times for different bucket sizes.
      68            0 :     #[clap(long, default_values_t = RateBucketInfo::DEFAULT_ENDPOINT_SET)]
      69            0 :     user_rps_limit: Vec<RateBucketInfo>,
      70              :     /// Whether the auth rate limiter actually takes effect (for testing)
      71            0 :     #[clap(long, default_value_t = false, value_parser = clap::builder::BoolishValueParser::new(), action = clap::ArgAction::Set)]
      72            0 :     auth_rate_limit_enabled: bool,
      73              :     /// Authentication rate limiter max number of hashes per second.
      74            0 :     #[clap(long, default_values_t = RateBucketInfo::DEFAULT_AUTH_SET)]
      75            0 :     auth_rate_limit: Vec<RateBucketInfo>,
      76              :     /// The IP subnet to use when considering whether two IP addresses are considered the same.
      77            0 :     #[clap(long, default_value_t = 64)]
      78            0 :     auth_rate_limit_ip_subnet: u8,
      79              :     /// Whether to retry the connection to the compute node
      80              :     #[clap(long, default_value = config::RetryConfig::CONNECT_TO_COMPUTE_DEFAULT_VALUES)]
      81            0 :     connect_to_compute_retry: String,
      82              :     /// Address of the postgres server
      83              :     #[clap(long, default_value = "127.0.0.1:5432")]
      84            0 :     postgres: SocketAddr,
      85              :     /// Address of the internal compute-ctl api service
      86              :     #[clap(long, default_value = "http://127.0.0.1:3081/")]
      87            0 :     compute_ctl: ApiUrl,
      88              :     /// Path of the local proxy config file
      89              :     #[clap(long, default_value = "./local_proxy.json")]
      90            0 :     config_path: Utf8PathBuf,
      91              :     /// Path of the local proxy PID file
      92              :     #[clap(long, default_value = "./local_proxy.pid")]
      93            0 :     pid_path: Utf8PathBuf,
      94              : }
      95              : 
      96              : #[derive(clap::Args, Clone, Copy, Debug)]
      97              : struct SqlOverHttpArgs {
      98              :     /// How many connections to pool for each endpoint. Excess connections are discarded
      99            0 :     #[clap(long, default_value_t = 200)]
     100            0 :     sql_over_http_pool_max_total_conns: usize,
     101              : 
     102              :     /// How long pooled connections should remain idle for before closing
     103              :     #[clap(long, default_value = "5m", value_parser = humantime::parse_duration)]
     104            0 :     sql_over_http_idle_timeout: tokio::time::Duration,
     105              : 
     106            0 :     #[clap(long, default_value_t = 100)]
     107            0 :     sql_over_http_client_conn_threshold: u64,
     108              : 
     109            0 :     #[clap(long, default_value_t = 16)]
     110            0 :     sql_over_http_cancel_set_shards: usize,
     111              : 
     112            0 :     #[clap(long, default_value_t = 10 * 1024 * 1024)] // 10 MiB
     113            0 :     sql_over_http_max_request_size_bytes: usize,
     114              : 
     115            0 :     #[clap(long, default_value_t = 10 * 1024 * 1024)] // 10 MiB
     116            0 :     sql_over_http_max_response_size_bytes: usize,
     117              : }
     118              : 
     119            0 : pub async fn run() -> anyhow::Result<()> {
     120            0 :     let _logging_guard = crate::logging::init_local_proxy()?;
     121            0 :     let _panic_hook_guard = utils::logging::replace_panic_hook_with_tracing_panic_hook();
     122            0 :     let _sentry_guard = init_sentry(Some(GIT_VERSION.into()), &[]);
     123            0 : 
     124            0 :     Metrics::install(Arc::new(ThreadPoolMetrics::new(0)));
     125            0 : 
     126            0 :     // TODO: refactor these to use labels
     127            0 :     debug!("Version: {GIT_VERSION}");
     128            0 :     debug!("Build_tag: {BUILD_TAG}");
     129            0 :     let neon_metrics = ::metrics::NeonMetrics::new(::metrics::BuildInfo {
     130            0 :         revision: GIT_VERSION,
     131            0 :         build_tag: BUILD_TAG,
     132            0 :     });
     133              : 
     134            0 :     let jemalloc = match crate::jemalloc::MetricRecorder::new() {
     135            0 :         Ok(t) => Some(t),
     136            0 :         Err(e) => {
     137            0 :             tracing::error!(error = ?e, "could not start jemalloc metrics loop");
     138            0 :             None
     139              :         }
     140              :     };
     141              : 
     142            0 :     let args = LocalProxyCliArgs::parse();
     143            0 :     let config = build_config(&args)?;
     144            0 :     let auth_backend = build_auth_backend(&args);
     145              : 
     146              :     // before we bind to any ports, write the process ID to a file
     147              :     // so that compute-ctl can find our process later
     148              :     // in order to trigger the appropriate SIGHUP on config change.
     149              :     //
     150              :     // This also claims a "lock" that makes sure only one instance
     151              :     // of local_proxy runs at a time.
     152            0 :     let _process_guard = loop {
     153            0 :         match pid_file::claim_for_current_process(&args.pid_path) {
     154            0 :             Ok(guard) => break guard,
     155            0 :             Err(e) => {
     156            0 :                 // compute-ctl might have tried to read the pid-file to let us
     157            0 :                 // know about some config change. We should try again.
     158            0 :                 error!(path=?args.pid_path, "could not claim PID file guard: {e:?}");
     159            0 :                 tokio::time::sleep(Duration::from_secs(1)).await;
     160              :             }
     161              :         }
     162              :     };
     163              : 
     164            0 :     let metrics_listener = TcpListener::bind(args.metrics).await?.into_std()?;
     165            0 :     let http_listener = TcpListener::bind(args.http).await?;
     166            0 :     let shutdown = CancellationToken::new();
     167            0 : 
     168            0 :     // todo: should scale with CU
     169            0 :     let endpoint_rate_limiter = Arc::new(EndpointRateLimiter::new_with_shards(
     170            0 :         LeakyBucketConfig {
     171            0 :             rps: 10.0,
     172            0 :             max: 100.0,
     173            0 :         },
     174            0 :         16,
     175            0 :     ));
     176            0 : 
     177            0 :     let mut maintenance_tasks = JoinSet::new();
     178            0 : 
     179            0 :     let refresh_config_notify = Arc::new(Notify::new());
     180            0 :     maintenance_tasks.spawn(crate::signals::handle(shutdown.clone(), {
     181            0 :         let refresh_config_notify = Arc::clone(&refresh_config_notify);
     182            0 :         move || {
     183            0 :             refresh_config_notify.notify_one();
     184            0 :         }
     185            0 :     }));
     186            0 : 
     187            0 :     // trigger the first config load **after** setting up the signal hook
     188            0 :     // to avoid the race condition where:
     189            0 :     // 1. No config file registered when local_proxy starts up
     190            0 :     // 2. The config file is written but the signal hook is not yet received
     191            0 :     // 3. local_proxy completes startup but has no config loaded, despite there being a registerd config.
     192            0 :     refresh_config_notify.notify_one();
     193            0 :     tokio::spawn(refresh_config_loop(args.config_path, refresh_config_notify));
     194            0 : 
     195            0 :     maintenance_tasks.spawn(crate::http::health_server::task_main(
     196            0 :         metrics_listener,
     197            0 :         AppMetrics {
     198            0 :             jemalloc,
     199            0 :             neon_metrics,
     200            0 :             proxy: crate::metrics::Metrics::get(),
     201            0 :         },
     202            0 :     ));
     203            0 : 
     204            0 :     let task = serverless::task_main(
     205            0 :         config,
     206            0 :         auth_backend,
     207            0 :         http_listener,
     208            0 :         shutdown.clone(),
     209            0 :         Arc::new(CancellationHandler::new(&config.connect_to_compute, None)),
     210            0 :         endpoint_rate_limiter,
     211            0 :     );
     212            0 : 
     213            0 :     match futures::future::select(pin!(maintenance_tasks.join_next()), pin!(task)).await {
     214              :         // exit immediately on maintenance task completion
     215            0 :         Either::Left((Some(res), _)) => match crate::error::flatten_err(res)? {},
     216              :         // exit with error immediately if all maintenance tasks have ceased (should be caught by branch above)
     217            0 :         Either::Left((None, _)) => bail!("no maintenance tasks running. invalid state"),
     218              :         // exit immediately on client task error
     219            0 :         Either::Right((res, _)) => res?,
     220              :     }
     221              : 
     222            0 :     Ok(())
     223            0 : }
     224              : 
     225              : /// ProxyConfig is created at proxy startup, and lives forever.
     226            0 : fn build_config(args: &LocalProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> {
     227              :     let config::ConcurrencyLockOptions {
     228            0 :         shards,
     229            0 :         limiter,
     230            0 :         epoch,
     231            0 :         timeout,
     232            0 :     } = args.connect_compute_lock.parse()?;
     233            0 :     info!(
     234              :         ?limiter,
     235              :         shards,
     236              :         ?epoch,
     237            0 :         "Using NodeLocks (connect_compute)"
     238              :     );
     239            0 :     let connect_compute_locks = ApiLocks::new(
     240            0 :         "connect_compute_lock",
     241            0 :         limiter,
     242            0 :         shards,
     243            0 :         timeout,
     244            0 :         epoch,
     245            0 :         &Metrics::get().proxy.connect_compute_lock,
     246            0 :     );
     247            0 : 
     248            0 :     let http_config = HttpConfig {
     249            0 :         accept_websockets: false,
     250            0 :         pool_options: GlobalConnPoolOptions {
     251            0 :             gc_epoch: Duration::from_secs(60),
     252            0 :             pool_shards: 2,
     253            0 :             idle_timeout: args.sql_over_http.sql_over_http_idle_timeout,
     254            0 :             opt_in: false,
     255            0 : 
     256            0 :             max_conns_per_endpoint: args.sql_over_http.sql_over_http_pool_max_total_conns,
     257            0 :             max_total_conns: args.sql_over_http.sql_over_http_pool_max_total_conns,
     258            0 :         },
     259            0 :         cancel_set: CancelSet::new(args.sql_over_http.sql_over_http_cancel_set_shards),
     260            0 :         client_conn_threshold: args.sql_over_http.sql_over_http_client_conn_threshold,
     261            0 :         max_request_size_bytes: args.sql_over_http.sql_over_http_max_request_size_bytes,
     262            0 :         max_response_size_bytes: args.sql_over_http.sql_over_http_max_response_size_bytes,
     263            0 :     };
     264              : 
     265            0 :     let compute_config = ComputeConfig {
     266            0 :         retry: RetryConfig::parse(RetryConfig::CONNECT_TO_COMPUTE_DEFAULT_VALUES)?,
     267            0 :         tls: Arc::new(compute_client_config_with_root_certs()?),
     268            0 :         timeout: Duration::from_secs(2),
     269            0 :     };
     270            0 : 
     271            0 :     Ok(Box::leak(Box::new(ProxyConfig {
     272            0 :         tls_config: None,
     273            0 :         metric_collection: None,
     274            0 :         http_config,
     275            0 :         authentication_config: AuthenticationConfig {
     276            0 :             jwks_cache: JwkCache::default(),
     277            0 :             thread_pool: ThreadPool::new(0),
     278            0 :             scram_protocol_timeout: Duration::from_secs(10),
     279            0 :             rate_limiter_enabled: false,
     280            0 :             rate_limiter: BucketRateLimiter::new(vec![]),
     281            0 :             rate_limit_ip_subnet: 64,
     282            0 :             ip_allowlist_check_enabled: true,
     283            0 :             is_vpc_acccess_proxy: false,
     284            0 :             is_auth_broker: false,
     285            0 :             accept_jwts: true,
     286            0 :             console_redirect_confirmation_timeout: Duration::ZERO,
     287            0 :         },
     288            0 :         proxy_protocol_v2: config::ProxyProtocolV2::Rejected,
     289            0 :         handshake_timeout: Duration::from_secs(10),
     290            0 :         region: "local".into(),
     291            0 :         wake_compute_retry_config: RetryConfig::parse(RetryConfig::WAKE_COMPUTE_DEFAULT_VALUES)?,
     292            0 :         connect_compute_locks,
     293            0 :         connect_to_compute: compute_config,
     294              :     })))
     295            0 : }
     296              : 
     297              : /// auth::Backend is created at proxy startup, and lives forever.
     298            0 : fn build_auth_backend(args: &LocalProxyCliArgs) -> &'static auth::Backend<'static, ()> {
     299            0 :     let auth_backend = crate::auth::Backend::Local(crate::auth::backend::MaybeOwned::Owned(
     300            0 :         LocalBackend::new(args.postgres, args.compute_ctl.clone()),
     301            0 :     ));
     302            0 : 
     303            0 :     Box::leak(Box::new(auth_backend))
     304            0 : }
     305              : 
     306              : #[derive(Error, Debug)]
     307              : enum RefreshConfigError {
     308              :     #[error(transparent)]
     309              :     Read(#[from] std::io::Error),
     310              :     #[error(transparent)]
     311              :     Parse(#[from] serde_json::Error),
     312              :     #[error(transparent)]
     313              :     Validate(anyhow::Error),
     314              : }
     315              : 
     316            0 : async fn refresh_config_loop(path: Utf8PathBuf, rx: Arc<Notify>) {
     317            0 :     let mut init = true;
     318              :     loop {
     319            0 :         rx.notified().await;
     320              : 
     321            0 :         match refresh_config_inner(&path).await {
     322            0 :             Ok(()) => {}
     323              :             // don't log for file not found errors if this is the first time we are checking
     324              :             // for computes that don't use local_proxy, this is not an error.
     325            0 :             Err(RefreshConfigError::Read(e))
     326            0 :                 if init && e.kind() == std::io::ErrorKind::NotFound =>
     327            0 :             {
     328            0 :                 debug!(error=?e, ?path, "could not read config file");
     329              :             }
     330            0 :             Err(e) => {
     331            0 :                 error!(error=?e, ?path, "could not read config file");
     332              :             }
     333              :         }
     334              : 
     335            0 :         init = false;
     336              :     }
     337              : }
     338              : 
     339            0 : async fn refresh_config_inner(path: &Utf8Path) -> Result<(), RefreshConfigError> {
     340            0 :     let bytes = tokio::fs::read(&path).await?;
     341            0 :     let data: LocalProxySpec = serde_json::from_slice(&bytes)?;
     342              : 
     343            0 :     let mut jwks_set = vec![];
     344              : 
     345            0 :     fn parse_jwks_settings(jwks: compute_api::spec::JwksSettings) -> anyhow::Result<JwksSettings> {
     346            0 :         let mut jwks_url = url::Url::from_str(&jwks.jwks_url).context("parsing JWKS url")?;
     347              : 
     348            0 :         ensure!(
     349            0 :             jwks_url.has_authority()
     350            0 :                 && (jwks_url.scheme() == "http" || jwks_url.scheme() == "https"),
     351            0 :             "Invalid JWKS url. Must be HTTP",
     352              :         );
     353              : 
     354            0 :         ensure!(
     355            0 :             jwks_url.host().is_some_and(|h| h != url::Host::Domain("")),
     356            0 :             "Invalid JWKS url. No domain listed",
     357              :         );
     358              : 
     359              :         // clear username, password and ports
     360            0 :         jwks_url
     361            0 :             .set_username("")
     362            0 :             .expect("url can be a base and has a valid host and is not a file. should not error");
     363            0 :         jwks_url
     364            0 :             .set_password(None)
     365            0 :             .expect("url can be a base and has a valid host and is not a file. should not error");
     366            0 :         // local testing is hard if we need to have a specific restricted port
     367            0 :         if cfg!(not(feature = "testing")) {
     368            0 :             jwks_url.set_port(None).expect(
     369            0 :                 "url can be a base and has a valid host and is not a file. should not error",
     370            0 :             );
     371            0 :         }
     372              : 
     373              :         // clear query params
     374            0 :         jwks_url.set_fragment(None);
     375            0 :         jwks_url.query_pairs_mut().clear().finish();
     376            0 : 
     377            0 :         if jwks_url.scheme() != "https" {
     378              :             // local testing is hard if we need to set up https support.
     379            0 :             if cfg!(not(feature = "testing")) {
     380            0 :                 jwks_url
     381            0 :                     .set_scheme("https")
     382            0 :                     .expect("should not error to set the scheme to https if it was http");
     383            0 :             } else {
     384            0 :                 warn!(scheme = jwks_url.scheme(), "JWKS url is not HTTPS");
     385              :             }
     386            0 :         }
     387              : 
     388            0 :         Ok(JwksSettings {
     389            0 :             id: jwks.id,
     390            0 :             jwks_url,
     391            0 :             _provider_name: jwks.provider_name,
     392            0 :             jwt_audience: jwks.jwt_audience,
     393            0 :             role_names: jwks
     394            0 :                 .role_names
     395            0 :                 .into_iter()
     396            0 :                 .map(RoleName::from)
     397            0 :                 .map(|s| RoleNameInt::from(&s))
     398            0 :                 .collect(),
     399            0 :         })
     400            0 :     }
     401              : 
     402            0 :     for jwks in data.jwks.into_iter().flatten() {
     403            0 :         jwks_set.push(parse_jwks_settings(jwks).map_err(RefreshConfigError::Validate)?);
     404              :     }
     405              : 
     406            0 :     info!("successfully loaded new config");
     407            0 :     JWKS_ROLE_MAP.store(Some(Arc::new(EndpointJwksResponse { jwks: jwks_set })));
     408            0 : 
     409            0 :     Ok(())
     410            0 : }
        

Generated by: LCOV version 2.1-beta