LCOV - differential code coverage report
Current view: top level - proxy/src/auth - backend.rs (source / functions) Coverage Total Hit UBC CBC
Current: f6946e90941b557c917ac98cd5a7e9506d180f3e.info Lines: 79.0 % 105 83 22 83
Current Date: 2023-10-19 02:04:12 Functions: 38.0 % 50 19 31 19
Baseline: c8637f37369098875162f194f92736355783b050.info
Baseline Date: 2023-10-18 20:25:20

           TLA  Line data    Source code
       1                 : mod classic;
       2                 : mod hacks;
       3                 : mod link;
       4                 : 
       5                 : pub use link::LinkAuthError;
       6                 : 
       7                 : use crate::{
       8                 :     auth::{self, ClientCredentials},
       9                 :     console::{
      10                 :         self,
      11                 :         provider::{CachedNodeInfo, ConsoleReqExtra},
      12                 :         Api,
      13                 :     },
      14                 :     stream, url,
      15                 : };
      16                 : use futures::TryFutureExt;
      17                 : use std::borrow::Cow;
      18                 : use tokio::io::{AsyncRead, AsyncWrite};
      19                 : use tracing::info;
      20                 : 
      21                 : /// A product of successful authentication.
      22                 : pub struct AuthSuccess<T> {
      23                 :     /// Did we send [`pq_proto::BeMessage::AuthenticationOk`] to client?
      24                 :     pub reported_auth_ok: bool,
      25                 :     /// Something to be considered a positive result.
      26                 :     pub value: T,
      27                 : }
      28                 : 
      29                 : impl<T> AuthSuccess<T> {
      30                 :     /// Very similar to [`std::option::Option::map`].
      31                 :     /// Maps [`AuthSuccess<T>`] to [`AuthSuccess<R>`] by applying
      32                 :     /// a function to a contained value.
      33 CBC           3 :     pub fn map<R>(self, f: impl FnOnce(T) -> R) -> AuthSuccess<R> {
      34               3 :         AuthSuccess {
      35               3 :             reported_auth_ok: self.reported_auth_ok,
      36               3 :             value: f(self.value),
      37               3 :         }
      38               3 :     }
      39                 : }
      40                 : 
      41                 : /// This type serves two purposes:
      42                 : ///
      43                 : /// * When `T` is `()`, it's just a regular auth backend selector
      44                 : ///   which we use in [`crate::config::ProxyConfig`].
      45                 : ///
      46                 : /// * However, when we substitute `T` with [`ClientCredentials`],
      47                 : ///   this helps us provide the credentials only to those auth
      48                 : ///   backends which require them for the authentication process.
      49                 : pub enum BackendType<'a, T> {
      50                 :     /// Current Cloud API (V2).
      51                 :     Console(Cow<'a, console::provider::neon::Api>, T),
      52                 :     /// Local mock of Cloud API (V2).
      53                 :     Postgres(Cow<'a, console::provider::mock::Api>, T),
      54                 :     /// Authentication via a web browser.
      55                 :     Link(Cow<'a, url::ApiUrl>),
      56                 :     /// Test backend.
      57                 :     Test(&'a dyn TestBackend),
      58                 : }
      59                 : 
      60                 : pub trait TestBackend: Send + Sync + 'static {
      61                 :     fn wake_compute(&self) -> Result<CachedNodeInfo, console::errors::WakeComputeError>;
      62                 : }
      63                 : 
      64                 : impl std::fmt::Display for BackendType<'_, ()> {
      65              16 :     fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
      66              16 :         use BackendType::*;
      67              16 :         match self {
      68 UBC           0 :             Console(endpoint, _) => fmt.debug_tuple("Console").field(&endpoint.url()).finish(),
      69 CBC          13 :             Postgres(endpoint, _) => fmt.debug_tuple("Postgres").field(&endpoint.url()).finish(),
      70               3 :             Link(url) => fmt.debug_tuple("Link").field(&url.as_str()).finish(),
      71 UBC           0 :             Test(_) => fmt.debug_tuple("Test").finish(),
      72                 :         }
      73 CBC          16 :     }
      74                 : }
      75                 : 
      76                 : impl<T> BackendType<'_, T> {
      77                 :     /// Very similar to [`std::option::Option::as_ref`].
      78                 :     /// This helps us pass structured config to async tasks.
      79              60 :     pub fn as_ref(&self) -> BackendType<'_, &T> {
      80              60 :         use BackendType::*;
      81              60 :         match self {
      82 UBC           0 :             Console(c, x) => Console(Cow::Borrowed(c), x),
      83 CBC          57 :             Postgres(c, x) => Postgres(Cow::Borrowed(c), x),
      84               3 :             Link(c) => Link(Cow::Borrowed(c)),
      85 UBC           0 :             Test(x) => Test(*x),
      86                 :         }
      87 CBC          60 :     }
      88                 : }
      89                 : 
      90                 : impl<'a, T> BackendType<'a, T> {
      91                 :     /// Very similar to [`std::option::Option::map`].
      92                 :     /// Maps [`BackendType<T>`] to [`BackendType<R>`] by applying
      93                 :     /// a function to a contained value.
      94              60 :     pub fn map<R>(self, f: impl FnOnce(T) -> R) -> BackendType<'a, R> {
      95              60 :         use BackendType::*;
      96              60 :         match self {
      97 UBC           0 :             Console(c, x) => Console(c, f(x)),
      98 CBC          57 :             Postgres(c, x) => Postgres(c, f(x)),
      99               3 :             Link(c) => Link(c),
     100 UBC           0 :             Test(x) => Test(x),
     101                 :         }
     102 CBC          60 :     }
     103                 : }
     104                 : 
     105                 : impl<'a, T, E> BackendType<'a, Result<T, E>> {
     106                 :     /// Very similar to [`std::option::Option::transpose`].
     107                 :     /// This is most useful for error handling.
     108              60 :     pub fn transpose(self) -> Result<BackendType<'a, T>, E> {
     109              60 :         use BackendType::*;
     110              60 :         match self {
     111 UBC           0 :             Console(c, x) => x.map(|x| Console(c, x)),
     112 CBC          57 :             Postgres(c, x) => x.map(|x| Postgres(c, x)),
     113               3 :             Link(c) => Ok(Link(c)),
     114 UBC           0 :             Test(x) => Ok(Test(x)),
     115                 :         }
     116 CBC          60 :     }
     117                 : }
     118                 : 
     119                 : /// True to its name, this function encapsulates our current auth trade-offs.
     120                 : /// Here, we choose the appropriate auth flow based on circumstances.
     121              30 : async fn auth_quirks(
     122              30 :     api: &impl console::Api,
     123              30 :     extra: &ConsoleReqExtra<'_>,
     124              30 :     creds: &mut ClientCredentials<'_>,
     125              30 :     client: &mut stream::PqStream<impl AsyncRead + AsyncWrite + Unpin>,
     126              30 :     allow_cleartext: bool,
     127              30 : ) -> auth::Result<AuthSuccess<CachedNodeInfo>> {
     128              30 :     // If there's no project so far, that entails that client doesn't
     129              30 :     // support SNI or other means of passing the endpoint (project) name.
     130              30 :     // We now expect to see a very specific payload in the place of password.
     131              30 :     if creds.project.is_none() {
     132                 :         // Password will be checked by the compute node later.
     133               3 :         return hacks::password_hack(api, extra, creds, client).await;
     134              27 :     }
     135              27 : 
     136              27 :     // Password hack should set the project name.
     137              27 :     // TODO: make `creds.project` more type-safe.
     138              27 :     assert!(creds.project.is_some());
     139                 : 
     140                 :     // Perform cleartext auth if we're allowed to do that.
     141                 :     // Currently, we use it for websocket connections (latency).
     142              27 :     if allow_cleartext {
     143                 :         // Password will be checked by the compute node later.
     144 UBC           0 :         return hacks::cleartext_hack(api, extra, creds, client).await;
     145 CBC          27 :     }
     146              27 : 
     147              27 :     // Finally, proceed with the main auth flow (SCRAM-based).
     148             189 :     classic::authenticate(api, extra, creds, client).await
     149              30 : }
     150                 : 
     151                 : impl BackendType<'_, ClientCredentials<'_>> {
     152                 :     /// Get compute endpoint name from the credentials.
     153              33 :     pub fn get_endpoint(&self) -> Option<String> {
     154              33 :         use BackendType::*;
     155              33 : 
     156              33 :         match self {
     157 UBC           0 :             Console(_, creds) => creds.project.clone(),
     158 CBC          30 :             Postgres(_, creds) => creds.project.clone(),
     159               3 :             Link(_) => Some("link".to_owned()),
     160 UBC           0 :             Test(_) => Some("test".to_owned()),
     161                 :         }
     162 CBC          33 :     }
     163                 : 
     164                 :     /// Get username from the credentials.
     165               4 :     pub fn get_user(&self) -> &str {
     166               4 :         use BackendType::*;
     167               4 : 
     168               4 :         match self {
     169 UBC           0 :             Console(_, creds) => creds.user,
     170 CBC           4 :             Postgres(_, creds) => creds.user,
     171 UBC           0 :             Link(_) => "link",
     172               0 :             Test(_) => "test",
     173                 :         }
     174 CBC           4 :     }
     175                 : 
     176                 :     /// Authenticate the client via the requested backend, possibly using credentials.
     177              99 :     #[tracing::instrument(fields(allow_cleartext = allow_cleartext), skip_all)]
     178                 :     pub async fn authenticate(
     179                 :         &mut self,
     180                 :         extra: &ConsoleReqExtra<'_>,
     181                 :         client: &mut stream::PqStream<impl AsyncRead + AsyncWrite + Unpin>,
     182                 :         allow_cleartext: bool,
     183                 :     ) -> auth::Result<AuthSuccess<CachedNodeInfo>> {
     184                 :         use BackendType::*;
     185                 : 
     186                 :         let res = match self {
     187                 :             Console(api, creds) => {
     188 UBC           0 :                 info!(
     189               0 :                     user = creds.user,
     190               0 :                     project = creds.project(),
     191               0 :                     "performing authentication using the console"
     192               0 :                 );
     193                 : 
     194                 :                 let api = api.as_ref();
     195                 :                 auth_quirks(api, extra, creds, client, allow_cleartext).await?
     196                 :             }
     197                 :             Postgres(api, creds) => {
     198 CBC          30 :                 info!(
     199              30 :                     user = creds.user,
     200              30 :                     project = creds.project(),
     201              30 :                     "performing authentication using a local postgres instance"
     202              30 :                 );
     203                 : 
     204                 :                 let api = api.as_ref();
     205                 :                 auth_quirks(api, extra, creds, client, allow_cleartext).await?
     206                 :             }
     207                 :             // NOTE: this auth backend doesn't use client credentials.
     208                 :             Link(url) => {
     209               3 :                 info!("performing link authentication");
     210                 : 
     211                 :                 link::authenticate(url, client)
     212                 :                     .await?
     213                 :                     .map(CachedNodeInfo::new_uncached)
     214                 :             }
     215                 :             Test(_) => {
     216                 :                 unreachable!("this function should never be called in the test backend")
     217                 :             }
     218                 :         };
     219                 : 
     220              29 :         info!("user successfully authenticated");
     221                 :         Ok(res)
     222                 :     }
     223                 : 
     224                 :     /// When applicable, wake the compute node, gaining its connection info in the process.
     225                 :     /// The link auth flow doesn't support this, so we return [`None`] in that case.
     226              27 :     pub async fn wake_compute(
     227              27 :         &self,
     228              27 :         extra: &ConsoleReqExtra<'_>,
     229              27 :     ) -> Result<Option<CachedNodeInfo>, console::errors::WakeComputeError> {
     230              27 :         use BackendType::*;
     231              27 : 
     232              27 :         match self {
     233 UBC           0 :             Console(api, creds) => api.wake_compute(extra, creds).map_ok(Some).await,
     234 CBC          27 :             Postgres(api, creds) => api.wake_compute(extra, creds).map_ok(Some).await,
     235 UBC           0 :             Link(_) => Ok(None),
     236               0 :             Test(x) => x.wake_compute().map(Some),
     237                 :         }
     238 CBC          27 :     }
     239                 : }
        

Generated by: LCOV version 2.1-beta