LCOV - code coverage report
Current view: top level - proxy/src - compute.rs (source / functions) Coverage Total Hit
Test: 07bee600374ccd486c69370d0972d9035964fe68.info Lines: 23.0 % 196 45
Test Date: 2025-02-20 13:11:02 Functions: 35.0 % 20 7

            Line data    Source code
       1              : use std::io;
       2              : use std::net::SocketAddr;
       3              : use std::time::Duration;
       4              : 
       5              : use futures::{FutureExt, TryFutureExt};
       6              : use itertools::Itertools;
       7              : use postgres_client::tls::MakeTlsConnect;
       8              : use postgres_client::{CancelToken, RawConnection};
       9              : use postgres_protocol::message::backend::NoticeResponseBody;
      10              : use pq_proto::StartupMessageParams;
      11              : use rustls::pki_types::InvalidDnsNameError;
      12              : use thiserror::Error;
      13              : use tokio::net::TcpStream;
      14              : use tracing::{debug, error, info, warn};
      15              : 
      16              : use crate::auth::backend::ComputeUserInfo;
      17              : use crate::auth::parse_endpoint_param;
      18              : use crate::cancellation::CancelClosure;
      19              : use crate::config::ComputeConfig;
      20              : use crate::context::RequestContext;
      21              : use crate::control_plane::client::ApiLockError;
      22              : use crate::control_plane::errors::WakeComputeError;
      23              : use crate::control_plane::messages::MetricsAuxInfo;
      24              : use crate::error::{ReportableError, UserFacingError};
      25              : use crate::metrics::{Metrics, NumDbConnectionsGuard};
      26              : use crate::proxy::neon_option;
      27              : use crate::tls::postgres_rustls::MakeRustlsConnect;
      28              : use crate::types::Host;
      29              : 
      30              : pub const COULD_NOT_CONNECT: &str = "Couldn't connect to compute node";
      31              : 
      32              : #[derive(Debug, Error)]
      33              : pub(crate) enum ConnectionError {
      34              :     /// This error doesn't seem to reveal any secrets; for instance,
      35              :     /// `postgres_client::error::Kind` doesn't contain ip addresses and such.
      36              :     #[error("{COULD_NOT_CONNECT}: {0}")]
      37              :     Postgres(#[from] postgres_client::Error),
      38              : 
      39              :     #[error("{COULD_NOT_CONNECT}: {0}")]
      40              :     CouldNotConnect(#[from] io::Error),
      41              : 
      42              :     #[error("{COULD_NOT_CONNECT}: {0}")]
      43              :     TlsError(#[from] InvalidDnsNameError),
      44              : 
      45              :     #[error("{COULD_NOT_CONNECT}: {0}")]
      46              :     WakeComputeError(#[from] WakeComputeError),
      47              : 
      48              :     #[error("error acquiring resource permit: {0}")]
      49              :     TooManyConnectionAttempts(#[from] ApiLockError),
      50              : }
      51              : 
      52              : impl UserFacingError for ConnectionError {
      53            0 :     fn to_string_client(&self) -> String {
      54            0 :         match self {
      55              :             // This helps us drop irrelevant library-specific prefixes.
      56              :             // TODO: propagate severity level and other parameters.
      57            0 :             ConnectionError::Postgres(err) => match err.as_db_error() {
      58            0 :                 Some(err) => {
      59            0 :                     let msg = err.message();
      60            0 : 
      61            0 :                     if msg.starts_with("unsupported startup parameter: ")
      62            0 :                         || msg.starts_with("unsupported startup parameter in options: ")
      63              :                     {
      64            0 :                         format!("{msg}. Please use unpooled connection or remove this parameter from the startup package. More details: https://neon.tech/docs/connect/connection-errors#unsupported-startup-parameter")
      65              :                     } else {
      66            0 :                         msg.to_owned()
      67              :                     }
      68              :                 }
      69            0 :                 None => err.to_string(),
      70              :             },
      71            0 :             ConnectionError::WakeComputeError(err) => err.to_string_client(),
      72              :             ConnectionError::TooManyConnectionAttempts(_) => {
      73            0 :                 "Failed to acquire permit to connect to the database. Too many database connection attempts are currently ongoing.".to_owned()
      74              :             }
      75            0 :             _ => COULD_NOT_CONNECT.to_owned(),
      76              :         }
      77            0 :     }
      78              : }
      79              : 
      80              : impl ReportableError for ConnectionError {
      81            0 :     fn get_error_kind(&self) -> crate::error::ErrorKind {
      82            0 :         match self {
      83            0 :             ConnectionError::Postgres(e) if e.as_db_error().is_some() => {
      84            0 :                 crate::error::ErrorKind::Postgres
      85              :             }
      86            0 :             ConnectionError::Postgres(_) => crate::error::ErrorKind::Compute,
      87            0 :             ConnectionError::CouldNotConnect(_) => crate::error::ErrorKind::Compute,
      88            0 :             ConnectionError::TlsError(_) => crate::error::ErrorKind::Compute,
      89            0 :             ConnectionError::WakeComputeError(e) => e.get_error_kind(),
      90            0 :             ConnectionError::TooManyConnectionAttempts(e) => e.get_error_kind(),
      91              :         }
      92            0 :     }
      93              : }
      94              : 
      95              : /// A pair of `ClientKey` & `ServerKey` for `SCRAM-SHA-256`.
      96              : pub(crate) type ScramKeys = postgres_client::config::ScramKeys<32>;
      97              : 
      98              : /// A config for establishing a connection to compute node.
      99              : /// Eventually, `postgres_client` will be replaced with something better.
     100              : /// Newtype allows us to implement methods on top of it.
     101              : #[derive(Clone)]
     102              : pub(crate) struct ConnCfg(Box<postgres_client::Config>);
     103              : 
     104              : /// Creation and initialization routines.
     105              : impl ConnCfg {
     106           10 :     pub(crate) fn new(host: String, port: u16) -> Self {
     107           10 :         Self(Box::new(postgres_client::Config::new(host, port)))
     108           10 :     }
     109              : 
     110              :     /// Reuse password or auth keys from the other config.
     111            4 :     pub(crate) fn reuse_password(&mut self, other: Self) {
     112            4 :         if let Some(password) = other.get_password() {
     113            4 :             self.password(password);
     114            4 :         }
     115              : 
     116            4 :         if let Some(keys) = other.get_auth_keys() {
     117            0 :             self.auth_keys(keys);
     118            4 :         }
     119            4 :     }
     120              : 
     121            0 :     pub(crate) fn get_host(&self) -> Host {
     122            0 :         match self.0.get_host() {
     123            0 :             postgres_client::config::Host::Tcp(s) => s.into(),
     124            0 :         }
     125            0 :     }
     126              : 
     127              :     /// Apply startup message params to the connection config.
     128            0 :     pub(crate) fn set_startup_params(
     129            0 :         &mut self,
     130            0 :         params: &StartupMessageParams,
     131            0 :         arbitrary_params: bool,
     132            0 :     ) {
     133            0 :         if !arbitrary_params {
     134            0 :             self.set_param("client_encoding", "UTF8");
     135            0 :         }
     136            0 :         for (k, v) in params.iter() {
     137            0 :             match k {
     138              :                 // Only set `user` if it's not present in the config.
     139              :                 // Console redirect auth flow takes username from the console's response.
     140            0 :                 "user" if self.user_is_set() => {}
     141            0 :                 "database" if self.db_is_set() => {}
     142            0 :                 "options" => {
     143            0 :                     if let Some(options) = filtered_options(v) {
     144            0 :                         self.set_param(k, &options);
     145            0 :                     }
     146              :                 }
     147            0 :                 "user" | "database" | "application_name" | "replication" => {
     148            0 :                     self.set_param(k, v);
     149            0 :                 }
     150              : 
     151              :                 // if we allow arbitrary params, then we forward them through.
     152              :                 // this is a flag for a period of backwards compatibility
     153            0 :                 k if arbitrary_params => {
     154            0 :                     self.set_param(k, v);
     155            0 :                 }
     156            0 :                 _ => {}
     157              :             }
     158              :         }
     159            0 :     }
     160              : }
     161              : 
     162              : impl std::ops::Deref for ConnCfg {
     163              :     type Target = postgres_client::Config;
     164              : 
     165            8 :     fn deref(&self) -> &Self::Target {
     166            8 :         &self.0
     167            8 :     }
     168              : }
     169              : 
     170              : /// For now, let's make it easier to setup the config.
     171              : impl std::ops::DerefMut for ConnCfg {
     172           10 :     fn deref_mut(&mut self) -> &mut Self::Target {
     173           10 :         &mut self.0
     174           10 :     }
     175              : }
     176              : 
     177              : impl ConnCfg {
     178              :     /// Establish a raw TCP connection to the compute node.
     179            0 :     async fn connect_raw(&self, timeout: Duration) -> io::Result<(SocketAddr, TcpStream, &str)> {
     180              :         use postgres_client::config::Host;
     181              : 
     182              :         // wrap TcpStream::connect with timeout
     183            0 :         let connect_with_timeout = |host, port| {
     184            0 :             tokio::time::timeout(timeout, TcpStream::connect((host, port))).map(
     185            0 :                 move |res| match res {
     186            0 :                     Ok(tcpstream_connect_res) => tcpstream_connect_res,
     187            0 :                     Err(_) => Err(io::Error::new(
     188            0 :                         io::ErrorKind::TimedOut,
     189            0 :                         format!("exceeded connection timeout {timeout:?}"),
     190            0 :                     )),
     191            0 :                 },
     192            0 :             )
     193            0 :         };
     194              : 
     195            0 :         let connect_once = |host, port| {
     196            0 :             debug!("trying to connect to compute node at {host}:{port}");
     197            0 :             connect_with_timeout(host, port).and_then(|stream| async {
     198            0 :                 let socket_addr = stream.peer_addr()?;
     199            0 :                 let socket = socket2::SockRef::from(&stream);
     200            0 :                 // Disable Nagle's algorithm to not introduce latency between
     201            0 :                 // client and compute.
     202            0 :                 socket.set_nodelay(true)?;
     203              :                 // This prevents load balancer from severing the connection.
     204            0 :                 socket.set_keepalive(true)?;
     205            0 :                 Ok((socket_addr, stream))
     206            0 :             })
     207            0 :         };
     208              : 
     209              :         // We can't reuse connection establishing logic from `postgres_client` here,
     210              :         // because it has no means for extracting the underlying socket which we
     211              :         // require for our business.
     212            0 :         let port = self.0.get_port();
     213            0 :         let host = self.0.get_host();
     214            0 : 
     215            0 :         let host = match host {
     216            0 :             Host::Tcp(host) => host.as_str(),
     217            0 :         };
     218            0 : 
     219            0 :         match connect_once(host, port).await {
     220            0 :             Ok((sockaddr, stream)) => Ok((sockaddr, stream, host)),
     221            0 :             Err(err) => {
     222            0 :                 warn!("couldn't connect to compute node at {host}:{port}: {err}");
     223            0 :                 Err(err)
     224              :             }
     225              :         }
     226            0 :     }
     227              : }
     228              : 
     229              : type RustlsStream = <MakeRustlsConnect as MakeTlsConnect<tokio::net::TcpStream>>::Stream;
     230              : 
     231              : pub(crate) struct PostgresConnection {
     232              :     /// Socket connected to a compute node.
     233              :     pub(crate) stream:
     234              :         postgres_client::maybe_tls_stream::MaybeTlsStream<tokio::net::TcpStream, RustlsStream>,
     235              :     /// PostgreSQL connection parameters.
     236              :     pub(crate) params: std::collections::HashMap<String, String>,
     237              :     /// Query cancellation token.
     238              :     pub(crate) cancel_closure: CancelClosure,
     239              :     /// Labels for proxy's metrics.
     240              :     pub(crate) aux: MetricsAuxInfo,
     241              :     /// Notices received from compute after authenticating
     242              :     pub(crate) delayed_notice: Vec<NoticeResponseBody>,
     243              : 
     244              :     _guage: NumDbConnectionsGuard<'static>,
     245              : }
     246              : 
     247              : impl ConnCfg {
     248              :     /// Connect to a corresponding compute node.
     249            0 :     pub(crate) async fn connect(
     250            0 :         &self,
     251            0 :         ctx: &RequestContext,
     252            0 :         aux: MetricsAuxInfo,
     253            0 :         config: &ComputeConfig,
     254            0 :         user_info: ComputeUserInfo,
     255            0 :     ) -> Result<PostgresConnection, ConnectionError> {
     256            0 :         let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
     257            0 :         let (socket_addr, stream, host) = self.connect_raw(config.timeout).await?;
     258            0 :         drop(pause);
     259            0 : 
     260            0 :         let mut mk_tls = crate::tls::postgres_rustls::MakeRustlsConnect::new(config.tls.clone());
     261            0 :         let tls = <MakeRustlsConnect as MakeTlsConnect<tokio::net::TcpStream>>::make_tls_connect(
     262            0 :             &mut mk_tls,
     263            0 :             host,
     264            0 :         )?;
     265              : 
     266              :         // connect_raw() will not use TLS if sslmode is "disable"
     267            0 :         let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
     268            0 :         let connection = self.0.connect_raw(stream, tls).await?;
     269            0 :         drop(pause);
     270            0 : 
     271            0 :         let RawConnection {
     272            0 :             stream,
     273            0 :             parameters,
     274            0 :             delayed_notice,
     275            0 :             process_id,
     276            0 :             secret_key,
     277            0 :         } = connection;
     278            0 : 
     279            0 :         tracing::Span::current().record("pid", tracing::field::display(process_id));
     280            0 :         let stream = stream.into_inner();
     281            0 : 
     282            0 :         // TODO: lots of useful info but maybe we can move it elsewhere (eg traces?)
     283            0 :         info!(
     284            0 :             cold_start_info = ctx.cold_start_info().as_str(),
     285            0 :             "connected to compute node at {host} ({socket_addr}) sslmode={:?}",
     286            0 :             self.0.get_ssl_mode()
     287              :         );
     288              : 
     289              :         // NB: CancelToken is supposed to hold socket_addr, but we use connect_raw.
     290              :         // Yet another reason to rework the connection establishing code.
     291            0 :         let cancel_closure = CancelClosure::new(
     292            0 :             socket_addr,
     293            0 :             CancelToken {
     294            0 :                 socket_config: None,
     295            0 :                 ssl_mode: self.0.get_ssl_mode(),
     296            0 :                 process_id,
     297            0 :                 secret_key,
     298            0 :             },
     299            0 :             host.to_string(),
     300            0 :             user_info,
     301            0 :         );
     302            0 : 
     303            0 :         let connection = PostgresConnection {
     304            0 :             stream,
     305            0 :             params: parameters,
     306            0 :             delayed_notice,
     307            0 :             cancel_closure,
     308            0 :             aux,
     309            0 :             _guage: Metrics::get().proxy.db_connections.guard(ctx.protocol()),
     310            0 :         };
     311            0 : 
     312            0 :         Ok(connection)
     313            0 :     }
     314              : }
     315              : 
     316              : /// Retrieve `options` from a startup message, dropping all proxy-secific flags.
     317            6 : fn filtered_options(options: &str) -> Option<String> {
     318            6 :     #[allow(unstable_name_collisions)]
     319            6 :     let options: String = StartupMessageParams::parse_options_raw(options)
     320           14 :         .filter(|opt| parse_endpoint_param(opt).is_none() && neon_option(opt).is_none())
     321            6 :         .intersperse(" ") // TODO: use impl from std once it's stabilized
     322            6 :         .collect();
     323            6 : 
     324            6 :     // Don't even bother with empty options.
     325            6 :     if options.is_empty() {
     326            3 :         return None;
     327            3 :     }
     328            3 : 
     329            3 :     Some(options)
     330            6 : }
     331              : 
     332              : #[cfg(test)]
     333              : mod tests {
     334              :     use super::*;
     335              : 
     336              :     #[test]
     337            1 :     fn test_filtered_options() {
     338            1 :         // Empty options is unlikely to be useful anyway.
     339            1 :         let params = "";
     340            1 :         assert_eq!(filtered_options(params), None);
     341              : 
     342              :         // It's likely that clients will only use options to specify endpoint/project.
     343            1 :         let params = "project=foo";
     344            1 :         assert_eq!(filtered_options(params), None);
     345              : 
     346              :         // Same, because unescaped whitespaces are no-op.
     347            1 :         let params = " project=foo ";
     348            1 :         assert_eq!(filtered_options(params).as_deref(), None);
     349              : 
     350            1 :         let params = r"\  project=foo \ ";
     351            1 :         assert_eq!(filtered_options(params).as_deref(), Some(r"\  \ "));
     352              : 
     353            1 :         let params = "project = foo";
     354            1 :         assert_eq!(filtered_options(params).as_deref(), Some("project = foo"));
     355              : 
     356            1 :         let params = "project = foo neon_endpoint_type:read_write   neon_lsn:0/2 neon_proxy_params_compat:true";
     357            1 :         assert_eq!(filtered_options(params).as_deref(), Some("project = foo"));
     358            1 :     }
     359              : }
        

Generated by: LCOV version 2.1-beta