From 6329b7ecb8c343d848d287aa0d7e71b460582421 Mon Sep 17 00:00:00 2001 From: Michael Cervantes <52178+emcee21@users.noreply.github.com> Date: Fri, 14 Mar 2025 21:06:56 +0000 Subject: [PATCH 01/11] implement rds iam token fetch --- postgres/Cargo.toml | 2 + postgres/src/auth.rs | 78 +++++++++++ postgres/src/aws.rs | 275 +++++++++++++++++++++++++++++++++++++ postgres/src/config.rs | 8 +- postgres/src/lib.rs | 40 +++++- postgres/tests/postgres.rs | 5 +- 6 files changed, 403 insertions(+), 5 deletions(-) create mode 100644 postgres/src/auth.rs create mode 100644 postgres/src/aws.rs diff --git a/postgres/Cargo.toml b/postgres/Cargo.toml index 1164298e..62176eb2 100644 --- a/postgres/Cargo.toml +++ b/postgres/Cargo.toml @@ -14,6 +14,7 @@ readme = "README.md" all-features = true [features] +aws = ["dep:aws-rds-signer"] default = ["rt_tokio_1"] rt_tokio_1 = ["deadpool/rt_tokio_1"] rt_async-std_1 = ["deadpool/rt_async-std_1"] @@ -42,6 +43,7 @@ with-uuid-1 = ["tokio-postgres/with-uuid-1"] [dependencies] async-trait = "0.1.80" +aws-rds-signer = { version = "0.2", optional = true } deadpool = { path = "../", version = "0.12.0", default-features = false, features = [ "managed", ] } diff --git a/postgres/src/auth.rs b/postgres/src/auth.rs new file mode 100644 index 00000000..3ce66770 --- /dev/null +++ b/postgres/src/auth.rs @@ -0,0 +1,78 @@ +/// Token fetcher for authentication. +/// +/// This enum represents different methods of fetching authentication tokens: +/// - Default: Uses a static password/token +/// - AWS RDS: Dynamically fetches tokens for AWS RDS IAM authentication +#[derive(Debug)] +pub(super) enum AuthTokenFetcher { + /// Default authentication using a static password/token + Default(Vec), + #[cfg(feature = "aws")] + /// AWS RDS IAM authentication token fetcher + AwsRds(crate::aws::AwsRdsInner), +} + +impl AuthTokenFetcher { + /// Creates a new AuthTokenFetcher based on the provided configuration. + /// + /// # Arguments + /// + /// * `config` - The manager configuration + /// * `pg_config` - The PostgreSQL connection configuration + /// + /// # Returns + /// + /// Returns an AuthTokenFetcher configured based on the provided settings + pub(super) fn for_config(config: &super::ManagerConfig, pg_config: &super::PgConfig) -> Self { + #[cfg(not(feature = "aws"))] { + let _ = config; + Self::default(pg_config) + } + #[cfg(feature = "aws")] { + crate::aws::for_config(config, pg_config) + } + } + + /// Creates a default token fetcher with a static password. + pub(super)fn default(pg_config: &super::PgConfig) -> Self { + Self::Default(pg_config.get_password().unwrap_or_default().to_vec()) + } + + /// Fetches a new token if needed. + /// + /// For AWS RDS authentication, this will fetch a new token if the current one + /// has expired. For default authentication, this is a no-op. + pub(super) async fn fetch_token(&self) { + match self { + #[cfg(feature = "aws")] + AuthTokenFetcher::AwsRds(inner) => crate::aws::fetch_token(inner).await, + _ => {} + } + } + + /// Checks if a new token needs to be fetched. + /// + /// # Returns + /// + /// Returns true if a new token should be fetched, false otherwise. + pub(super) async fn is_fetch_needed(&self) -> bool { + match self { + #[cfg(feature = "aws")] + AuthTokenFetcher::AwsRds(inner) => inner.read().await.is_fetch_needed(), + _ => false, + } + } + + /// Gets the current authentication token. + /// + /// # Returns + /// + /// Returns the current authentication token as a byte vector. + pub(super) async fn token(&self) -> Vec { + match self { + AuthTokenFetcher::Default(token) => token.clone(), + #[cfg(feature = "aws")] + AuthTokenFetcher::AwsRds(inner) => inner.read().await.token().as_bytes().to_vec(), + } + } +} diff --git a/postgres/src/aws.rs b/postgres/src/aws.rs new file mode 100644 index 00000000..3243437b --- /dev/null +++ b/postgres/src/aws.rs @@ -0,0 +1,275 @@ +use std::time::Duration; + +use tokio::{sync::RwLock, time::Instant}; +use tokio_postgres::config::Host; + +use crate::auth::AuthTokenFetcher; + +use super::ManagerConfig; + +/// Default expiration time in seconds for AWS RDS IAM authentication tokens. +/// AWS recommends using a token lifetime between 5 minutes and 15 minutes. +const DEFAULT_EXPIRES_IN: u64 = 900; + +pub(super) type AwsRdsInner = RwLock; + +/// Configuration for AWS RDS IAM authentication. +/// +/// This struct holds configuration options for AWS RDS IAM authentication, +/// including the AWS region and token expiration duration. +#[derive(Debug, Clone, Default)] +#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] +pub struct AwsRdsSignerConfig { + region: Option, + expires_in: Duration, +} + +impl AwsRdsSignerConfig { + /// Creates a new builder for AwsRdsSignerConfig. + /// + /// Use this method to start building a new configuration with custom settings. + /// + /// # Returns + /// + /// Returns a new `AwsRdsSignerConfigBuilder` instance. + pub fn builder() -> AwsRdsSignerConfigBuilder { + AwsRdsSignerConfigBuilder::default() + } + + /// Gets the configured AWS region. + /// + /// # Returns + /// + /// Returns an `Option<&str>` containing the configured AWS region if set, + /// or `None` if no region was configured. + pub fn region(&self) -> Option<&str> { + self.region.as_deref() + } + + /// Gets the configured token expiration duration. + /// + /// # Returns + /// + /// Returns a `Duration` representing how long generated tokens will be valid for. + /// If not explicitly set during configuration, this will be the default duration + /// of 900 seconds (15 minutes). + pub fn expires_in(&self) -> Duration { + self.expires_in + } +} + +/// Builder for creating an `AwsRdsSignerConfig` instance. +/// +/// This builder provides a fluent interface for configuring AWS RDS IAM authentication +/// settings. Use this to create a new `AwsRdsSignerConfig` with custom settings. +#[derive(Debug, Default)] +pub struct AwsRdsSignerConfigBuilder { + region: Option, + expires_in: Option, +} + +impl AwsRdsSignerConfigBuilder { + /// Sets the AWS region for the signer. + /// + /// # Arguments + /// + /// * `region` - The AWS region where the RDS instance is located (e.g., "us-east-1") + /// + /// # Returns + /// + /// Returns the builder instance for method chaining. + pub fn region(mut self, region: impl Into) -> Self { + self.region = Some(region.into()); + self + } + + /// Sets the expiration duration for generated authentication tokens. + /// + /// # Arguments + /// + /// * `duration` - The duration after which the generated token will expire. + /// If not set, defaults to 900 seconds (15 minutes). + /// + /// # Returns + /// + /// Returns the builder instance for method chaining. + pub fn expires_in(mut self, duration: impl Into) -> Self { + self.expires_in = Some(duration.into()); + self + } + + /// Builds and returns the final `AwsRdsSignerConfig` instance. + /// + /// If `expires_in` was not set, it will use the default value of 900 seconds (15 minutes). + /// The region is optional and will be `None` if not set. + /// + /// # Returns + /// + /// Returns a configured `AwsRdsSignerConfig` instance. + pub fn build(self) -> AwsRdsSignerConfig { + AwsRdsSignerConfig { + region: self.region, + expires_in: self + .expires_in + .unwrap_or(Duration::from_secs(DEFAULT_EXPIRES_IN)), + } + } +} + +impl ManagerConfig { + /// Creates an AWS RDS signer configured with the connection details. + /// + /// # Arguments + /// + /// * `config` - The Postgres connection configuration + /// + /// # Returns + /// + /// Returns a configured `aws_rds_signer::Signer` instance ready for generating authentication tokens. + pub(super) fn get_rds_signer(&self, config: &tokio_postgres::Config) -> aws_rds_signer::Signer { + let Some(signer_config) = &self.aws_rds_signer_config else { + tracing::warn!("AWS RDS signer config is not set, using default signer"); + return aws_rds_signer::Signer::default(); + }; + let host = host_to_string(&config.get_hosts()[0]); + let port = config.get_ports()[0]; + let mut signer = aws_rds_signer::Signer::builder().host(host).port(port); + if let Some(region) = signer_config.region() { + signer = signer.region(region); + } + if let Some(user) = &config.get_user() { + signer = signer.user(user.to_string()); + } + signer.build() + } + + /// Checks if AWS RDS IAM authentication is enabled. + /// + /// # Returns + /// + /// Returns `true` if AWS RDS IAM authentication is configured and enabled, + /// `false` otherwise. + pub fn is_rds_signer_enabled(&self) -> bool { + self.aws_rds_signer_config.is_some() + } +} + +/// Converts a Postgres host configuration to a string representation. +/// +/// # Arguments +/// +/// * `host` - The host configuration, either TCP or Unix socket +/// +/// # Returns +/// +/// Returns a string representation of the host: +/// * For TCP hosts, returns the hostname +/// * For Unix sockets, returns the path as a string +fn host_to_string(host: &Host) -> String { + match host { + Host::Tcp(host) => host.to_string(), + Host::Unix(path) => path.to_string_lossy().to_string(), + } +} + +/// Internal state for AWS RDS authentication token fetching. +#[derive(Debug)] +pub(super) struct AuthTokenFetcherInner { + /// Duration after which a token expires and needs to be refreshed + expires_in: Duration, + /// Timestamp of when the last token was fetched + last_token_fetch: Option, + /// AWS RDS signer for generating authentication tokens + signer: aws_rds_signer::Signer, + /// The current authentication token + token: String, +} + +impl AuthTokenFetcher { + /// Creates a new AWS RDS token fetcher. + /// + /// # Arguments + /// + /// * `expires_in` - Duration after which the token expires + /// * `signer` - AWS RDS signer for generating tokens + /// + /// # Returns + /// + /// Returns a new AuthTokenFetcher configured for AWS RDS authentication + pub(super) fn aws_rds( + expires_in: Duration, + signer: aws_rds_signer::Signer, + ) -> AuthTokenFetcher { + AuthTokenFetcher::AwsRds(RwLock::new(AuthTokenFetcherInner { + expires_in, + last_token_fetch: None, + signer, + token: String::new(), + })) + } +} + +impl AuthTokenFetcherInner { + /// Checks if a new token needs to be fetched based on expiration time. + /// + /// # Returns + /// + /// Returns true if the current token has expired or no token exists, + /// false otherwise. + pub(super) fn is_fetch_needed(&self) -> bool { + if let Some(last_token_fetch) = self.last_token_fetch { + if last_token_fetch.elapsed() < self.expires_in { + return false; + } + } + true + } + + /// Fetches a new authentication token from AWS RDS. + /// + /// Updates the internal token and last fetch timestamp if successful. + /// Logs an error if token fetching fails. + pub(super) async fn fetch_token(&mut self) { + match self.signer.fetch_token().await { + Ok(token) => { + self.token = token; + self.last_token_fetch = Some(Instant::now()); + } + Err(e) => { + tracing::error!(target: "deadpool.postgres", "Failed to fetch RDS signer token: {}", e); + } + } + } + + /// Gets the current authentication token. + /// + /// # Returns + /// + /// Returns the current authentication token as a string slice. + pub(super) fn token(&self) -> &str { + &self.token + } +} + +/// Fetches a new token if needed by checking expiration and updating the token. +/// +/// # Arguments +/// +/// * `inner` - The AuthTokenFetcherInner containing token state +pub(super) async fn fetch_token(inner: &AwsRdsInner) { + if inner.read().await.is_fetch_needed() { + inner.write().await.fetch_token().await; + } +} + +pub(super) fn for_config( + config: &ManagerConfig, + pg_config: &tokio_postgres::Config, +) -> AuthTokenFetcher { + config.aws_rds_signer_config.as_ref().map_or( + AuthTokenFetcher::default(pg_config), + |signer_config| { + AuthTokenFetcher::aws_rds(signer_config.expires_in(), config.get_rds_signer(pg_config)) + }, + ) +} diff --git a/postgres/src/config.rs b/postgres/src/config.rs index 3a1d0188..133478e3 100644 --- a/postgres/src/config.rs +++ b/postgres/src/config.rs @@ -1,5 +1,4 @@ //! Configuration used for [`Pool`] creation. - use std::{env, fmt, net::IpAddr, str::FromStr, time::Duration}; use tokio_postgres::config::{ @@ -387,6 +386,13 @@ impl RecyclingMethod { #[derive(Clone, Debug, Default)] #[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))] pub struct ManagerConfig { + #[cfg(feature = "aws")] + /// AWS RDS signer mode. + pub aws_rds_signer_config: Option, + + /// Maximum age of a connection. + pub max_age: Option, + /// Method of how a connection is recycled. See [`RecyclingMethod`]. pub recycling_method: RecyclingMethod, } diff --git a/postgres/src/lib.rs b/postgres/src/lib.rs index 2bb4d8e1..fcc1c86f 100644 --- a/postgres/src/lib.rs +++ b/postgres/src/lib.rs @@ -20,9 +20,14 @@ unused_results )] +mod auth; mod config; mod generic_client; +#[cfg(feature = "aws")] +/// AWS RDS IAM authentication support for PostgreSQL connections. +pub mod aws; + use std::{ borrow::Cow, collections::HashMap, @@ -36,10 +41,11 @@ use std::{ }, }; +use auth::AuthTokenFetcher; use deadpool::managed; #[cfg(not(target_arch = "wasm32"))] use tokio::spawn; -use tokio::task::JoinHandle; +use tokio::{task::JoinHandle, time::Instant}; use tokio_postgres::{ types::Type, Client as PgClient, Config as PgConfig, Error, IsolationLevel, Statement, Transaction as PgTransaction, TransactionBuilder as PgTransactionBuilder, @@ -112,7 +118,15 @@ impl Manager { T::TlsConnect: Sync + Send, >::Future: Send, { - Self::from_connect(pg_config, ConfigConnectImpl { tls }, config) + let auth_token_fetcher = AuthTokenFetcher::for_config(&config, &pg_config); + Self::from_connect( + pg_config, + ConfigConnectImpl { + auth_token_fetcher, + tls, + }, + config, + ) } /// Create a new [`Manager`] using the given [`tokio_postgres::Config`], and @@ -159,6 +173,12 @@ impl managed::Manager for Manager { tracing::warn!(target: "deadpool.postgres", "Connection could not be recycled: Connection closed"); return Err(RecycleError::message("Connection closed")); } + if let Some(max_age) = self.config.max_age { + if client.created_at().elapsed() > max_age { + tracing::debug!(target: "deadpool.postgres", "Connection could not be recycled: Connection expired"); + return Err(RecycleError::message("Connection expired")); + } + } match self.config.recycling_method.query() { Some(sql) => match client.simple_query(sql).await { Ok(_) => Ok(()), @@ -201,6 +221,7 @@ where { /// The TLS connector to use for the connection. pub tls: T, + auth_token_fetcher: AuthTokenFetcher, } #[cfg(not(target_arch = "wasm32"))] @@ -216,8 +237,13 @@ where pg_config: &PgConfig, ) -> BoxFuture<'_, Result<(PgClient, JoinHandle<()>), Error>> { let tls = self.tls.clone(); - let pg_config = pg_config.clone(); + let mut pg_config = pg_config.clone(); Box::pin(async move { + if self.auth_token_fetcher.is_fetch_needed().await { + tracing::debug!(target: "deadpool.postgres", "Fetching token"); + self.auth_token_fetcher.fetch_token().await; + } + let _ = pg_config.password(self.auth_token_fetcher.token().await); let fut = pg_config.connect(tls); let (client, connection) = fut.await?; let conn_task = spawn(async move { @@ -415,6 +441,8 @@ pub struct ClientWrapper { /// wrapper is dropped. conn_task: JoinHandle<()>, + created_at: Instant, + /// [`StatementCache`] of this client. pub statement_cache: Arc, } @@ -427,10 +455,16 @@ impl ClientWrapper { Self { client, conn_task, + created_at: Instant::now(), statement_cache: Arc::new(StatementCache::new()), } } + /// Returns the timestamp when this client was created. + pub fn created_at(&self) -> Instant { + self.created_at + } + /// Like [`tokio_postgres::Client::prepare()`], but uses an existing /// [`Statement`] from the [`StatementCache`] if possible. pub async fn prepare_cached(&self, query: &str) -> Result { diff --git a/postgres/tests/postgres.rs b/postgres/tests/postgres.rs index 7e64b42a..f4add265 100644 --- a/postgres/tests/postgres.rs +++ b/postgres/tests/postgres.rs @@ -164,7 +164,10 @@ async fn recycling_methods() { ]; let mut cfg = Config::from_env(); for recycling_method in recycling_methods { - cfg.pg.manager = Some(ManagerConfig { recycling_method }); + cfg.pg.manager = Some(ManagerConfig { + recycling_method, + ..Default::default() + }); let pool = cfg .pg .create_pool(Some(Runtime::Tokio1), tokio_postgres::NoTls) From d6ef57afb67a5dfd85f0b1125195a8fdf08f54ca Mon Sep 17 00:00:00 2001 From: Michael Cervantes <52178+emcee21@users.noreply.github.com> Date: Fri, 14 Mar 2025 21:23:17 +0000 Subject: [PATCH 02/11] no empty regions --- postgres/src/aws.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/postgres/src/aws.rs b/postgres/src/aws.rs index 3243437b..7bca4ac3 100644 --- a/postgres/src/aws.rs +++ b/postgres/src/aws.rs @@ -79,7 +79,12 @@ impl AwsRdsSignerConfigBuilder { /// /// Returns the builder instance for method chaining. pub fn region(mut self, region: impl Into) -> Self { - self.region = Some(region.into()); + let region = region.into(); + if region.is_empty() { + self.region = None; + } else { + self.region = Some(region); + } self } From a18b6fb781f66e7d2cbd83ba5c9014a1ef2540a0 Mon Sep 17 00:00:00 2001 From: Michael Cervantes <52178+emcee21@users.noreply.github.com> Date: Wed, 19 Mar 2025 14:45:42 +0000 Subject: [PATCH 03/11] add trace logging --- postgres/src/aws.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/postgres/src/aws.rs b/postgres/src/aws.rs index 7bca4ac3..96d74d5d 100644 --- a/postgres/src/aws.rs +++ b/postgres/src/aws.rs @@ -145,6 +145,7 @@ impl ManagerConfig { if let Some(user) = &config.get_user() { signer = signer.user(user.to_string()); } + tracing::debug!(target: "deadpool.postgres", "AWS RDS signer: {:?}", signer); signer.build() } @@ -226,6 +227,9 @@ impl AuthTokenFetcherInner { if last_token_fetch.elapsed() < self.expires_in { return false; } + tracing::debug!(target: "deadpool.postgres", "Token expired, fetch needed"); + } else { + tracing::debug!(target: "deadpool.postgres", "No token found, fetch needed"); } true } @@ -235,10 +239,16 @@ impl AuthTokenFetcherInner { /// Updates the internal token and last fetch timestamp if successful. /// Logs an error if token fetching fails. pub(super) async fn fetch_token(&mut self) { + tracing::debug!(target: "deadpool.postgres", "Fetching RDS token"); match self.signer.fetch_token().await { Ok(token) => { self.token = token; self.last_token_fetch = Some(Instant::now()); + tracing::debug!( + target: "deadpool.postgres", + "RDS token fetched successfully at {:?}", + self.last_token_fetch + ); } Err(e) => { tracing::error!(target: "deadpool.postgres", "Failed to fetch RDS signer token: {}", e); @@ -274,6 +284,7 @@ pub(super) fn for_config( config.aws_rds_signer_config.as_ref().map_or( AuthTokenFetcher::default(pg_config), |signer_config| { + tracing::debug!(target: "deadpool.postgres", "Creating AuthTokenFetcher with config: {:?}", signer_config); AuthTokenFetcher::aws_rds(signer_config.expires_in(), config.get_rds_signer(pg_config)) }, ) From 92ec655d1b2559fadc177fb92070b8fb82e8e09e Mon Sep 17 00:00:00 2001 From: Michael Cervantes <52178+emcee21@users.noreply.github.com> Date: Wed, 19 Mar 2025 15:45:39 +0000 Subject: [PATCH 04/11] Refactor authentication token fetching: rename methods for clarity and improve token handling. Update `fetch_token` to `fetch_token_if_needed` and introduce `with_token` for executing functions with the current token. --- postgres/src/auth.rs | 64 ++++++++++++++++++++++---------------------- postgres/src/aws.rs | 2 +- postgres/src/lib.rs | 9 +++---- 3 files changed, 37 insertions(+), 38 deletions(-) diff --git a/postgres/src/auth.rs b/postgres/src/auth.rs index 3ce66770..cc25f362 100644 --- a/postgres/src/auth.rs +++ b/postgres/src/auth.rs @@ -1,5 +1,5 @@ /// Token fetcher for authentication. -/// +/// /// This enum represents different methods of fetching authentication tokens: /// - Default: Uses a static password/token /// - AWS RDS: Dynamically fetches tokens for AWS RDS IAM authentication @@ -14,65 +14,65 @@ pub(super) enum AuthTokenFetcher { impl AuthTokenFetcher { /// Creates a new AuthTokenFetcher based on the provided configuration. - /// + /// /// # Arguments - /// + /// /// * `config` - The manager configuration /// * `pg_config` - The PostgreSQL connection configuration - /// + /// /// # Returns - /// + /// /// Returns an AuthTokenFetcher configured based on the provided settings pub(super) fn for_config(config: &super::ManagerConfig, pg_config: &super::PgConfig) -> Self { - #[cfg(not(feature = "aws"))] { + #[cfg(not(feature = "aws"))] + { let _ = config; Self::default(pg_config) } - #[cfg(feature = "aws")] { + #[cfg(feature = "aws")] + { crate::aws::for_config(config, pg_config) } } /// Creates a default token fetcher with a static password. - pub(super)fn default(pg_config: &super::PgConfig) -> Self { + pub(super) fn default(pg_config: &super::PgConfig) -> Self { Self::Default(pg_config.get_password().unwrap_or_default().to_vec()) } /// Fetches a new token if needed. - /// + /// /// For AWS RDS authentication, this will fetch a new token if the current one /// has expired. For default authentication, this is a no-op. - pub(super) async fn fetch_token(&self) { + pub(super) async fn fetch_token_if_needed(&self) { match self { #[cfg(feature = "aws")] - AuthTokenFetcher::AwsRds(inner) => crate::aws::fetch_token(inner).await, - _ => {} + AuthTokenFetcher::AwsRds(inner) => crate::aws::fetch_token_if_needed(inner).await, + _ => {}, } } - /// Checks if a new token needs to be fetched. - /// - /// # Returns - /// - /// Returns true if a new token should be fetched, false otherwise. - pub(super) async fn is_fetch_needed(&self) -> bool { - match self { - #[cfg(feature = "aws")] - AuthTokenFetcher::AwsRds(inner) => inner.read().await.is_fetch_needed(), - _ => false, - } - } - - /// Gets the current authentication token. - /// + /// Executes a provided function with the current authentication token. + /// + /// This method retrieves the current authentication token and passes it to the provided + /// function `f`. The function `f` is then executed with the token as its argument. + /// + /// # Arguments + /// + /// * `f` - A closure or function that takes a byte slice (`&[u8]`) representing the + /// authentication token and returns a value of type `R`. + /// /// # Returns - /// - /// Returns the current authentication token as a byte vector. - pub(super) async fn token(&self) -> Vec { + /// + /// Returns the result of the provided function `f` executed with the current authentication token. + pub(super) async fn with_token(&self, f: F) -> R + where + F: FnOnce(&[u8]) -> R, + { match self { - AuthTokenFetcher::Default(token) => token.clone(), + AuthTokenFetcher::Default(token) => f(token), #[cfg(feature = "aws")] - AuthTokenFetcher::AwsRds(inner) => inner.read().await.token().as_bytes().to_vec(), + AuthTokenFetcher::AwsRds(inner) => f(inner.read().await.token().as_bytes()), } } } diff --git a/postgres/src/aws.rs b/postgres/src/aws.rs index 96d74d5d..4e0f57e0 100644 --- a/postgres/src/aws.rs +++ b/postgres/src/aws.rs @@ -271,7 +271,7 @@ impl AuthTokenFetcherInner { /// # Arguments /// /// * `inner` - The AuthTokenFetcherInner containing token state -pub(super) async fn fetch_token(inner: &AwsRdsInner) { +pub(super) async fn fetch_token_if_needed(inner: &AwsRdsInner) { if inner.read().await.is_fetch_needed() { inner.write().await.fetch_token().await; } diff --git a/postgres/src/lib.rs b/postgres/src/lib.rs index fcc1c86f..9d12d261 100644 --- a/postgres/src/lib.rs +++ b/postgres/src/lib.rs @@ -239,11 +239,10 @@ where let tls = self.tls.clone(); let mut pg_config = pg_config.clone(); Box::pin(async move { - if self.auth_token_fetcher.is_fetch_needed().await { - tracing::debug!(target: "deadpool.postgres", "Fetching token"); - self.auth_token_fetcher.fetch_token().await; - } - let _ = pg_config.password(self.auth_token_fetcher.token().await); + self.auth_token_fetcher.fetch_token_if_needed().await; + self.auth_token_fetcher.with_token(|token| { + let _ = pg_config.password(token); + }).await; let fut = pg_config.connect(tls); let (client, connection) = fut.await?; let conn_task = spawn(async move { From 479a0b6336dcebabd1e7ff017ece894949149183 Mon Sep 17 00:00:00 2001 From: Michael Cervantes <52178+emcee21@users.noreply.github.com> Date: Wed, 19 Mar 2025 15:47:03 +0000 Subject: [PATCH 05/11] Update CI workflow to include 'iam' branch in push triggers --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index df717912..aa0a557d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -2,7 +2,7 @@ name: CI on: push: - branches: [ "master" ] + branches: [ "iam", "master" ] tags: [ "deadpool-*" ] pull_request: branches: [ "master" ] From 131ffad6743a10802a4281ef24b334f3b4939377 Mon Sep 17 00:00:00 2001 From: Michael Cervantes <52178+emcee21@users.noreply.github.com> Date: Wed, 19 Mar 2025 15:49:44 +0000 Subject: [PATCH 06/11] rustfmt --- postgres/src/auth.rs | 2 +- postgres/src/lib.rs | 8 +++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/postgres/src/auth.rs b/postgres/src/auth.rs index cc25f362..b4394d90 100644 --- a/postgres/src/auth.rs +++ b/postgres/src/auth.rs @@ -48,7 +48,7 @@ impl AuthTokenFetcher { match self { #[cfg(feature = "aws")] AuthTokenFetcher::AwsRds(inner) => crate::aws::fetch_token_if_needed(inner).await, - _ => {}, + _ => {} } } diff --git a/postgres/src/lib.rs b/postgres/src/lib.rs index 9d12d261..2862bdd3 100644 --- a/postgres/src/lib.rs +++ b/postgres/src/lib.rs @@ -240,9 +240,11 @@ where let mut pg_config = pg_config.clone(); Box::pin(async move { self.auth_token_fetcher.fetch_token_if_needed().await; - self.auth_token_fetcher.with_token(|token| { - let _ = pg_config.password(token); - }).await; + self.auth_token_fetcher + .with_token(|token| { + let _ = pg_config.password(token); + }) + .await; let fut = pg_config.connect(tls); let (client, connection) = fut.await?; let conn_task = spawn(async move { From 4760060c1028d3eb74df4c69a63ee64f989559f9 Mon Sep 17 00:00:00 2001 From: Michael Cervantes <52178+emcee21@users.noreply.github.com> Date: Wed, 19 Mar 2025 15:56:03 +0000 Subject: [PATCH 07/11] Add default configuration for ManagerConfig in README examples --- postgres/README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/postgres/README.md b/postgres/README.md index 8553a2cc..537e3eb7 100644 --- a/postgres/README.md +++ b/postgres/README.md @@ -40,6 +40,7 @@ async fn main() { cfg.dbname = Some("deadpool".to_string()); cfg.manager = Some(ManagerConfig { recycling_method: RecyclingMethod::Fast, + ..Default::default() }); let pool = cfg.create_pool(Some(Runtime::Tokio1), NoTls).unwrap(); for i in 1..10i32 { @@ -114,6 +115,7 @@ async fn main() { pg_config.dbname("deadpool"); let mgr_config = ManagerConfig { recycling_method: RecyclingMethod::Fast, + ..Default::default() }; let mgr = Manager::from_config(pg_config, NoTls, mgr_config); let pool = Pool::builder(mgr).max_size(16).build().unwrap(); From b986c2b3a10af6a5b1fa702ce6d3d2a90d6dd91b Mon Sep 17 00:00:00 2001 From: Michael Cervantes <52178+emcee21@users.noreply.github.com> Date: Wed, 19 Mar 2025 17:34:50 +0000 Subject: [PATCH 08/11] Fix token handling in AwsRds variant: change token retrieval to use `as_ref()` for improved efficiency. --- postgres/src/auth.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/postgres/src/auth.rs b/postgres/src/auth.rs index b4394d90..fa754d83 100644 --- a/postgres/src/auth.rs +++ b/postgres/src/auth.rs @@ -72,7 +72,7 @@ impl AuthTokenFetcher { match self { AuthTokenFetcher::Default(token) => f(token), #[cfg(feature = "aws")] - AuthTokenFetcher::AwsRds(inner) => f(inner.read().await.token().as_bytes()), + AuthTokenFetcher::AwsRds(inner) => f(inner.read().await.token().as_ref()), } } } From b38dcd26c4f64c3a1dd0805ede07531915aef52f Mon Sep 17 00:00:00 2001 From: Michael Cervantes <52178+emcee21@users.noreply.github.com> Date: Wed, 19 Mar 2025 18:25:03 +0000 Subject: [PATCH 09/11] Refactor authentication token fetching: rename `fetch_token_if_needed` to `fetch_token` and streamline token handling for AWS RDS and default configurations. --- postgres/src/auth.rs | 35 ++++++++--------------------------- postgres/src/aws.rs | 7 ++++++- postgres/src/lib.rs | 7 +------ 3 files changed, 15 insertions(+), 34 deletions(-) diff --git a/postgres/src/auth.rs b/postgres/src/auth.rs index fa754d83..9ac0cd9e 100644 --- a/postgres/src/auth.rs +++ b/postgres/src/auth.rs @@ -44,35 +44,16 @@ impl AuthTokenFetcher { /// /// For AWS RDS authentication, this will fetch a new token if the current one /// has expired. For default authentication, this is a no-op. - pub(super) async fn fetch_token_if_needed(&self) { + /// + pub(super) async fn fetch_token(&self, pg_config: &mut super::PgConfig) { match self { #[cfg(feature = "aws")] - AuthTokenFetcher::AwsRds(inner) => crate::aws::fetch_token_if_needed(inner).await, - _ => {} - } - } - - /// Executes a provided function with the current authentication token. - /// - /// This method retrieves the current authentication token and passes it to the provided - /// function `f`. The function `f` is then executed with the token as its argument. - /// - /// # Arguments - /// - /// * `f` - A closure or function that takes a byte slice (`&[u8]`) representing the - /// authentication token and returns a value of type `R`. - /// - /// # Returns - /// - /// Returns the result of the provided function `f` executed with the current authentication token. - pub(super) async fn with_token(&self, f: F) -> R - where - F: FnOnce(&[u8]) -> R, - { - match self { - AuthTokenFetcher::Default(token) => f(token), - #[cfg(feature = "aws")] - AuthTokenFetcher::AwsRds(inner) => f(inner.read().await.token().as_ref()), + AuthTokenFetcher::AwsRds(inner) => { + crate::aws::fetch_token(inner, pg_config).await; + } + AuthTokenFetcher::Default(token) => { + let _ = pg_config.password(token); + } } } } diff --git a/postgres/src/aws.rs b/postgres/src/aws.rs index 4e0f57e0..32414867 100644 --- a/postgres/src/aws.rs +++ b/postgres/src/aws.rs @@ -271,10 +271,15 @@ impl AuthTokenFetcherInner { /// # Arguments /// /// * `inner` - The AuthTokenFetcherInner containing token state -pub(super) async fn fetch_token_if_needed(inner: &AwsRdsInner) { +pub(super) async fn fetch_token( + inner: &AwsRdsInner, + pg_config: &mut tokio_postgres::Config, +) { if inner.read().await.is_fetch_needed() { inner.write().await.fetch_token().await; } + let inner = inner.read().await; + let _ = pg_config.password(inner.token()); } pub(super) fn for_config( diff --git a/postgres/src/lib.rs b/postgres/src/lib.rs index 2862bdd3..83763839 100644 --- a/postgres/src/lib.rs +++ b/postgres/src/lib.rs @@ -239,12 +239,7 @@ where let tls = self.tls.clone(); let mut pg_config = pg_config.clone(); Box::pin(async move { - self.auth_token_fetcher.fetch_token_if_needed().await; - self.auth_token_fetcher - .with_token(|token| { - let _ = pg_config.password(token); - }) - .await; + self.auth_token_fetcher.fetch_token(&mut pg_config).await; let fut = pg_config.connect(tls); let (client, connection) = fut.await?; let conn_task = spawn(async move { From 9576448b37784483360cd507b6e5b0223058d5d1 Mon Sep 17 00:00:00 2001 From: Michael Cervantes <52178+emcee21@users.noreply.github.com> Date: Wed, 19 Mar 2025 19:04:36 +0000 Subject: [PATCH 10/11] Add debug logging for RDS token password setting in `fetch_token`: log the first few characters of the token for tracing purposes. --- postgres/src/aws.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/postgres/src/aws.rs b/postgres/src/aws.rs index 32414867..8f4e6978 100644 --- a/postgres/src/aws.rs +++ b/postgres/src/aws.rs @@ -279,7 +279,9 @@ pub(super) async fn fetch_token( inner.write().await.fetch_token().await; } let inner = inner.read().await; - let _ = pg_config.password(inner.token()); + let token = inner.token(); + tracing::debug!(target: "deadpool.postgres", "Setting password to RDS token: {}", &token[..5]); + let _ = pg_config.password(token); } pub(super) fn for_config( From 18e904ca597212b163580e7988e9f33e80380fe2 Mon Sep 17 00:00:00 2001 From: Michael Cervantes <52178+emcee21@users.noreply.github.com> Date: Wed, 19 Mar 2025 19:05:40 +0000 Subject: [PATCH 11/11] Update debug logging in `fetch_token` to display the first 10 characters of the RDS token for improved traceability. --- postgres/src/aws.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/postgres/src/aws.rs b/postgres/src/aws.rs index 8f4e6978..2e024641 100644 --- a/postgres/src/aws.rs +++ b/postgres/src/aws.rs @@ -280,7 +280,7 @@ pub(super) async fn fetch_token( } let inner = inner.read().await; let token = inner.token(); - tracing::debug!(target: "deadpool.postgres", "Setting password to RDS token: {}", &token[..5]); + tracing::debug!(target: "deadpool.postgres", "Setting password to RDS token: {}", &token[..10]); let _ = pg_config.password(token); }