LCOV - code coverage report
Current view: top level - proxy/src/auth - backend.rs (source / functions) Coverage Total Hit
Test: c639aa5f7ab62b43d647b10f40d15a15686ce8a9.info Lines: 94.0 % 182 171
Test Date: 2024-02-12 20:26:03 Functions: 45.7 % 70 32

            Line data    Source code
       1              : mod classic;
       2              : mod hacks;
       3              : mod link;
       4              : 
       5              : pub use link::LinkAuthError;
       6              : use tokio_postgres::config::AuthKeys;
       7              : 
       8              : use crate::auth::credentials::check_peer_addr_is_in_list;
       9              : use crate::auth::validate_password_and_exchange;
      10              : use crate::cache::Cached;
      11              : use crate::console::errors::GetAuthInfoError;
      12              : use crate::console::provider::{CachedRoleSecret, ConsoleBackend};
      13              : use crate::console::AuthSecret;
      14              : use crate::context::RequestMonitoring;
      15              : use crate::proxy::wake_compute::wake_compute;
      16              : use crate::proxy::NeonOptions;
      17              : use crate::stream::Stream;
      18              : use crate::{
      19              :     auth::{self, ComputeUserInfoMaybeEndpoint},
      20              :     config::AuthenticationConfig,
      21              :     console::{
      22              :         self,
      23              :         provider::{CachedAllowedIps, CachedNodeInfo},
      24              :         Api,
      25              :     },
      26              :     stream, url,
      27              : };
      28              : use crate::{scram, EndpointCacheKey, EndpointId, RoleName};
      29              : use futures::TryFutureExt;
      30              : use std::sync::Arc;
      31              : use tokio::io::{AsyncRead, AsyncWrite};
      32              : use tracing::info;
      33              : 
      34              : /// Alternative to [`std::borrow::Cow`] but doesn't need `T: ToOwned` as we don't need that functionality
      35              : pub enum MaybeOwned<'a, T> {
      36              :     Owned(T),
      37              :     Borrowed(&'a T),
      38              : }
      39              : 
      40              : impl<T> std::ops::Deref for MaybeOwned<'_, T> {
      41              :     type Target = T;
      42              : 
      43          421 :     fn deref(&self) -> &Self::Target {
      44          421 :         match self {
      45          198 :             MaybeOwned::Owned(t) => t,
      46          223 :             MaybeOwned::Borrowed(t) => t,
      47              :         }
      48          421 :     }
      49              : }
      50              : 
      51              : /// This type serves two purposes:
      52              : ///
      53              : /// * When `T` is `()`, it's just a regular auth backend selector
      54              : ///   which we use in [`crate::config::ProxyConfig`].
      55              : ///
      56              : /// * However, when we substitute `T` with [`ComputeUserInfoMaybeEndpoint`],
      57              : ///   this helps us provide the credentials only to those auth
      58              : ///   backends which require them for the authentication process.
      59              : pub enum BackendType<'a, T> {
      60              :     /// Cloud API (V2).
      61              :     Console(MaybeOwned<'a, ConsoleBackend>, T),
      62              :     /// Authentication via a web browser.
      63              :     Link(MaybeOwned<'a, url::ApiUrl>),
      64              : }
      65              : 
      66              : pub trait TestBackend: Send + Sync + 'static {
      67              :     fn wake_compute(&self) -> Result<CachedNodeInfo, console::errors::WakeComputeError>;
      68              :     fn get_allowed_ips_and_secret(
      69              :         &self,
      70              :     ) -> Result<(CachedAllowedIps, Option<CachedRoleSecret>), console::errors::GetAuthInfoError>;
      71              :     fn get_role_secret(&self) -> Result<CachedRoleSecret, console::errors::GetAuthInfoError>;
      72              : }
      73              : 
      74              : impl std::fmt::Display for BackendType<'_, ()> {
      75           25 :     fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
      76           25 :         use BackendType::*;
      77           25 :         match self {
      78           22 :             Console(api, _) => match &**api {
      79            1 :                 ConsoleBackend::Console(endpoint) => {
      80            1 :                     fmt.debug_tuple("Console").field(&endpoint.url()).finish()
      81              :                 }
      82              :                 #[cfg(any(test, feature = "testing"))]
      83           21 :                 ConsoleBackend::Postgres(endpoint) => {
      84           21 :                     fmt.debug_tuple("Postgres").field(&endpoint.url()).finish()
      85              :                 }
      86              :                 #[cfg(test)]
      87            0 :                 ConsoleBackend::Test(_) => fmt.debug_tuple("Test").finish(),
      88              :             },
      89            3 :             Link(url) => fmt.debug_tuple("Link").field(&url.as_str()).finish(),
      90              :         }
      91           25 :     }
      92              : }
      93              : 
      94              : impl<T> BackendType<'_, T> {
      95              :     /// Very similar to [`std::option::Option::as_ref`].
      96              :     /// This helps us pass structured config to async tasks.
      97          139 :     pub fn as_ref(&self) -> BackendType<'_, &T> {
      98          139 :         use BackendType::*;
      99          139 :         match self {
     100          136 :             Console(c, x) => Console(MaybeOwned::Borrowed(c), x),
     101            3 :             Link(c) => Link(MaybeOwned::Borrowed(c)),
     102              :         }
     103          139 :     }
     104              : }
     105              : 
     106              : impl<'a, T> BackendType<'a, T> {
     107              :     /// Very similar to [`std::option::Option::map`].
     108              :     /// Maps [`BackendType<T>`] to [`BackendType<R>`] by applying
     109              :     /// a function to a contained value.
     110          139 :     pub fn map<R>(self, f: impl FnOnce(T) -> R) -> BackendType<'a, R> {
     111          139 :         use BackendType::*;
     112          139 :         match self {
     113          136 :             Console(c, x) => Console(c, f(x)),
     114            3 :             Link(c) => Link(c),
     115              :         }
     116          139 :     }
     117              : }
     118              : 
     119              : impl<'a, T, E> BackendType<'a, Result<T, E>> {
     120              :     /// Very similar to [`std::option::Option::transpose`].
     121              :     /// This is most useful for error handling.
     122           52 :     pub fn transpose(self) -> Result<BackendType<'a, T>, E> {
     123           52 :         use BackendType::*;
     124           52 :         match self {
     125           49 :             Console(c, x) => x.map(|x| Console(c, x)),
     126            3 :             Link(c) => Ok(Link(c)),
     127              :         }
     128           52 :     }
     129              : }
     130              : 
     131              : pub struct ComputeCredentials<T> {
     132              :     pub info: ComputeUserInfo,
     133              :     pub keys: T,
     134              : }
     135              : 
     136            0 : #[derive(Debug, Clone)]
     137              : pub struct ComputeUserInfoNoEndpoint {
     138              :     pub user: RoleName,
     139              :     pub options: NeonOptions,
     140              : }
     141              : 
     142          248 : #[derive(Debug, Clone)]
     143              : pub struct ComputeUserInfo {
     144              :     pub endpoint: EndpointId,
     145              :     pub user: RoleName,
     146              :     pub options: NeonOptions,
     147              : }
     148              : 
     149              : impl ComputeUserInfo {
     150           68 :     pub fn endpoint_cache_key(&self) -> EndpointCacheKey {
     151           68 :         self.options.get_cache_key(&self.endpoint)
     152           68 :     }
     153              : }
     154              : 
     155              : pub enum ComputeCredentialKeys {
     156              :     #[cfg(any(test, feature = "testing"))]
     157              :     Password(Vec<u8>),
     158              :     AuthKeys(AuthKeys),
     159              : }
     160              : 
     161              : impl TryFrom<ComputeUserInfoMaybeEndpoint> for ComputeUserInfo {
     162              :     // user name
     163              :     type Error = ComputeUserInfoNoEndpoint;
     164              : 
     165           49 :     fn try_from(user_info: ComputeUserInfoMaybeEndpoint) -> Result<Self, Self::Error> {
     166           49 :         match user_info.endpoint_id {
     167            3 :             None => Err(ComputeUserInfoNoEndpoint {
     168            3 :                 user: user_info.user,
     169            3 :                 options: user_info.options,
     170            3 :             }),
     171           46 :             Some(endpoint) => Ok(ComputeUserInfo {
     172           46 :                 endpoint,
     173           46 :                 user: user_info.user,
     174           46 :                 options: user_info.options,
     175           46 :             }),
     176              :         }
     177           49 :     }
     178              : }
     179              : 
     180              : /// True to its name, this function encapsulates our current auth trade-offs.
     181              : /// Here, we choose the appropriate auth flow based on circumstances.
     182              : ///
     183              : /// All authentication flows will emit an AuthenticationOk message if successful.
     184           49 : async fn auth_quirks(
     185           49 :     ctx: &mut RequestMonitoring,
     186           49 :     api: &impl console::Api,
     187           49 :     user_info: ComputeUserInfoMaybeEndpoint,
     188           49 :     client: &mut stream::PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
     189           49 :     allow_cleartext: bool,
     190           49 :     config: &'static AuthenticationConfig,
     191           49 : ) -> auth::Result<ComputeCredentials<ComputeCredentialKeys>> {
     192              :     // If there's no project so far, that entails that client doesn't
     193              :     // support SNI or other means of passing the endpoint (project) name.
     194              :     // We now expect to see a very specific payload in the place of password.
     195           49 :     let (info, unauthenticated_password) = match user_info.try_into() {
     196            3 :         Err(info) => {
     197            3 :             let res = hacks::password_hack_no_authentication(ctx, info, client).await?;
     198              : 
     199            2 :             ctx.set_endpoint_id(res.info.endpoint.clone());
     200            2 :             tracing::Span::current().record("ep", &tracing::field::display(&res.info.endpoint));
     201            2 : 
     202            2 :             (res.info, Some(res.keys))
     203              :         }
     204           46 :         Ok(info) => (info, None),
     205              :     };
     206              : 
     207           48 :     info!("fetching user's authentication info");
     208          297 :     let (allowed_ips, maybe_secret) = api.get_allowed_ips_and_secret(ctx, &info).await?;
     209              : 
     210              :     // check allowed list
     211           44 :     if !check_peer_addr_is_in_list(&ctx.peer_addr, &allowed_ips) {
     212            3 :         return Err(auth::AuthError::ip_address_not_allowed());
     213           41 :     }
     214           41 :     let cached_secret = match maybe_secret {
     215            0 :         Some(secret) => secret,
     216          285 :         None => api.get_role_secret(ctx, &info).await?,
     217              :     };
     218              : 
     219           41 :     let secret = cached_secret.value.clone().unwrap_or_else(|| {
     220            1 :         // If we don't have an authentication secret, we mock one to
     221            1 :         // prevent malicious probing (possible due to missing protocol steps).
     222            1 :         // This mocked secret will never lead to successful authentication.
     223            1 :         info!("authentication info not found, mocking it");
     224            1 :         AuthSecret::Scram(scram::ServerSecret::mock(&info.user, rand::random()))
     225           41 :     });
     226           41 :     match authenticate_with_secret(
     227           41 :         ctx,
     228           41 :         secret,
     229           41 :         info,
     230           41 :         client,
     231           41 :         unauthenticated_password,
     232           41 :         allow_cleartext,
     233           41 :         config,
     234           41 :     )
     235           78 :     .await
     236              :     {
     237           38 :         Ok(keys) => Ok(keys),
     238            3 :         Err(e) => {
     239            3 :             if e.is_auth_failed() {
     240            3 :                 // The password could have been changed, so we invalidate the cache.
     241            3 :                 cached_secret.invalidate();
     242            3 :             }
     243            3 :             Err(e)
     244              :         }
     245              :     }
     246           49 : }
     247              : 
     248           41 : async fn authenticate_with_secret(
     249           41 :     ctx: &mut RequestMonitoring,
     250           41 :     secret: AuthSecret,
     251           41 :     info: ComputeUserInfo,
     252           41 :     client: &mut stream::PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
     253           41 :     unauthenticated_password: Option<Vec<u8>>,
     254           41 :     allow_cleartext: bool,
     255           41 :     config: &'static AuthenticationConfig,
     256           41 : ) -> auth::Result<ComputeCredentials<ComputeCredentialKeys>> {
     257           41 :     if let Some(password) = unauthenticated_password {
     258            2 :         let auth_outcome = validate_password_and_exchange(&password, secret)?;
     259            2 :         let keys = match auth_outcome {
     260            2 :             crate::sasl::Outcome::Success(key) => key,
     261            0 :             crate::sasl::Outcome::Failure(reason) => {
     262            0 :                 info!("auth backend failed with an error: {reason}");
     263            0 :                 return Err(auth::AuthError::auth_failed(&*info.user));
     264              :             }
     265              :         };
     266              : 
     267              :         // we have authenticated the password
     268            2 :         client.write_message_noflush(&pq_proto::BeMessage::AuthenticationOk)?;
     269              : 
     270            2 :         return Ok(ComputeCredentials { info, keys });
     271           39 :     }
     272           39 : 
     273           39 :     // -- the remaining flows are self-authenticating --
     274           39 : 
     275           39 :     // Perform cleartext auth if we're allowed to do that.
     276           39 :     // Currently, we use it for websocket connections (latency).
     277           39 :     if allow_cleartext {
     278            0 :         ctx.set_auth_method(crate::context::AuthMethod::Cleartext);
     279            0 :         return hacks::authenticate_cleartext(ctx, info, client, secret).await;
     280           39 :     }
     281           39 : 
     282           39 :     // Finally, proceed with the main auth flow (SCRAM-based).
     283           78 :     classic::authenticate(ctx, info, client, config, secret).await
     284           41 : }
     285              : 
     286              : impl<'a> BackendType<'a, ComputeUserInfoMaybeEndpoint> {
     287              :     /// Get compute endpoint name from the credentials.
     288           52 :     pub fn get_endpoint(&self) -> Option<EndpointId> {
     289           52 :         use BackendType::*;
     290           52 : 
     291           52 :         match self {
     292           49 :             Console(_, user_info) => user_info.endpoint_id.clone(),
     293            3 :             Link(_) => Some("link".into()),
     294              :         }
     295           52 :     }
     296              : 
     297              :     /// Get username from the credentials.
     298           52 :     pub fn get_user(&self) -> &str {
     299           52 :         use BackendType::*;
     300           52 : 
     301           52 :         match self {
     302           49 :             Console(_, user_info) => &user_info.user,
     303            3 :             Link(_) => "link",
     304              :         }
     305           52 :     }
     306              : 
     307              :     /// Authenticate the client via the requested backend, possibly using credentials.
     308          104 :     #[tracing::instrument(fields(allow_cleartext = allow_cleartext), skip_all)]
     309              :     pub async fn authenticate(
     310              :         self,
     311              :         ctx: &mut RequestMonitoring,
     312              :         client: &mut stream::PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
     313              :         allow_cleartext: bool,
     314              :         config: &'static AuthenticationConfig,
     315              :     ) -> auth::Result<(CachedNodeInfo, BackendType<'a, ComputeUserInfo>)> {
     316              :         use BackendType::*;
     317              : 
     318              :         let res = match self {
     319              :             Console(api, user_info) => {
     320           49 :                 info!(
     321           49 :                     user = &*user_info.user,
     322           49 :                     project = user_info.endpoint(),
     323           49 :                     "performing authentication using the console"
     324           49 :                 );
     325              : 
     326              :                 let compute_credentials =
     327              :                     auth_quirks(ctx, &*api, user_info, client, allow_cleartext, config).await?;
     328              : 
     329              :                 let mut num_retries = 0;
     330              :                 let mut node =
     331              :                     wake_compute(&mut num_retries, ctx, &api, &compute_credentials.info).await?;
     332              : 
     333              :                 ctx.set_project(node.aux.clone());
     334              : 
     335              :                 match compute_credentials.keys {
     336              :                     #[cfg(any(test, feature = "testing"))]
     337              :                     ComputeCredentialKeys::Password(password) => node.config.password(password),
     338              :                     ComputeCredentialKeys::AuthKeys(auth_keys) => node.config.auth_keys(auth_keys),
     339              :                 };
     340              : 
     341              :                 (node, BackendType::Console(api, compute_credentials.info))
     342              :             }
     343              :             // NOTE: this auth backend doesn't use client credentials.
     344              :             Link(url) => {
     345            3 :                 info!("performing link authentication");
     346              : 
     347              :                 let node_info = link::authenticate(ctx, &url, client).await?;
     348              : 
     349              :                 (
     350              :                     CachedNodeInfo::new_uncached(node_info),
     351              :                     BackendType::Link(url),
     352              :                 )
     353              :             }
     354              :         };
     355              : 
     356           41 :         info!("user successfully authenticated");
     357              :         Ok(res)
     358              :     }
     359              : }
     360              : 
     361              : impl BackendType<'_, ComputeUserInfo> {
     362           46 :     pub async fn get_role_secret(
     363           46 :         &self,
     364           46 :         ctx: &mut RequestMonitoring,
     365           46 :     ) -> Result<CachedRoleSecret, GetAuthInfoError> {
     366           46 :         use BackendType::*;
     367           46 :         match self {
     368          317 :             Console(api, user_info) => api.get_role_secret(ctx, user_info).await,
     369            0 :             Link(_) => Ok(Cached::new_uncached(None)),
     370              :         }
     371           46 :     }
     372              : 
     373           47 :     pub async fn get_allowed_ips_and_secret(
     374           47 :         &self,
     375           47 :         ctx: &mut RequestMonitoring,
     376           47 :     ) -> Result<(CachedAllowedIps, Option<CachedRoleSecret>), GetAuthInfoError> {
     377           47 :         use BackendType::*;
     378           47 :         match self {
     379          366 :             Console(api, user_info) => api.get_allowed_ips_and_secret(ctx, user_info).await,
     380            0 :             Link(_) => Ok((Cached::new_uncached(Arc::new(vec![])), None)),
     381              :         }
     382           47 :     }
     383              : 
     384              :     /// When applicable, wake the compute node, gaining its connection info in the process.
     385              :     /// The link auth flow doesn't support this, so we return [`None`] in that case.
     386           40 :     pub async fn wake_compute(
     387           40 :         &self,
     388           40 :         ctx: &mut RequestMonitoring,
     389           40 :     ) -> Result<Option<CachedNodeInfo>, console::errors::WakeComputeError> {
     390           40 :         use BackendType::*;
     391           40 : 
     392           40 :         match self {
     393           40 :             Console(api, user_info) => api.wake_compute(ctx, user_info).map_ok(Some).await,
     394            0 :             Link(_) => Ok(None),
     395              :         }
     396           40 :     }
     397              : }
        

Generated by: LCOV version 2.1-beta