LCOV - code coverage report
Current view: top level - proxy/src/auth/backend - console_redirect.rs (source / functions) Coverage Total Hit
Test: f5f94ec0366b63fd2cbbe02edc2087dbd893d04d.info Lines: 0.0 % 112 0
Test Date: 2024-11-20 05:34:23 Functions: 0.0 % 23 0

            Line data    Source code
       1              : use async_trait::async_trait;
       2              : use pq_proto::BeMessage as Be;
       3              : use thiserror::Error;
       4              : use tokio::io::{AsyncRead, AsyncWrite};
       5              : use tokio_postgres::config::SslMode;
       6              : use tracing::{info, info_span};
       7              : 
       8              : use super::ComputeCredentialKeys;
       9              : use crate::cache::Cached;
      10              : use crate::config::AuthenticationConfig;
      11              : use crate::context::RequestMonitoring;
      12              : use crate::control_plane::{self, CachedNodeInfo, NodeInfo};
      13              : use crate::error::{ReportableError, UserFacingError};
      14              : use crate::proxy::connect_compute::ComputeConnectBackend;
      15              : use crate::stream::PqStream;
      16              : use crate::{auth, compute, waiters};
      17              : 
      18            0 : #[derive(Debug, Error)]
      19              : pub(crate) enum ConsoleRedirectError {
      20              :     #[error(transparent)]
      21              :     WaiterRegister(#[from] waiters::RegisterError),
      22              : 
      23              :     #[error(transparent)]
      24              :     WaiterWait(#[from] waiters::WaitError),
      25              : 
      26              :     #[error(transparent)]
      27              :     Io(#[from] std::io::Error),
      28              : }
      29              : 
      30              : #[derive(Debug)]
      31              : pub struct ConsoleRedirectBackend {
      32              :     console_uri: reqwest::Url,
      33              : }
      34              : 
      35              : impl UserFacingError for ConsoleRedirectError {
      36            0 :     fn to_string_client(&self) -> String {
      37            0 :         "Internal error".to_string()
      38            0 :     }
      39              : }
      40              : 
      41              : impl ReportableError for ConsoleRedirectError {
      42            0 :     fn get_error_kind(&self) -> crate::error::ErrorKind {
      43            0 :         match self {
      44            0 :             Self::WaiterRegister(_) => crate::error::ErrorKind::Service,
      45            0 :             Self::WaiterWait(_) => crate::error::ErrorKind::Service,
      46            0 :             Self::Io(_) => crate::error::ErrorKind::ClientDisconnect,
      47              :         }
      48            0 :     }
      49              : }
      50              : 
      51            0 : fn hello_message(redirect_uri: &reqwest::Url, session_id: &str) -> String {
      52            0 :     format!(
      53            0 :         concat![
      54            0 :             "Welcome to Neon!\n",
      55            0 :             "Authenticate by visiting:\n",
      56            0 :             "    {redirect_uri}{session_id}\n\n",
      57            0 :         ],
      58            0 :         redirect_uri = redirect_uri,
      59            0 :         session_id = session_id,
      60            0 :     )
      61            0 : }
      62              : 
      63            0 : pub(crate) fn new_psql_session_id() -> String {
      64            0 :     hex::encode(rand::random::<[u8; 8]>())
      65            0 : }
      66              : 
      67              : impl ConsoleRedirectBackend {
      68            0 :     pub fn new(console_uri: reqwest::Url) -> Self {
      69            0 :         Self { console_uri }
      70            0 :     }
      71              : 
      72            0 :     pub(crate) async fn authenticate(
      73            0 :         &self,
      74            0 :         ctx: &RequestMonitoring,
      75            0 :         auth_config: &'static AuthenticationConfig,
      76            0 :         client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
      77            0 :     ) -> auth::Result<ConsoleRedirectNodeInfo> {
      78            0 :         authenticate(ctx, auth_config, &self.console_uri, client)
      79            0 :             .await
      80            0 :             .map(ConsoleRedirectNodeInfo)
      81            0 :     }
      82              : }
      83              : 
      84              : pub struct ConsoleRedirectNodeInfo(pub(super) NodeInfo);
      85              : 
      86              : #[async_trait]
      87              : impl ComputeConnectBackend for ConsoleRedirectNodeInfo {
      88            0 :     async fn wake_compute(
      89            0 :         &self,
      90            0 :         _ctx: &RequestMonitoring,
      91            0 :     ) -> Result<CachedNodeInfo, control_plane::errors::WakeComputeError> {
      92            0 :         Ok(Cached::new_uncached(self.0.clone()))
      93            0 :     }
      94              : 
      95            0 :     fn get_keys(&self) -> &ComputeCredentialKeys {
      96            0 :         &ComputeCredentialKeys::None
      97            0 :     }
      98              : }
      99              : 
     100            0 : async fn authenticate(
     101            0 :     ctx: &RequestMonitoring,
     102            0 :     auth_config: &'static AuthenticationConfig,
     103            0 :     link_uri: &reqwest::Url,
     104            0 :     client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
     105            0 : ) -> auth::Result<NodeInfo> {
     106            0 :     ctx.set_auth_method(crate::context::AuthMethod::ConsoleRedirect);
     107              : 
     108              :     // registering waiter can fail if we get unlucky with rng.
     109              :     // just try again.
     110            0 :     let (psql_session_id, waiter) = loop {
     111            0 :         let psql_session_id = new_psql_session_id();
     112            0 : 
     113            0 :         match control_plane::mgmt::get_waiter(&psql_session_id) {
     114            0 :             Ok(waiter) => break (psql_session_id, waiter),
     115            0 :             Err(_e) => continue,
     116              :         }
     117              :     };
     118              : 
     119            0 :     let span = info_span!("console_redirect", psql_session_id = &psql_session_id);
     120            0 :     let greeting = hello_message(link_uri, &psql_session_id);
     121            0 : 
     122            0 :     // Give user a URL to spawn a new database.
     123            0 :     info!(parent: &span, "sending the auth URL to the user");
     124            0 :     client
     125            0 :         .write_message_noflush(&Be::AuthenticationOk)?
     126            0 :         .write_message_noflush(&Be::CLIENT_ENCODING)?
     127            0 :         .write_message(&Be::NoticeResponse(&greeting))
     128            0 :         .await?;
     129              : 
     130              :     // Wait for console response via control plane (see `mgmt`).
     131            0 :     info!(parent: &span, "waiting for console's reply...");
     132            0 :     let db_info = tokio::time::timeout(auth_config.console_redirect_confirmation_timeout, waiter)
     133            0 :         .await
     134            0 :         .map_err(|_elapsed| {
     135            0 :             auth::AuthError::confirmation_timeout(
     136            0 :                 auth_config.console_redirect_confirmation_timeout.into(),
     137            0 :             )
     138            0 :         })?
     139            0 :         .map_err(ConsoleRedirectError::from)?;
     140              : 
     141            0 :     if auth_config.ip_allowlist_check_enabled {
     142            0 :         if let Some(allowed_ips) = &db_info.allowed_ips {
     143            0 :             if !auth::check_peer_addr_is_in_list(&ctx.peer_addr(), allowed_ips) {
     144            0 :                 return Err(auth::AuthError::ip_address_not_allowed(ctx.peer_addr()));
     145            0 :             }
     146            0 :         }
     147            0 :     }
     148              : 
     149            0 :     client.write_message_noflush(&Be::NoticeResponse("Connecting to database."))?;
     150              : 
     151              :     // This config should be self-contained, because we won't
     152              :     // take username or dbname from client's startup message.
     153            0 :     let mut config = compute::ConnCfg::new();
     154            0 :     config
     155            0 :         .host(&db_info.host)
     156            0 :         .port(db_info.port)
     157            0 :         .dbname(&db_info.dbname)
     158            0 :         .user(&db_info.user);
     159            0 : 
     160            0 :     ctx.set_dbname(db_info.dbname.into());
     161            0 :     ctx.set_user(db_info.user.into());
     162            0 :     ctx.set_project(db_info.aux.clone());
     163            0 :     info!("woken up a compute node");
     164              : 
     165              :     // Backwards compatibility. pg_sni_proxy uses "--" in domain names
     166              :     // while direct connections do not. Once we migrate to pg_sni_proxy
     167              :     // everywhere, we can remove this.
     168            0 :     if db_info.host.contains("--") {
     169            0 :         // we need TLS connection with SNI info to properly route it
     170            0 :         config.ssl_mode(SslMode::Require);
     171            0 :     } else {
     172            0 :         config.ssl_mode(SslMode::Disable);
     173            0 :     }
     174              : 
     175            0 :     if let Some(password) = db_info.password {
     176            0 :         config.password(password.as_ref());
     177            0 :     }
     178              : 
     179            0 :     Ok(NodeInfo {
     180            0 :         config,
     181            0 :         aux: db_info.aux,
     182            0 :         allow_self_signed_compute: false, // caller may override
     183            0 :     })
     184            0 : }
        

Generated by: LCOV version 2.1-beta