From beef6e295cb4e01af00c279f406012a82d91c8c4 Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Mon, 3 Aug 2026 18:33:53 +0000 Subject: [PATCH] feat: azure workload identity auth for mssql and postgres resources (#10470) * feat: azure workload identity auth for mssql and postgres resources * refactor: keep mssql config lines untouched by the auth-mode change * fix: single-flight token refresh, cache eviction and identity-aware pg cache key * fix: back off after a failed entra id refresh and normalize blank pg identity fields * fix: re-check the fallback token lifetime after a failed refresh * refactor: select workload identity with a sentinel password instead of resource fields * fix: log the workload identity mode on the postgres path too --- .../src/azure_workload_identity.rs | 317 ++++++++++++++++++ backend/windmill-common/src/lib.rs | 48 ++- backend/windmill-worker/src/mssql_executor.rs | 14 +- backend/windmill-worker/src/pg_executor.rs | 129 +++++-- 4 files changed, 475 insertions(+), 33 deletions(-) create mode 100644 backend/windmill-common/src/azure_workload_identity.rs diff --git a/backend/windmill-common/src/azure_workload_identity.rs b/backend/windmill-common/src/azure_workload_identity.rs new file mode 100644 index 0000000000..fdbdd1f6b5 --- /dev/null +++ b/backend/windmill-common/src/azure_workload_identity.rs @@ -0,0 +1,317 @@ +//! Azure Workload Identity Federation. +//! +//! The Kubernetes-projected service account token of the pod is exchanged with Entra +//! ID for an access token, which Azure-hosted databases accept in place of a password. +//! Nothing long-lived is stored on the Windmill instance. + +use std::collections::HashMap; +use std::sync::{Arc, Mutex, RwLock}; +use std::time::{Duration, Instant}; + +use serde_json::Value; + +use crate::error::{Error, Result}; +use crate::utils::HTTP_CLIENT; + +/// Entra ID scope of Azure SQL / SQL Server. +pub const AZURE_SQL_SCOPE: &str = "https://database.windows.net/.default"; + +/// Entra ID scope of Azure Database for PostgreSQL / MySQL. +pub const AZURE_OSSRDBMS_SCOPE: &str = "https://ossrdbms-aad.database.windows.net/.default"; + +/// A database resource whose password is this authenticates as the worker's workload +/// identity, the way a `DATABASE_URL` whose password is `entraid` does for the instance +/// database. Carrying the mode in the password is what lets an existing deployment turn +/// it on: the resource type schemas live on the hub, and every database form already has +/// a password field. Anyone whose password is literally this string authenticates as the +/// worker's identity instead of failing to log in, which is why it is deliberately less +/// password-shaped than the instance's `entraid`: that one is set by whoever runs the +/// instance, this one sits where users keep their own secrets. +pub const WORKLOAD_IDENTITY_PASSWORD: &str = "ms_entraid"; + +/// Renew an access token this long before it expires. +const TOKEN_REFRESH_BUFFER: Duration = Duration::from_secs(5 * 60); + +/// A token with less life than this left is not worth handing to a connection: opening +/// one takes up to 20 seconds, and it authenticates at the far end of that. +const TOKEN_MIN_LIFETIME: Duration = Duration::from_secs(60); + +/// How long a failed renewal suppresses the next attempt, as long as the current token +/// still works. Without it every queued job retries the exchange in turn, so a +/// throttling Entra ID would cost each of them the request's full latency. +const REFRESH_RETRY_BACKOFF: Duration = Duration::from_secs(30); + +lazy_static::lazy_static! { + /// Access tokens keyed by identity and scope, shared by every job on the worker. + static ref TOKEN_CACHE: Mutex>> = Mutex::new(HashMap::new()); +} + +#[derive(Default)] +struct TokenSlot { + /// Read on the hit path, so it must never be held across the exchange. + token: RwLock>, + /// Held for the whole exchange, so a burst of jobs on a cold or expiring entry + /// makes one request to Entra ID instead of one per job. Guards the time of the + /// last failed exchange, which is what the waiting jobs need to see. + refreshing: tokio::sync::Mutex>, +} + +struct CachedToken { + token: String, + expires_at: Instant, +} + +impl TokenSlot { + fn token_valid_for(&self, remaining: Duration) -> Option { + let cached = self.token.read().unwrap(); + cached + .as_ref() + .filter(|cached| Instant::now() + remaining < cached.expires_at) + .map(|cached| cached.token.clone()) + } +} + +/// Whether to go to Entra ID, given the last failed exchange and whether the current +/// token would still serve. A cold slot always tries: there is nothing to fall back on. +fn should_attempt_refresh(last_failure: Option, has_usable_token: bool) -> bool { + match last_failure { + Some(at) if has_usable_token => at.elapsed() >= REFRESH_RETRY_BACKOFF, + _ => true, + } +} + +/// The federated credentials of the identity the worker authenticates as, all injected +/// into the pod by the Azure workload identity webhook. The identity is the worker's, +/// not the resource's: reaching two databases as two identities means two worker groups. +pub struct WorkloadIdentityConfig { + tenant_id: String, + client_id: String, + federated_token_file: String, + authority_host: String, +} + +impl WorkloadIdentityConfig { + pub fn resolve() -> Result { + fn required(env_var: &str) -> Result { + std::env::var(env_var) + .ok() + .filter(|v| !v.is_empty()) + .ok_or_else(|| { + Error::BadConfig(format!( + "Workload identity authentication requires the {} env var on the worker, \ + injected by the Azure workload identity webhook", + env_var + )) + }) + } + + Ok(Self { + tenant_id: required("AZURE_TENANT_ID")?, + client_id: required("AZURE_CLIENT_ID")?, + federated_token_file: required("AZURE_FEDERATED_TOKEN_FILE")?, + authority_host: std::env::var("AZURE_AUTHORITY_HOST") + .ok() + .filter(|v| !v.is_empty()) + .unwrap_or_else(|| "login.microsoftonline.com".to_string()), + }) + } + + pub fn client_id(&self) -> &str { + &self.client_id + } + + fn cache_key(&self, scope: &str) -> String { + format!( + "{}|{}|{}|{}", + self.authority_host, self.tenant_id, self.client_id, scope + ) + } + + /// AZURE_AUTHORITY_HOST is injected with a scheme and a trailing slash + /// (`https://login.microsoftonline.com/`), neither of which belongs in the path. + fn token_endpoint(&self) -> String { + let authority = self + .authority_host + .trim_start_matches("https://") + .trim_start_matches("http://") + .trim_end_matches('/'); + format!("https://{}/{}/oauth2/v2.0/token", authority, self.tenant_id) + } + + /// An Entra ID access token for `scope`, reusing the cached one while it is valid. + pub async fn access_token(&self, scope: &str) -> Result { + let slot = self.slot(scope); + if let Some(token) = slot.token_valid_for(TOKEN_REFRESH_BUFFER) { + return Ok(token); + } + + let mut last_failure = slot.refreshing.lock().await; + // Whoever held the lock may have just refreshed it. + if let Some(token) = slot.token_valid_for(TOKEN_REFRESH_BUFFER) { + return Ok(token); + } + + // Inside the refresh buffer the previous token still works, and Entra ID + // throttles bursts: a failed renewal must not fail an otherwise fine job. + let usable = slot.token_valid_for(TOKEN_MIN_LIFETIME); + if !should_attempt_refresh(*last_failure, usable.is_some()) { + return Ok(usable.unwrap()); + } + + match self.request_token(scope).await { + Ok(fresh) => { + let token = fresh.token.clone(); + *slot.token.write().unwrap() = Some(fresh); + *last_failure = None; + Ok(token) + } + Err(e) => { + *last_failure = Some(Instant::now()); + // Check again rather than trusting the pre-request value: a request that + // times out eats 20 seconds of whatever the token had left. + match slot.token_valid_for(TOKEN_MIN_LIFETIME) { + Some(token) => { + tracing::warn!("Keeping the current Entra ID token, renewal failed: {e:#}"); + Ok(token) + } + None => Err(e), + } + } + } + } + + fn slot(&self, scope: &str) -> Arc { + let mut cache = TOKEN_CACHE.lock().unwrap(); + // Drop what has expired and that no in-flight request is holding. + cache.retain(|_, slot| { + Arc::strong_count(slot) > 1 || slot.token_valid_for(TOKEN_MIN_LIFETIME).is_some() + }); + cache.entry(self.cache_key(scope)).or_default().clone() + } + + async fn request_token(&self, scope: &str) -> Result { + // The projected token rotates on disk, so it must be re-read on every exchange. + let assertion = tokio::fs::read_to_string(&self.federated_token_file) + .await + .map_err(|e| { + Error::ExecutionErr(format!( + "Failed to read the federated token at {}: {}", + self.federated_token_file, e + )) + })?; + + let url = self.token_endpoint(); + let response = HTTP_CLIENT + .post(&url) + .form(&[ + ("grant_type", "client_credentials"), + ("client_id", self.client_id.as_str()), + ( + "client_assertion_type", + "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", + ), + ("client_assertion", assertion.trim()), + ("scope", scope), + ]) + .send() + .await + .map_err(|e| { + Error::ExecutionErr(format!( + "Failed to request an Entra ID token from {url}: {e}" + )) + })?; + + let status = response.status(); + let body: Value = response.json().await.map_err(|e| { + Error::ExecutionErr(format!("Failed to parse the Entra ID token response: {e}")) + })?; + + if !status.is_success() { + return Err(Error::ExecutionErr(format!( + "Entra ID token request failed ({}): {} - {}", + status, + body["error"].as_str().unwrap_or("unknown"), + body["error_description"] + .as_str() + .unwrap_or("no description") + ))); + } + + let token = body["access_token"] + .as_str() + .ok_or_else(|| { + Error::ExecutionErr("Entra ID token response is missing access_token".to_string()) + })? + .to_string(); + let expires_in = body["expires_in"].as_u64().unwrap_or(3600); + + Ok(CachedToken { token, expires_at: Instant::now() + Duration::from_secs(expires_in) }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn config(authority_host: &str) -> WorkloadIdentityConfig { + WorkloadIdentityConfig { + tenant_id: "tenant".to_string(), + client_id: "client".to_string(), + federated_token_file: "/var/run/secrets/azure/tokens/azure-identity-token".to_string(), + authority_host: authority_host.to_string(), + } + } + + #[test] + fn test_token_endpoint() { + let expected = "https://login.microsoftonline.com/tenant/oauth2/v2.0/token"; + assert_eq!( + config("login.microsoftonline.com").token_endpoint(), + expected + ); + // The shape the workload identity webhook actually injects. + assert_eq!( + config("https://login.microsoftonline.com/").token_endpoint(), + expected + ); + } + + #[test] + fn test_cache_key_is_scoped() { + let config = config("login.microsoftonline.com"); + assert_ne!( + config.cache_key(AZURE_SQL_SCOPE), + config.cache_key(AZURE_OSSRDBMS_SCOPE) + ); + } + + /// A recent failure must not be re-attempted by every job queued behind the + /// refresh lock, but only while there is a token left to serve them. + #[test] + fn test_failed_refresh_is_not_retried_by_every_caller() { + assert!(!should_attempt_refresh(Some(Instant::now()), true)); + assert!(should_attempt_refresh(Some(Instant::now()), false)); + assert!(should_attempt_refresh( + Instant::now().checked_sub(REFRESH_RETRY_BACKOFF), + true + )); + assert!(should_attempt_refresh(None, true)); + } + + /// A token inside the refresh buffer is no longer served as fresh, but is still + /// good enough to fall back on when the renewal itself fails. + #[test] + fn test_token_within_refresh_buffer_is_stale_but_usable() { + let slot = TokenSlot::default(); + *slot.token.write().unwrap() = Some(CachedToken { + token: "tok".to_string(), + expires_at: Instant::now() + TOKEN_REFRESH_BUFFER - Duration::from_secs(60), + }); + + assert_eq!(slot.token_valid_for(TOKEN_REFRESH_BUFFER), None); + assert_eq!( + slot.token_valid_for(TOKEN_MIN_LIFETIME), + Some("tok".to_string()) + ); + } +} diff --git a/backend/windmill-common/src/lib.rs b/backend/windmill-common/src/lib.rs index 1789ddef1b..b0b7e720c7 100644 --- a/backend/windmill-common/src/lib.rs +++ b/backend/windmill-common/src/lib.rs @@ -29,6 +29,7 @@ use sqlx::{Acquire, Postgres}; pub mod agent_workers; pub mod apps; pub mod assets; +pub mod azure_workload_identity; pub mod dbt_manifest; pub mod audit; pub mod auth; @@ -1053,9 +1054,6 @@ impl PgDatabase { 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(), @@ -1075,7 +1073,41 @@ impl PgDatabase { error::Error::InternalErr(format!("IAM token generation failed: {e:#}")) })?; - // RDS IAM auth requires SSL. + self.connect_with_token("IAM RDS", &token).await + } + + /// Connect to Azure Database for PostgreSQL as the worker's federated identity. + /// The Entra ID access token replaces the password; `user` must be the Entra + /// principal name the server knows the identity by. + #[cfg(feature = "enterprise")] + pub async fn connect_with_workload_identity( + &self, + ) -> Result<(tokio_postgres::Client, TokioPgConnection), error::Error> { + let workload_identity = azure_workload_identity::WorkloadIdentityConfig::resolve()?; + let token = workload_identity + .access_token(azure_workload_identity::AZURE_OSSRDBMS_SCOPE) + .await?; + + self.connect_with_token("Azure workload identity", &token) + .await + } + + /// Connect with an externally issued access token in place of the password. + /// Both issuers (AWS IAM, Entra ID) mandate TLS, so encryption is forced on + /// regardless of the resource's sslmode; the sslmode still selects how far the + /// server's certificate is verified. + #[cfg(feature = "enterprise")] + async fn connect_with_token( + &self, + auth_kind: &str, + token: &str, + ) -> Result<(tokio_postgres::Client, TokioPgConnection), error::Error> { + use native_tls::TlsConnector; + use postgres_native_tls::MakeTlsConnector; + + let port = self.port.unwrap_or(5432); + let user = self.user.as_deref().unwrap_or("postgres"); + let mut connector = TlsConnector::builder(); let verified = Self::configure_pg_tls_verification( &mut connector, @@ -1084,19 +1116,19 @@ impl PgDatabase { self.accept_invalid_certs, )?; if !verified { - tracing::warn!("IAM RDS auth without certificate verification: TLS certificate verification is disabled. Provide root_certificate_pem (and set sslmode=verify-full) to enforce verification."); + tracing::warn!("{auth_kind} auth without certificate verification: TLS certificate verification is disabled. Provide root_certificate_pem (and set sslmode=verify-full) to enforce verification."); } - tracing::info!("Creating new IAM RDS connection to {}", &self.host); + tracing::info!("Creating new {auth_kind} connection to {}", &self.host); - // Use Config builder directly to pass the IAM token as the password. + // Use Config builder directly to pass the 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) + .password(token) .dbname(&self.dbname) .ssl_mode(tokio_postgres::config::SslMode::Require); diff --git a/backend/windmill-worker/src/mssql_executor.rs b/backend/windmill-worker/src/mssql_executor.rs index a5b4a06509..f3ffcfd103 100644 --- a/backend/windmill-worker/src/mssql_executor.rs +++ b/backend/windmill-worker/src/mssql_executor.rs @@ -11,6 +11,9 @@ use tiberius::{ use tokio::net::TcpStream; use tokio_util::compat::TokioAsyncWriteCompatExt; use uuid::Uuid; +use windmill_common::azure_workload_identity::{ + WorkloadIdentityConfig, AZURE_SQL_SCOPE, WORKLOAD_IDENTITY_PASSWORD, +}; use windmill_common::utils::merge_raw_values_to_object; use windmill_common::worker::SqlResultCollectionStrategy; use windmill_common::{ @@ -148,6 +151,15 @@ pub async fn do_mssql( "Integrated authentication is not available in this build. Requires mssql-kerberos (Linux) or mssql-winauth (Windows) feature.".to_string(), )); } + } else if database.password.as_deref() == Some(WORKLOAD_IDENTITY_PASSWORD) { + let workload_identity = WorkloadIdentityConfig::resolve()?; + let logs = format!( + "\nUsing Azure Workload Identity (client id {})", + workload_identity.client_id() + ); + append_logs(&job.id, &job.workspace_id, logs, conn).await; + let token = workload_identity.access_token(AZURE_SQL_SCOPE).await?; + config.authentication(AuthMethod::aad_token(token)); } else if let Some(token_value) = &database.aad_token { if let Some(token) = &token_value.token { config.authentication(AuthMethod::aad_token(token)); @@ -160,7 +172,7 @@ pub async fn do_mssql( config.authentication(AuthMethod::sql_server(user.clone(), password.clone())); } else { return Err(Error::BadRequest( - "No authentication method configured. Set integrated_auth, aad_token, or user/password.".to_string(), + "No authentication method configured. Set integrated_auth, aad_token, or user/password (password `ms_entraid` for Azure workload identity).".to_string(), )); } diff --git a/backend/windmill-worker/src/pg_executor.rs b/backend/windmill-worker/src/pg_executor.rs index 71e78f92fe..320b237d77 100644 --- a/backend/windmill-worker/src/pg_executor.rs +++ b/backend/windmill-worker/src/pg_executor.rs @@ -22,6 +22,7 @@ use tokio_postgres::{ Column, }; use uuid::Uuid; +use windmill_common::azure_workload_identity::WORKLOAD_IDENTITY_PASSWORD; use windmill_common::error::to_anyhow; use windmill_common::error::{self, Error}; use windmill_common::worker::{ @@ -65,24 +66,74 @@ pub async fn clear_pg_cache() { CONNECTION_COUNTER.write().await.clear(); } +/// How the connection authenticates, which also keys the connection cache: a +/// connection established under one mode must never be handed to a request asking +/// for another. +#[derive(Clone, Copy, Debug, PartialEq)] +enum PgAuthMode { + Password, + /// AWS RDS IAM. + Iam, + /// Azure Entra ID, via the worker's federated identity. + WorkloadIdentity, +} + +impl PgAuthMode { + fn of(database: &PgDatabase) -> error::Result { + let workload_identity = database.password.as_deref() == Some(WORKLOAD_IDENTITY_PASSWORD); + match (database.use_iam_auth == Some(true), workload_identity) { + (true, true) => Err(Error::BadRequest( + "IAM RDS authentication cannot use the Azure workload identity password" + .to_string(), + )), + (true, false) => Ok(PgAuthMode::Iam), + (false, true) => Ok(PgAuthMode::WorkloadIdentity), + (false, false) => Ok(PgAuthMode::Password), + } + } + + fn cache_key_segment(&self) -> &'static str { + match self { + // Workload identity needs no segment of its own: what selects it, the + // sentinel password, is already part of to_uri(). + PgAuthMode::Password | PgAuthMode::WorkloadIdentity => "", + PgAuthMode::Iam => "&iam=true", + } + } +} + async fn new_pg_connection( database: &PgDatabase, - _use_iam_auth: bool, + auth_mode: PgAuthMode, main_db: Option<&DB>, ) -> error::Result<(tokio_postgres::Client, tokio::task::JoinHandle<()>)> { - let (client, connection) = if _use_iam_auth { - #[cfg(all(feature = "enterprise", feature = "private"))] - { - database.connect_with_iam().await? + let (client, connection) = match auth_mode { + PgAuthMode::Iam => { + #[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(), + )); + } } - #[cfg(not(all(feature = "enterprise", feature = "private")))] - { - return Err(Error::ExecutionErr( - "IAM RDS authentication requires Windmill Enterprise Edition".to_string(), - )); + PgAuthMode::WorkloadIdentity => { + #[cfg(feature = "enterprise")] + { + database.connect_with_workload_identity().await? + } + #[cfg(not(feature = "enterprise"))] + { + return Err(Error::ExecutionErr( + "Azure workload identity authentication requires Windmill Enterprise Edition" + .to_string(), + )); + } } - } else { - database.connect(main_db).await? + PgAuthMode::Password => database.connect(main_db).await?, }; let handle = tokio::spawn(async move { if let Err(e) = connection.await { @@ -643,11 +694,24 @@ pub async fn do_postgresql( annotations.result_collection }; - let use_iam_auth = database.use_iam_auth == Some(true); + let auth_mode = PgAuthMode::of(&database)?; - // 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. + // The sentinel password is easy to set by accident, so say in the job logs which + // identity the query actually ran as rather than only in the worker logs. + if matches!(auth_mode, PgAuthMode::WorkloadIdentity) { + windmill_queue::append_logs( + &job.id, + &job.workspace_id, + "Using Azure Workload Identity\n", + conn, + ) + .await; + } + + // Include the auth mode in the cache key to distinguish connections to the same host + // authenticated differently. 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. // // to_uri() collapses require/verify-ca/verify-full to the same string, so the TLS verification // inputs are folded into the key separately. Without this a connection established under a @@ -664,11 +728,11 @@ pub async fn do_postgresql( // to_uri() already ends with `?sslmode=...`, so append further key segments // with `&` to keep database_string a well-formed URI (it is only ever a cache // key, but a malformed one would mislead anyone who later logs or parses it). - let database_string = if use_iam_auth { - format!("{}&iam=true&tls={tls_disc:x}", database.to_uri()) - } else { - format!("{}&tls={tls_disc:x}", database.to_uri()) - }; + let database_string = format!( + "{}{}&tls={tls_disc:x}", + database.to_uri(), + auth_mode.cache_key_segment() + ); let database_string_clone = database_string.clone(); let cached_client; @@ -748,18 +812,18 @@ pub async fn do_postgresql( } drop(guard); cached_client = None; - new_client = Some(new_pg_connection(&database, use_iam_auth, conn.as_sql()).await?); + new_client = Some(new_pg_connection(&database, auth_mode, conn.as_sql()).await?); } } else { // Release the lock before connecting so the post-query caching // code can re-acquire it. drop(guard); cached_client = None; - new_client = Some(new_pg_connection(&database, use_iam_auth, conn.as_sql()).await?); + new_client = Some(new_pg_connection(&database, auth_mode, conn.as_sql()).await?); } } else { cached_client = None; - new_client = Some(new_pg_connection(&database, use_iam_auth, conn.as_sql()).await?); + new_client = Some(new_pg_connection(&database, auth_mode, conn.as_sql()).await?); } let (mut sig, _) = parse_pgsql_sig_with_typed_schema(&query) @@ -2058,6 +2122,23 @@ impl FromSql<'_> for StringCollector { mod tests { use super::*; + /// The sentinel password is the whole opt-in: nothing else marks the resource, so a + /// resource carrying it must not fall through to password auth. + #[test] + fn test_workload_identity_password_selects_the_auth_mode() { + let db = |password: &str| { + PgDatabase::parse_uri(&format!("postgres://someuser:{password}@host:5432/db")).unwrap() + }; + assert_eq!( + PgAuthMode::of(&db(WORKLOAD_IDENTITY_PASSWORD)).unwrap(), + PgAuthMode::WorkloadIdentity + ); + assert_eq!( + PgAuthMode::of(&db("hunter2")).unwrap(), + PgAuthMode::Password + ); + } + #[test] fn test_map_s3object_jsonb_overflow() { let pg_err = Error::ExecutionErr(