Line data    Source code

       1              : use crate::{
       2              :     auth, compute,
       3              :     console::{self, provider::NodeInfo},
       4              :     context::RequestMonitoring,
       5              :     error::UserFacingError,
       6              :     stream::PqStream,
       7              :     waiters,
       8              : };
       9              : use pq_proto::BeMessage as Be;
      10              : use thiserror::Error;
      11              : use tokio::io::{AsyncRead, AsyncWrite};
      12              : use tokio_postgres::config::SslMode;
      13              : use tracing::{info, info_span};
      14              : 
      15            0 : #[derive(Debug, Error)]
      16              : pub enum LinkAuthError {
      17              :     /// Authentication error reported by the console.
      18              :     #[error("Authentication failed: {0}")]
      19              :     AuthFailed(String),
      20              : 
      21              :     #[error(transparent)]
      22              :     WaiterRegister(#[from] waiters::RegisterError),
      23              : 
      24              :     #[error(transparent)]
      25              :     WaiterWait(#[from] waiters::WaitError),
      26              : 
      27              :     #[error(transparent)]
      28              :     Io(#[from] std::io::Error),
      29              : }
      30              : 
      31              : impl UserFacingError for LinkAuthError {
      32            0 :     fn to_string_client(&self) -> String {
      33            0 :         use LinkAuthError::*;
      34            0 :         match self {
      35            0 :             AuthFailed(_) => self.to_string(),
      36            0 :             _ => "Internal error".to_string(),
      37              :         }
      38            0 :     }
      39              : }
      40              : 
      41            3 : fn hello_message(redirect_uri: &reqwest::Url, session_id: &str) -> String {
      42            3 :     format!(
      43            3 :         concat![
      44            3 :             "Welcome to Neon!\n",
      45            3 :             "Authenticate by visiting:\n",
      46            3 :             "    {redirect_uri}{session_id}\n\n",
      47            3 :         ],
      48            3 :         redirect_uri = redirect_uri,
      49            3 :         session_id = session_id,
      50            3 :     )
      51            3 : }
      52              : 
      53            3 : pub fn new_psql_session_id() -> String {
      54            3 :     hex::encode(rand::random::<[u8; 8]>())
      55            3 : }
      56              : 
      57            3 : pub(super) async fn authenticate(
      58            3 :     ctx: &mut RequestMonitoring,
      59            3 :     link_uri: &reqwest::Url,
      60            3 :     client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
      61            3 : ) -> auth::Result<NodeInfo> {
      62              :     // registering waiter can fail if we get unlucky with rng.
      63              :     // just try again.
      64            3 :     let (psql_session_id, waiter) = loop {
      65            3 :         let psql_session_id = new_psql_session_id();
      66            3 : 
      67            3 :         match console::mgmt::get_waiter(&psql_session_id) {
      68            3 :             Ok(waiter) => break (psql_session_id, waiter),
      69            0 :             Err(_e) => continue,
      70              :         }
      71              :     };
      72              : 
      73            3 :     let span = info_span!("link", psql_session_id = &psql_session_id);
      74            3 :     let greeting = hello_message(link_uri, &psql_session_id);
      75            3 : 
      76            3 :     // Give user a URL to spawn a new database.
      77            3 :     info!(parent: &span, "sending the auth URL to the user");
      78            3 :     client
      79            3 :         .write_message_noflush(&Be::AuthenticationOk)?
      80            3 :         .write_message_noflush(&Be::CLIENT_ENCODING)?
      81            3 :         .write_message(&Be::NoticeResponse(&greeting))
      82            0 :         .await?;
      83              : 
      84              :     // Wait for web console response (see `mgmt`).
      85            3 :     info!(parent: &span, "waiting for console's reply...");
      86            3 :     let db_info = waiter.await.map_err(LinkAuthError::from)?;
      87              : 
      88            3 :     client.write_message_noflush(&Be::NoticeResponse("Connecting to database."))?;
      89              : 
      90              :     // This config should be self-contained, because we won't
      91              :     // take username or dbname from client's startup message.
      92            3 :     let mut config = compute::ConnCfg::new();
      93            3 :     config
      94            3 :         .host(&db_info.host)
      95            3 :         .port(db_info.port)
      96            3 :         .dbname(&db_info.dbname)
      97            3 :         .user(&db_info.user);
      98            3 : 
      99            3 :     ctx.set_user(db_info.user.into());
     100            3 :     ctx.set_project(db_info.aux.clone());
     101            3 :     tracing::Span::current().record("ep", &tracing::field::display(&db_info.aux.endpoint_id));
     102            3 : 
     103            3 :     // Backwards compatibility. pg_sni_proxy uses "--" in domain names
     104            3 :     // while direct connections do not. Once we migrate to pg_sni_proxy
     105            3 :     // everywhere, we can remove this.
     106            3 :     if db_info.host.contains("--") {
     107            0 :         // we need TLS connection with SNI info to properly route it
     108            0 :         config.ssl_mode(SslMode::Require);
     109            3 :     } else {
     110            3 :         config.ssl_mode(SslMode::Disable);
     111            3 :     }
     112              : 
     113            3 :     if let Some(password) = db_info.password {
     114            0 :         config.password(password.as_ref());
     115            3 :     }
     116              : 
     117            3 :     Ok(NodeInfo {
     118            3 :         config,
     119            3 :         aux: db_info.aux,
     120            3 :         allow_self_signed_compute: false, // caller may override
     121            3 :     })
     122            3 : }