diff --git a/backend/ee-repo-ref.txt b/backend/ee-repo-ref.txt index 07aa65d1f5..8d70a431af 100644 --- a/backend/ee-repo-ref.txt +++ b/backend/ee-repo-ref.txt @@ -1 +1 @@ -e08a87450627bef9013498e40ee93a47bedda7ee +0e001bab643e449b3310b0692dd3598ee0902ecc diff --git a/backend/src/db_connect.rs b/backend/src/db_connect.rs index abb3378efc..3860c18909 100644 --- a/backend/src/db_connect.rs +++ b/backend/src/db_connect.rs @@ -41,50 +41,69 @@ pub async fn connect_db( }; let pool = connect(database_url.clone(), max_connections, worker_mode).await?; + #[cfg(all(feature = "enterprise", feature = "private"))] - let pool2 = pool.clone(); - #[cfg(all(feature = "enterprise", feature = "private"))] - if let DatabaseUrl::IamRds(database_url) = database_url { - tokio::spawn(async move { - loop { - tokio::select! { - _ = killpill_rx.recv() => { - break; - } - _ = tokio::time::sleep(std::time::Duration::from_secs(10)) => { - let needs_refresh = { - let read_guard = database_url.read().await; - read_guard.needs_refresh() - }; - if needs_refresh { - let new_url = tokio::time::timeout(std::time::Duration::from_secs(10), get_database_url()).await; + { + let needs_token_refresh = matches!( + database_url, + DatabaseUrl::IamRds(_) | DatabaseUrl::EntraId(_) + ); + let label = match &database_url { + DatabaseUrl::IamRds(_) => "IAM RDS", + DatabaseUrl::EntraId(_) => "Entra ID", + DatabaseUrl::Static(_) => "", + }; + if needs_token_refresh { + let pool2 = pool.clone(); + let database_url2 = database_url.clone(); + tokio::spawn(async move { + loop { + tokio::select! { + _ = killpill_rx.recv() => { + break; + } + _ = tokio::time::sleep(std::time::Duration::from_secs(10)) => { + if !database_url2.needs_refresh().await { + continue; + } + let new_url = tokio::time::timeout( + std::time::Duration::from_secs(10), + get_database_url(), + ) + .await; match new_url { Ok(Ok(new_url)) => { match new_url.connect_options().await { Ok(connect_options) => { pool2.set_connect_options(connect_options); - tracing::info!("Refreshed IAM RDS URL successfully"); + tracing::info!("Refreshed {label} URL successfully"); } Err(e) => { - tracing::error!("Error getting IAM RDS connect options, retrying in 10s: {}", e); + tracing::error!( + "Error getting {label} connect options, retrying in 10s: {e}" + ); continue; } } } Ok(Err(e)) => { - tracing::error!("Error refreshing IAM RDS URL, trying again in 10s: {}", e); + tracing::error!( + "Error refreshing {label} URL, trying again in 10s: {e}" + ); continue; } Err(e) => { - tracing::error!("Timeout after 10s refreshing IAM RDS URL, trying again in 10 seconds: {}", e); + tracing::error!( + "Timeout after 10s refreshing {label} URL, trying again in 10s: {e}" + ); continue; } } } } } - } - }); + }); + } } Ok(pool) diff --git a/backend/windmill-common/src/db_params.rs b/backend/windmill-common/src/db_params.rs new file mode 100644 index 0000000000..d52700376d --- /dev/null +++ b/backend/windmill-common/src/db_params.rs @@ -0,0 +1,45 @@ +use anyhow::Result; + +/// Parsed database connection parameters, shared across DB auth providers (IAM RDS, Entra ID, etc.) +#[derive(Debug, Clone)] +pub struct DatabaseParams { + pub hostname: String, + pub port: u64, + pub username: String, + pub database: String, +} + +/// Extract database connection parameters from a PostgreSQL URL +pub fn extract_database_params(database_url: &str) -> Result { + let url = url::Url::parse(database_url) + .map_err(|e| anyhow::anyhow!("Failed to parse database URL: {}", e))?; + + let hostname = url + .host_str() + .ok_or_else(|| anyhow::anyhow!("Database URL missing hostname"))? + .to_string(); + + let port = url.port().unwrap_or(5432) as u64; + + let username = if url.username().is_empty() { + return Err(anyhow::anyhow!("Database URL missing username")); + } else { + urlencoding::decode(url.username())?.to_string() + }; + + let database = url + .path() + .trim_start_matches('/') + .split('/') + .next() + .filter(|s| !s.is_empty()) + .ok_or_else(|| anyhow::anyhow!("Database URL missing database name"))? + .to_string(); + + Ok(DatabaseParams { + hostname, + port, + username, + database: urlencoding::decode(&database)?.to_string(), + }) +} diff --git a/backend/windmill-common/src/lib.rs b/backend/windmill-common/src/lib.rs index 98fb7bdebf..70bf614596 100644 --- a/backend/windmill-common/src/lib.rs +++ b/backend/windmill-common/src/lib.rs @@ -43,7 +43,10 @@ pub mod cache; pub mod client; pub mod db; #[cfg(all(feature = "enterprise", feature = "private"))] +mod db_entra_ee; +#[cfg(all(feature = "enterprise", feature = "private"))] mod db_iam_ee; +pub mod db_params; #[cfg(feature = "private")] pub mod ee; pub mod ee_oss; @@ -633,12 +636,14 @@ impl PgDatabase { pub enum DatabaseUrl { #[cfg(all(feature = "enterprise", feature = "private"))] IamRds(std::sync::Arc>), + #[cfg(all(feature = "enterprise", feature = "private"))] + EntraId(std::sync::Arc>), Static(String), } impl DatabaseUrl { /// Get the database URL as a string. - /// Note: For IAM RDS, this returns the original URL (for metadata extraction). + /// For token-based auth, this returns the original URL (for metadata extraction). /// For actual database connections, use connect_options() instead. pub async fn as_str(&self) -> String { match self { @@ -647,13 +652,18 @@ impl DatabaseUrl { let guard = rds_url.read().await; guard.as_str().to_string() } + #[cfg(all(feature = "enterprise", feature = "private"))] + DatabaseUrl::EntraId(entra_url) => { + let guard = entra_url.read().await; + guard.as_str().to_string() + } DatabaseUrl::Static(url) => url.clone(), } } /// Get PgConnectOptions for this database URL. - /// For IAM RDS, this returns options built directly from the token to avoid double-encoding - /// issues with temporary credentials (IRSA/Pod Identity). + /// For token-based auth (IAM RDS, Entra ID), this returns options built directly from the + /// token to avoid double-encoding issues with temporary credentials. /// For static URLs, this parses the URL string. pub async fn connect_options(&self) -> Result { match self { @@ -662,6 +672,11 @@ impl DatabaseUrl { let guard = rds_url.read().await; Ok(guard.connect_options()) } + #[cfg(all(feature = "enterprise", feature = "private"))] + DatabaseUrl::EntraId(entra_url) => { + let guard = entra_url.read().await; + Ok(guard.connect_options()) + } DatabaseUrl::Static(url) => sqlx::postgres::PgConnectOptions::from_str(url) .map_err(|e| Error::InternalErr(format!("Failed to parse database URL: {}", e))), } @@ -671,9 +686,31 @@ impl DatabaseUrl { match self { #[cfg(all(feature = "enterprise", feature = "private"))] DatabaseUrl::IamRds(rds_url) => rds_url.write().await.refresh().await, + #[cfg(all(feature = "enterprise", feature = "private"))] + DatabaseUrl::EntraId(entra_url) => entra_url.write().await.refresh().await, DatabaseUrl::Static(_) => Ok(()), } } + + pub async fn needs_refresh(&self) -> bool { + match self { + #[cfg(all(feature = "enterprise", feature = "private"))] + DatabaseUrl::IamRds(rds_url) => rds_url.read().await.needs_refresh(), + #[cfg(all(feature = "enterprise", feature = "private"))] + DatabaseUrl::EntraId(entra_url) => entra_url.read().await.needs_refresh(), + DatabaseUrl::Static(_) => false, + } + } + + /// Double-checked refresh: read-lock to check, then write-lock to refresh if still needed. + pub async fn refresh_if_needed(&self) -> Result<(), Error> { + if self.needs_refresh().await { + self.refresh().await.map_err(|e| { + Error::InternalErr(format!("Failed to refresh database token: {}", e)) + })?; + } + Ok(()) + } } static DATABASE_URL_CACHE: tokio::sync::OnceCell = tokio::sync::OnceCell::const_new(); @@ -701,7 +738,9 @@ pub async fn get_database_url() -> Result { let parsed_url = url::Url::parse(&url)?; - if parsed_url.password().is_some_and(|x| x == "iamrds") { + let password = parsed_url.password().unwrap_or_default(); + + if password == "iamrds" { let region = var("AWS_REGION").map_err(|_| { Error::BadConfig( "AWS_REGION env var is required for IAM RDS authentication".to_string(), @@ -731,34 +770,68 @@ pub async fn get_database_url() -> Result { "IAM RDS authentication is not enabled in OSS mode".to_string(), )); } + } else if password == "entraid" { + let tenant_id = var("AZURE_TENANT_ID").map_err(|_| { + Error::BadConfig( + "AZURE_TENANT_ID env var is required for Entra ID authentication" + .to_string(), + ) + })?; + + tracing::info!( + "entraid mode detected, generating Entra ID URL for tenant: {tenant_id}" + ); + #[cfg(all(feature = "enterprise", feature = "private"))] + { + let client_id = var("AZURE_CLIENT_ID").map_err(|_| { + Error::BadConfig( + "AZURE_CLIENT_ID env var is required for Entra ID authentication" + .to_string(), + ) + })?; + let federated_token_file = + var("AZURE_FEDERATED_TOKEN_FILE").map_err(|_| { + Error::BadConfig( + "AZURE_FEDERATED_TOKEN_FILE env var is required for Entra ID authentication".to_string(), + ) + })?; + let authority_host = var("AZURE_AUTHORITY_HOST") + .unwrap_or_else(|_| "login.microsoftonline.com".to_string()); + + let entra_url = db_entra_ee::generate_database_url( + &url, + &tenant_id, + &client_id, + &federated_token_file, + &authority_host, + ) + .await + .map_err(|e| { + Error::InternalErr(format!( + "Failed to generate Entra ID database URL: {}", + e + )) + })?; + tracing::info!("Entra ID URL generated successfully"); + Ok::(DatabaseUrl::EntraId(std::sync::Arc::new( + tokio::sync::RwLock::new(entra_url), + ))) + } + + #[cfg(not(all(feature = "enterprise", feature = "private")))] + { + return Err(Error::BadConfig( + "Entra ID authentication is not enabled in OSS mode".to_string(), + )); + } } else { Ok::(DatabaseUrl::Static(url.to_string())) } }) .await?; - // Check if we need to refresh and do so if necessary - #[cfg(all(feature = "enterprise", feature = "private"))] - if let DatabaseUrl::IamRds(ref rds_url_lock) = database_url { - // Check if refresh is needed - let needs_refresh = { - let read_guard = rds_url_lock.read().await; - read_guard.needs_refresh() - }; + database_url.refresh_if_needed().await?; - // If refresh is needed, acquire write lock and refresh - if needs_refresh { - let mut write_guard = rds_url_lock.write().await; - // Double-check after acquiring write lock (another task might have refreshed) - if write_guard.needs_refresh() { - write_guard.refresh().await.map_err(|e| { - Error::InternalErr(format!("Failed to refresh IAM token: {}", e)) - })?; - } - } - } - - // Return the URL string Ok(database_url.clone()) }