Line data    Source code

       1              : use tokio::io::{AsyncRead, AsyncWrite};
       2              : use tracing::{info, warn};
       3              : 
       4              : use super::{ComputeCredentials, ComputeUserInfo};
       5              : use crate::auth::backend::ComputeCredentialKeys;
       6              : use crate::auth::{self, AuthFlow};
       7              : use crate::config::AuthenticationConfig;
       8              : use crate::context::RequestMonitoring;
       9              : use crate::control_plane::AuthSecret;
      10              : use crate::stream::{PqStream, Stream};
      11              : use crate::{compute, sasl};
      12              : 
      13            1 : pub(super) async fn authenticate(
      14            1 :     ctx: &RequestMonitoring,
      15            1 :     creds: ComputeUserInfo,
      16            1 :     client: &mut PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
      17            1 :     config: &'static AuthenticationConfig,
      18            1 :     secret: AuthSecret,
      19            1 : ) -> auth::Result<ComputeCredentials> {
      20            1 :     let flow = AuthFlow::new(client);
      21            1 :     let scram_keys = match secret {
      22              :         #[cfg(any(test, feature = "testing"))]
      23              :         AuthSecret::Md5(_) => {
      24            0 :             info!("auth endpoint chooses MD5");
      25            0 :             return Err(auth::AuthError::bad_auth_method("MD5"));
      26              :         }
      27            1 :         AuthSecret::Scram(secret) => {
      28            1 :             info!("auth endpoint chooses SCRAM");
      29            1 :             let scram = auth::Scram(&secret, ctx);
      30              : 
      31            1 :             let auth_outcome = tokio::time::timeout(
      32            1 :                 config.scram_protocol_timeout,
      33            1 :                 async {
      34            1 : 
      35            1 :                     flow.begin(scram).await.map_err(|error| {
      36            0 :                         warn!(?error, "error sending scram acknowledgement");
      37            0 :                         error
      38            2 :                     })?.authenticate().await.map_err(|error| {
      39            0 :                         warn!(?error, "error processing scram messages");
      40            0 :                         error
      41            1 :                     })
      42            1 :                 }
      43            1 :             )
      44            2 :             .await
      45            1 :             .map_err(|e| {
      46            0 :                 warn!("error processing scram messages error = authentication timed out, execution time exceeded {} seconds", config.scram_protocol_timeout.as_secs());
      47            0 :                 auth::AuthError::user_timeout(e)
      48            1 :             })??;
      49              : 
      50            1 :             let client_key = match auth_outcome {
      51            1 :                 sasl::Outcome::Success(key) => key,
      52            0 :                 sasl::Outcome::Failure(reason) => {
      53            0 :                     info!("auth backend failed with an error: {reason}");
      54            0 :                     return Err(auth::AuthError::auth_failed(&*creds.user));
      55              :                 }
      56              :             };
      57              : 
      58            1 :             compute::ScramKeys {
      59            1 :                 client_key: client_key.as_bytes(),
      60            1 :                 server_key: secret.server_key.as_bytes(),
      61            1 :             }
      62            1 :         }
      63            1 :     };
      64            1 : 
      65            1 :     Ok(ComputeCredentials {
      66            1 :         info: creds,
      67            1 :         keys: ComputeCredentialKeys::AuthKeys(tokio_postgres::config::AuthKeys::ScramSha256(
      68            1 :             scram_keys,
      69            1 :         )),
      70            1 :     })
      71            1 : }