sqlx-postgres, sqlx-aws: impl dynamic passwords

Cargo.toml

@@ -9,6 +9,7 @@ members = [
     #    "sqlx-bench",
     "sqlx-mysql",
     "sqlx-postgres",
+    "sqlx-aws",
     "sqlx-sqlite",
     "examples/mysql/todos",
     "examples/postgres/axum-social-with-tests",

sqlx-aws/Cargo.toml

@@ -0,0 +1,36 @@
+[package]
+name = "sqlx-aws"
+documentation = "https://docs.rs/sqlx"
+description = "AWS specific helpers for using sqlx."
+
+version.workspace = true
+license.workspace = true
+edition.workspace = true
+repository.workspace = true
+keywords.workspace = true
+categories.workspace = true
+authors.workspace = true
+
+[features]
+default = ["dsql"]
+dsql = ["sqlx-postgres", "aws-sdk-dsql"]
+
+[dependencies]
+aws-config = { version = "1.6.2", default-features = false, features = ["behavior-version-latest"] }
+aws-sdk-dsql = { version = "1.16.0", default-features = false, optional = true }
+sqlx-core = { workspace = true }
+sqlx-postgres = { workspace = true, optional = true }
+tokio = { workspace = true, features = ["rt", "sync"] }
+
+[dev-dependencies]
+sqlx = { workspace = true, features = ["runtime-tokio"] }
+tokio = { workspace = true, features = ["rt", "rt-multi-thread", "sync", "macros"] }
+aws-config = { version = "1.6.2", features = ["behavior-version-latest"] }
+aws-sdk-dsql = { version = "1.16.0" }
+
+[lints]
+workspace = true
+
+[[example]]
+name = "dsql"
+features = ["dsql"]

sqlx-aws/examples/dsql.rs

@@ -0,0 +1,17 @@
+use aws_sdk_dsql::error::BoxError;
+use sqlx_aws::iam::dsql::DsqlIamProvider;
+use sqlx_postgres::{PgConnectOptions, PgPoolOptions};
+
+#[tokio::main]
+async fn main() -> Result<(), BoxError> {
+    let hostname = std::env::var("DSQL_CLUSTER_ENDPOINT")
+        .expect("please set DSQL_CLUSTER_ENDPOINT is your environment");
+
+    let provider = DsqlIamProvider::new(hostname).await?;
+    let opts = PgConnectOptions::new_without_pgpass()
+        .password(provider)
+        .database("postgres");
+    let _pool = PgPoolOptions::new().connect_with(opts).await?;
+
+    Ok(())
+}

sqlx-aws/src/iam/dsql.rs

@@ -0,0 +1,170 @@
+//! 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 {}

sqlx-aws/src/iam/mod.rs

@@ -0,0 +1,2 @@
+#[cfg(feature = "dsql")]
+pub mod dsql;

sqlx-aws/src/lib.rs

@@ -0,0 +1 @@
+pub mod iam;

sqlx-core/src/error.rs

@@ -52,6 +52,10 @@ pub enum Error {
     #[error("error occurred while attempting to establish a TLS connection: {0}")]
     Tls(#[source] BoxDynError),
 
+    /// Password provider failed.
+    #[error("unable to provide a password: {0}")]
+    PasswordProvider(#[source] BoxDynError),
+
     /// Unexpected or invalid data encountered while communicating with the database.
     ///
     /// This should indicate there is a programming error in a SQLx driver or there

sqlx-postgres/src/connection/establish.rs

@@ -76,9 +76,7 @@ impl PgConnection {
                         // password in clear-text form.
 
                         stream
-                            .send(Password::Cleartext(
-                                options.password.as_deref().unwrap_or_default(),
-                            ))
+                            .send(Password::Cleartext(options.password.get()?.as_ref()))
                             .await?;
                     }
 
@@ -91,7 +89,7 @@ impl PgConnection {
                         stream
                             .send(Password::Md5 {
                                 username: &options.username,
-                                password: options.password.as_deref().unwrap_or_default(),
+                                password: options.password.get()?.as_ref(),
                                 salt: body.salt,
                             })
                             .await?;

sqlx-postgres/src/connection/sasl.rs

@@ -87,7 +87,7 @@ pub(crate) async fn authenticate(
 
     // SaltedPassword := Hi(Normalize(password), salt, i)
     let salted_password = hi(
-        options.password.as_deref().unwrap_or_default(),
+        options.password.get()?.as_ref(),
         &cont.salt,
         cont.iterations,
     )?;

sqlx-postgres/src/lib.rs

@@ -54,7 +54,7 @@ pub use database::Postgres;
 pub use error::{PgDatabaseError, PgErrorPosition};
 pub use listener::{PgListener, PgNotification};
 pub use message::PgSeverity;
-pub use options::{PgConnectOptions, PgSslMode};
+pub use options::{PasswordProvider, PgConnectOptions, PgSslMode};
 pub use query_result::PgQueryResult;
 pub use row::PgRow;
 pub use statement::PgStatement;