dsql.rs

//! Aurora DSQL requires an IAM token in place of a password. Tokens are
//! generated by the AWS SDK using your AWS credentials.

use std::{
    borrow::Cow,
    fmt,
    sync::{Arc, RwLock},
    time::Duration,
};

use aws_config::{BehaviorVersion, SdkConfig};
use aws_sdk_dsql::{
    auth_token::{AuthToken, AuthTokenGenerator, Config},
    error::BoxError,
};
use sqlx_postgres::PasswordProvider;
use tokio::{task::JoinHandle, time::sleep};

/// A builder type to get you build a customized [`DsqlIamProvider`], in case
/// the AWS SDK defaults aren't what you're looking for.
///
/// If you're happy with the AWS SDK defaults, prefer using
/// [`DsqlIamProvider::new`].
///
///
/// ```ignore
/// use sqlx_aws::iam::dsql::*;
///
/// let b = DsqlIamProviderBuilder::defaults().await;
/// let my_config = Config::builder().hostname("...").build()?;
/// let provider = b.with_generator_config(my_config).await?;
/// ```
pub struct DsqlIamProviderBuilder {
    cfg: SdkConfig,
    is_admin: bool,
}

impl DsqlIamProviderBuilder {
    /// A new builder. The AWS SDK is automatically configured.
    pub async fn defaults() -> Self {
        let cfg = aws_config::load_defaults(BehaviorVersion::latest()).await;
        Self::new_with_sdk_cfg(cfg)
    }

    /// A new builder with custom SDK config.
    pub fn new_with_sdk_cfg(cfg: SdkConfig) -> Self {
        Self {
            cfg,
            is_admin: false,
        }
    }

    /// Build a provider with the given [`auth_token::Config`].
    pub async fn with_generator_config(self, config: Config) -> Result<DsqlIamProvider, BoxError> {
        let DsqlIamProviderBuilder { cfg, is_admin } = self;

        // This default value is hardcoded in the AuthTokenGenerator. There is
        // no way to share the value.
        let expires_in = config.expires_in().unwrap_or(900);

        // Token generation is fast (because it is a local operation). However,
        // there is some coordination involved (such as loading AWS credentials,
        // or tokio scheduling). We want to avoid ever having stale tokens, and so schedule refreshes slightly ahead of expiry.
        let refresh_interval = Duration::from_secs(if expires_in > 60 {
            expires_in - 60
        } else {
            expires_in
        });

        let generator = AuthTokenGenerator::new(config);

        // Boostrap: try once. This allows for failing fast for the case where
        // things haven't been correctly configured.
        let auth_token = match is_admin {
            true => generator.db_connect_admin_auth_token(&cfg).await,
            false => generator.db_connect_auth_token(&cfg).await,
        }?;

        let token = Arc::new(RwLock::new(Ok(auth_token)));
        let _token = token.clone();

        let task = tokio::spawn(async move {
            sleep(refresh_interval).await;

            loop {
                let res = match is_admin {
                    true => generator.db_connect_admin_auth_token(&cfg).await,
                    false => generator.db_connect_auth_token(&cfg).await,
                };
                match res {
                    Ok(auth_token) => {
                        *_token.write().expect("never poisoned") = Ok(auth_token);
                        sleep(refresh_interval).await;
                    }
                    // XXX: In theory, this should almost never happen, because
                    // we did a boostrap token generation, which should catch
                    // nearly all errors. However, it is possible that the
                    // underlying credential provider has failed in some way.
                    Err(err) => {
                        // Refreshes are eager, which means it may be possible
                        // that we're about to replace perfectly good token with
                        // an error. It doesn't seem worthwhile to guard against
                        // that, since tokens are short lived and are likely to
                        // expire shortly anyways.
                        *_token.write().expect("never poisoned") = Err(err);

                        // sleep an arbitrary amount of time to prevent busy
                        // loops, but not so long that we don't try again (if
                        // the underlying error has been resolved).
                        sleep(Duration::from_secs(1)).await;
                    }
                }
            }
        });

        Ok(DsqlIamProvider { token, task })
    }
}

/// A sqlx [`PasswordProvider`] that automatically manages IAM tokens.
///
/// ```ignore
/// use sqlx_postgres::PgConnectOptions;
/// use sqlx_aws::iam::dsql::*;
///
/// let provider = DsqlIamProvider::new("peccy.dsql.us-east-1.on.aws").await?;
/// let opts = PgConnectOptions::new_without_pgpass()
///     .password(provider);
/// ```
pub struct DsqlIamProvider {
    token: Arc<RwLock<Result<AuthToken, BoxError>>>,
    task: JoinHandle<()>,
}

impl Drop for DsqlIamProvider {
    fn drop(&mut self) {
        self.task.abort();
    }
}

impl DsqlIamProvider {
    pub async fn new(hostname: impl Into<String>) -> Result<Self, BoxError> {
        let builder = DsqlIamProviderBuilder::defaults().await;
        let config = Config::builder()
            .hostname(hostname)
            .build()
            .expect("hostname was provided");
        builder.with_generator_config(config).await
    }
}

impl PasswordProvider for DsqlIamProvider {
    fn password<'a>(&'a self) -> Result<Cow<'a, str>, sqlx_core::error::BoxDynError> {
        match &*self.token.read().expect("never poisoned") {
            Ok(auth_token) => Ok(Cow::Owned(auth_token.as_str().to_string())),
            Err(err) => Err(Box::new(RefreshError(format!("{err}")))),
        }
    }
}

#[derive(Debug)]
pub struct RefreshError(String);

impl fmt::Display for RefreshError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "unable to refresh auth token: {}", self.0)
    }
}

impl std::error::Error for RefreshError {}