feat: rds iam passthrough auth

integration/complex/rds_iam/dev.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+set -euo pipefail
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+
+cd "${SCRIPT_DIR}"
+
+if [[ ! -x venv/bin/python ]]; then
+    python3 -m venv venv
+fi
+
+venv/bin/python -m pip install -r requirements.txt
+
+# Requires PgDog on 127.0.0.1:6432 and configured AWS credentials.
+venv/bin/python -m pytest test_rds_iam.py -v --tb=short "$@"

integration/complex/rds_iam/pgdog.toml

@@ -0,0 +1,11 @@
+[general]
+host = "127.0.0.1"
+port = 6432
+auth_type = "external_token"
+
+[[databases]]
+name = "postgres"
+host = "staging-iam.cluster-c5icciqq4b0q.us-west-2.rds.amazonaws.com"
+
+[admin]
+password = "pgdog"

integration/complex/rds_iam/requirements.txt

@@ -0,0 +1,3 @@
+boto3
+psycopg==3.2.6
+pytest==8.3.5

integration/complex/rds_iam/test_rds_iam.py

@@ -0,0 +1,54 @@
+from urllib.parse import parse_qs, urlsplit
+
+import boto3
+import psycopg
+import pytest
+from botocore.exceptions import NoCredentialsError
+
+
+@pytest.fixture(scope="module")
+def rds_iam_token():
+    """Generate a staging-iam token using the default AWS credential chain."""
+    hostname = "staging-iam.cluster-c5icciqq4b0q.us-west-2.rds.amazonaws.com"
+    client = boto3.client("rds", region_name="us-west-2")
+
+    try:
+        return client.generate_db_auth_token(
+            DBHostname=hostname,
+            Port=5432,
+            DBUsername="postgres",
+            Region="us-west-2",
+        )
+    except NoCredentialsError:
+        pytest.skip("AWS credentials are required to generate an RDS IAM token")
+
+
+def test_generate_rds_iam_token(rds_iam_token):
+    url = urlsplit(f"https://{rds_iam_token}")
+    params = parse_qs(url.query)
+    assert (
+        url.hostname == "staging-iam.cluster-c5icciqq4b0q.us-west-2.rds.amazonaws.com"
+    )
+    assert url.port == 5432
+    assert params["Action"] == ["connect"]
+    assert params["DBUser"] == ["postgres"]
+    assert params["X-Amz-Expires"] == ["900"]
+    assert params["X-Amz-Credential"][0].endswith("/us-west-2/rds-db/aws4_request")
+    assert params["X-Amz-Signature"][0]
+
+
+def test_connect_to_pgdog(rds_iam_token):
+    """Reuse one IAM token across ten separate client connections."""
+    for attempt in range(10):
+        with psycopg.connect(
+            host="127.0.0.1",
+            port=6432,
+            dbname="postgres",
+            user="postgres",
+            password=rds_iam_token,
+            sslmode="disable",
+            connect_timeout=10,
+            options="-c statement_timeout=10000",
+        ) as conn:
+            row = conn.execute("SELECT current_user, current_database()").fetchone()
+            assert row == ("postgres", "postgres"), f"Connection {attempt + 1}"

integration/complex/rds_iam/users.toml

@@ -0,0 +1,4 @@
+[[users]]
+name = "postgres"
+database = "postgres"
+server_auth = "rds_iam"

pgdog-config/src/auth.rs

@@ -37,7 +37,7 @@ impl PassthroughAuth {
 /// See [authentication](https://docs.pgdog.dev/features/authentication/).
 ///
 /// <https://docs.pgdog.dev/configuration/pgdog.toml/general/#auth_type>
-#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, JsonSchema)]
+#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, JsonSchema, Copy)]
 #[serde(rename_all = "snake_case")]
 pub enum AuthType {
     /// MD5 password hashing; very quick but not secure.
@@ -49,6 +49,8 @@ pub enum AuthType {
     Trust,
     /// Plaintext password.
     Plain,
+    /// RDS IAM or Azure Workload identity.
+    ExternalToken,
 }
 
 impl Display for AuthType {
@@ -58,6 +60,7 @@ impl Display for AuthType {
             Self::Scram => write!(f, "scram"),
             Self::Trust => write!(f, "trust"),
             Self::Plain => write!(f, "plain"),
+            Self::ExternalToken => write!(f, "external_token"),
         }
     }
 }
@@ -74,6 +77,10 @@ impl AuthType {
     pub fn trust(&self) -> bool {
         matches!(self, Self::Trust)
     }
+
+    pub fn external_token(&self) -> bool {
+        matches!(self, Self::ExternalToken)
+    }
 }
 
 impl FromStr for AuthType {

pgdog-config/src/users.rs

@@ -51,9 +51,12 @@ impl Users {
     pub fn check(&mut self, config: &Config) {
         for user in &mut self.users {
             if user.passwords().is_empty() {
-                if !config.general.passthrough_auth() && user.identity.is_none() {
+                if !config.general.passthrough_auth()
+                    && user.identity.is_none()
+                    && !config.general.auth_type.external_token()
+                {
                     warn!(
-                        r#"user "{}" (database "{}") doesn't have a password, passthrough auth and mTLS are disabled"#,
+                        r#"user "{}" (database "{}") doesn't have a password, while passthrough auth, external token, and mTLS auth are disabled"#,
                         user.name, user.database,
                     );
                 }

pgdog/src/auth/mod.rs

@@ -4,6 +4,7 @@ pub(crate) mod auth_result;
 pub(crate) mod error;
 pub(crate) mod md5;
 pub(crate) mod scram;
+pub(crate) mod token_cache;
 pub(crate) mod vault;
 
 pub(crate) use auth_result::AuthResult;

pgdog/src/auth/token_cache.rs

@@ -0,0 +1,108 @@
+//! Token cache that validates a token we received
+//! from the client is actually valid with whatever mechanism
+//! we authenticate to Postgres with.
+//!
+//! This is effectively passthrough auth for RDS IAM, Azure Workload Identity, etc.
+
+use moka::sync::Cache;
+use once_cell::sync::Lazy;
+use rand::seq::{IndexedRandom, IteratorRandom};
+use std::{
+    sync::Arc,
+    time::{Duration, SystemTime},
+};
+use tokio::{sync::Mutex, time::sleep};
+
+use crate::backend::pool::Connection;
+use crate::tasks::{shutdown_signal, spawn};
+
+pub static AUTH_TOKEN_CACHE: Lazy<AuthTokenCache> = Lazy::new(|| {
+    let cache = AuthTokenCache::default();
+    let background = cache.clone();
+
+    spawn("auth token cache", async move {
+        let shutdown = shutdown_signal();
+        loop {
+            tokio::select! {
+                _ = sleep(Duration::from_millis(333)) => background.evict(),
+                _ = shutdown.cancelled() => break,
+            }
+        }
+    });
+
+    cache
+});
+
+#[derive(Debug, Clone)]
+struct Entry {
+    expires_at: SystemTime,
+    lock: Arc<Mutex<()>>,
+}
+
+#[derive(Clone, Debug)]
+pub(crate) struct AuthTokenCache {
+    tokens: Cache<String, Entry>,
+}
+
+impl Default for AuthTokenCache {
+    fn default() -> Self {
+        Self {
+            tokens: Cache::builder()
+                .max_capacity(1_000)
+                .support_invalidation_closures()
+                .build(),
+        }
+    }
+}
+
+impl AuthTokenCache {
+    fn evict(&self) {
+        let now = SystemTime::now();
+        self.tokens
+            .invalidate_entries_if(move |_, value| value.expires_at <= now)
+            .expect("invalidation closures are enabled");
+    }
+
+    /// Check a token against the external token provider, e.g., RDS.
+    ///
+    /// Takes a lock on the token so a connection storm from clients
+    /// only creates one connection to the actual database.
+    pub(crate) async fn check(
+        &self,
+        user: &str,
+        database: &str,
+        token: &str,
+    ) -> Result<bool, crate::frontend::Error> {
+        if self.tokens.contains_key(token) {
+            return Ok(true);
+        }
+
+        let lock = {
+            let entry = self.tokens.entry(token.to_owned()).or_insert(Entry {
+                expires_at: crate::backend::auth::rds_iam::expires_at(),
+                lock: Arc::new(Mutex::new(())),
+            });
+            entry.value().lock.clone()
+        };
+
+        let _guard = lock.lock().await;
+
+        let conn = Connection::new(user, database, false)?;
+        let shard = conn
+            .cluster()?
+            .shards()
+            .choose(&mut rand::rng())
+            .expect("to have at least one shard");
+        let pool = shard
+            .pool_iter()
+            .choose(&mut rand::rng())
+            .expect("to have at least one pool");
+
+        if let Ok(()) = pool.validate_token(token).await {
+            Ok(true)
+        } else {
+            self.tokens.remove(token);
+            Ok(false)
+        }
+    }
+}

pgdog/src/backend/auth/rds_iam.rs

@@ -135,9 +135,12 @@ pub(crate) async fn token(addr: Address) -> Result<(String, SystemTime), Error>
             ))
         })?;
 
-    // RDS IAM tokens are valid for 15 minutes.
-    let expires_at = SystemTime::now() + Duration::from_secs(900);
-    Ok((token, expires_at))
+    Ok((token, expires_at()))
+}
+
+// RDS IAM tokens are valid for 15 minutes.
+pub(crate) fn expires_at() -> SystemTime {
+    SystemTime::now() + Duration::from_secs(900)
 }
 
 #[cfg(test)]

pgdog/src/backend/connect_reason.rs

@@ -8,6 +8,7 @@ pub(crate) enum ConnectReason {
     PubSub,
     Probe,
     Healthcheck,
+    ValidateToken,
     #[default]
     Other,
 }

pgdog/src/backend/databases.rs

@@ -470,7 +470,10 @@ impl Databases {
 
         // Launch all clusters
         for cluster in self.all().values() {
-            if cluster.passwords().is_empty() && cluster.identity().is_none() {
+            if cluster.passwords().is_empty()
+                && cluster.identity().is_none()
+                && !cluster.auth_type().external_token()
+            {
                 warn!(
                     r#"disabling pool for user "{}" and database "{}", password not set"#,
                     cluster.user(),

pgdog/src/backend/disconnect_reason.rs

@@ -12,6 +12,7 @@ pub(crate) enum DisconnectReason {
     Unhealthy,
     Healthcheck,
     CredentialsRefresh,
+    CredentialsCheck,
     ServerClosed,
     #[default]
     Other,
@@ -24,14 +25,15 @@ impl Display for DisconnectReason {
             Self::Old => "max age",
             Self::Error => "error",
             Self::Other => "other",
-            Self::ForceClose => "force close",
-            Self::Offline => "pool offline",
+            Self::ForceClose => "force_close",
+            Self::Offline => "pool_offline",
             Self::OutOfSync => "out of sync",
-            Self::ReplicationMode => "in replication mode",
+            Self::ReplicationMode => "in_replication_mode",
             Self::Unhealthy => "unhealthy",
-            Self::Healthcheck => "standalone healthcheck",
-            Self::CredentialsRefresh => "credentials refresh",
-            Self::ServerClosed => "server closed",
+            Self::Healthcheck => "standalone_healthcheck",
+            Self::CredentialsRefresh => "credentials_refresh",
+            Self::ServerClosed => "server_closed",
+            Self::CredentialsCheck => "credentials_check",
         };
 
         write!(f, "{}", reason)

pgdog/src/backend/pool/address.rs

@@ -193,6 +193,19 @@ impl Address {
         }
     }
 
+    /// Create an instance of the same address, but swap the password
+    /// and authentication mechanism so the server is forced to use this password
+    /// no matter what `server_auth` is configured.
+    ///
+    /// This is used for validating RDS IAM tokens we receive from clients.
+    pub(super) fn into_client_token(self, token: &str) -> Self {
+        Self {
+            passwords: vec![Password::new(token, PasswordSource::ClientToken)],
+            server_auth: ServerAuth::Password,
+            ..self
+        }
+    }
+
     /// Test convention: `new_test()` represents a primary. Tests that need
     /// a replica do `Address { configured_role: Role::Replica, ..new_test() }`.
     #[cfg(test)]