From 56253c04cb679c58d00750da699a6cb62ed52aca Mon Sep 17 00:00:00 2001 From: Alexander Petric Date: Fri, 27 Mar 2026 17:50:29 -0400 Subject: [PATCH] feat: IAM RDS auth for PostgreSQL worker resources (#8573) * feat: add IAM RDS auth support for PostgreSQL worker resources Co-Authored-By: Claude Opus 4.6 * refactor: use Config builder for IAM RDS connections Co-Authored-By: Claude Opus 4.6 * fix: address PR review feedback for IAM RDS auth Co-Authored-By: Claude Opus 4.6 * chore: update ee-repo-ref to ebea6ef1e5bfcfc3f0151da9687dac6c61bbfab6 This commit updates the EE repository reference after PR #493 was merged in windmill-ee-private. Previous ee-repo-ref: 1228561a98c5195bb97a81d4a57ce2bb2ecfca79 New ee-repo-ref: ebea6ef1e5bfcfc3f0151da9687dac6c61bbfab6 Automated by sync-ee-ref workflow. --------- Co-authored-by: Claude Opus 4.6 Co-authored-by: windmill-internal-app[bot] --- backend/ee-repo-ref.txt | 2 +- backend/windmill-common/src/lib.rs | 73 ++++++++++++++++++++++ backend/windmill-worker/Cargo.toml | 2 +- backend/windmill-worker/src/pg_executor.rs | 26 +++++++- 4 files changed, 99 insertions(+), 4 deletions(-) diff --git a/backend/ee-repo-ref.txt b/backend/ee-repo-ref.txt index 4751795cd9..6a22b66a17 100644 --- a/backend/ee-repo-ref.txt +++ b/backend/ee-repo-ref.txt @@ -1 +1 @@ -208da6989ef606e4068663246903acbcaa90a9dc \ No newline at end of file +ebea6ef1e5bfcfc3f0151da9687dac6c61bbfab6 diff --git a/backend/windmill-common/src/lib.rs b/backend/windmill-common/src/lib.rs index b5ec518315..deb1c38a03 100644 --- a/backend/windmill-common/src/lib.rs +++ b/backend/windmill-common/src/lib.rs @@ -406,6 +406,8 @@ pub struct PgDatabase { pub sslmode: Option, pub dbname: String, pub root_certificate_pem: Option, + pub use_iam_auth: Option, + pub region: Option, } // Wrapper enum to hold either Tls or NoTls connection @@ -513,6 +515,75 @@ impl PgDatabase { } } + #[cfg(all(feature = "enterprise", feature = "private"))] + pub async fn connect_with_iam( + &self, + ) -> Result<(tokio_postgres::Client, TokioPgConnection), error::Error> { + use native_tls::TlsConnector; + use postgres_native_tls::MakeTlsConnector; + + // Resolve region: resource field takes priority, then env var + let region = match self.region.as_deref() { + Some(r) => r.to_string(), + None => std::env::var("AWS_REGION").map_err(|_| { + error::Error::BadConfig( + "Region is required for IAM RDS auth. Set 'region' on the resource or AWS_REGION env var".to_string(), + ) + })?, + }; + + let port = self.port.unwrap_or(5432); + let user = self.user.as_deref().unwrap_or("postgres"); + + let token = db_iam_ee::generate_auth_token(®ion, &self.host, port as u64, user) + .await + .map_err(|e| { + error::Error::InternalErr(format!("IAM token generation failed: {e:#}")) + })?; + + // RDS IAM auth requires SSL + let mut connector = TlsConnector::builder(); + if let Some(root_certificate_pem) = &self.root_certificate_pem { + if !root_certificate_pem.is_empty() { + connector.add_root_certificate( + native_tls::Certificate::from_pem(root_certificate_pem.as_bytes()) + .map_err(|e| error::Error::BadConfig(format!("Invalid Certs: {e:#}")))?, + ); + } else { + connector.danger_accept_invalid_certs(true); + connector.danger_accept_invalid_hostnames(true); + } + } else { + tracing::warn!("IAM RDS auth without root certificate: TLS certificate verification is disabled. Consider providing root_certificate_pem for production use."); + connector + .danger_accept_invalid_certs(true) + .danger_accept_invalid_hostnames(true); + } + + tracing::info!("Creating new IAM RDS connection to {}", &self.host); + + // Use Config builder directly to pass the IAM token as the password. + // This avoids needing to URL-encode the token into a connection string. + let mut config = tokio_postgres::Config::new(); + config + .host(&self.host) + .port(port as u16) + .user(user) + .password(&token) + .dbname(&self.dbname) + .ssl_mode(tokio_postgres::config::SslMode::Require); + + let (client, connection) = tokio::time::timeout( + std::time::Duration::from_secs(20), + config.connect(MakeTlsConnector::new(connector.build().map_err(to_anyhow)?)), + ) + .await + .map_err(to_anyhow)? + .map_err(to_anyhow)?; + + Ok((client, TokioPgConnection::Tls(connection))) + } + pub fn parse_uri(url: &str) -> Result { let parsed_url = url::Url::parse(url) .map_err(|_| Error::BadConfig("Invalid PostgreSQL URL".to_string()))?; @@ -551,6 +622,8 @@ impl PgDatabase { dbname, sslmode, root_certificate_pem: None, + use_iam_auth: None, + region: None, }) } } diff --git a/backend/windmill-worker/Cargo.toml b/backend/windmill-worker/Cargo.toml index c1e2a4927a..b65a2489ac 100644 --- a/backend/windmill-worker/Cargo.toml +++ b/backend/windmill-worker/Cargo.toml @@ -10,7 +10,7 @@ path = "src/lib.rs" [features] default = [] -private = ["windmill-worker-volumes/private", "windmill-queue/private"] +private = ["windmill-worker-volumes/private", "windmill-queue/private", "windmill-common/private"] mcp = ["dep:windmill-mcp"] prometheus = ["dep:prometheus", "windmill-common/prometheus"] enterprise = ["windmill-queue/enterprise", "windmill-git-sync/enterprise", "windmill-common/enterprise", "windmill-worker-volumes/enterprise", "dep:pem", "dep:tokio-util", "dep:opentelemetry-proto", "dep:prost", "dep:hudsucker", "dep:rcgen", "dep:hyper-http-proxy", "dep:hyper-tls", "dep:hyper-util"] diff --git a/backend/windmill-worker/src/pg_executor.rs b/backend/windmill-worker/src/pg_executor.rs index 5b309a611c..f769e1034d 100644 --- a/backend/windmill-worker/src/pg_executor.rs +++ b/backend/windmill-worker/src/pg_executor.rs @@ -285,7 +285,16 @@ pub async fn do_postgresql( annotations.result_collection }; - let database_string = database.to_uri(); + let use_iam_auth = database.use_iam_auth == Some(true); + + // Include use_iam_auth in cache key to distinguish IAM vs non-IAM connections to the same host. + // The cache key is static (doesn't include the token), which is correct because PostgreSQL + // connections remain valid after initial auth — fresh tokens are generated on cache miss. + let database_string = if use_iam_auth { + format!("{}?iam=true", database.to_uri()) + } else { + database.to_uri() + }; let database_string_clone = database_string.clone(); let mtex; @@ -309,7 +318,20 @@ pub async fn do_postgresql( ); (None, mtex) } else { - let (client, connection) = database.connect().await?; + let (client, connection) = if use_iam_auth { + #[cfg(all(feature = "enterprise", feature = "private"))] + { + database.connect_with_iam().await? + } + #[cfg(not(all(feature = "enterprise", feature = "private")))] + { + return Err(Error::ExecutionErr( + "IAM RDS authentication requires Windmill Enterprise Edition".to_string(), + )); + } + } else { + database.connect().await? + }; let handle = tokio::spawn(async move { if let Err(e) = connection.await { let mut mtex = CONNECTION_CACHE.lock().await;