From f2df6cbc9964629647d0fc37307aba878f5867db Mon Sep 17 00:00:00 2001
From: dieriba <dieriba.pro@gmail.com>
Date: Sat, 26 Jul 2025 15:25:25 +0200
Subject: [PATCH] improve: add retry logic for sqs oidc fetch credentials
 (#6275)

* improve

* update dependency

* push ee ref

---------

Co-authored-by: Ruben Fiszel <ruben@windmill.dev>
---
 backend/Cargo.lock                  |  1 +
 backend/ee-repo-ref.txt             |  2 +-
 backend/windmill-api/Cargo.toml     |  4 ++--
 backend/windmill-common/src/auth.rs | 28 +++++++++++-----------------
 4 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/backend/Cargo.lock b/backend/Cargo.lock
index bb4e1bf1db..edd2461c51 100644
--- a/backend/Cargo.lock
+++ b/backend/Cargo.lock
@@ -14750,6 +14750,7 @@ dependencies = [
  "aws-sdk-sqs",
  "aws-sdk-sts",
  "axum",
+ "backon",
  "base32",
  "base64 0.22.1",
  "byteorder",
diff --git a/backend/ee-repo-ref.txt b/backend/ee-repo-ref.txt
index d6ac3d82b7..fb1a717f9c 100644
--- a/backend/ee-repo-ref.txt
+++ b/backend/ee-repo-ref.txt
@@ -1 +1 @@
-9c3b54214685e6489646f28ba81e22f0d8a0cbfc
\ No newline at end of file
+98bae7801a0d21d6a1608cd150533cad365d84ce
\ No newline at end of file
diff --git a/backend/windmill-api/Cargo.toml b/backend/windmill-api/Cargo.toml
index 3e3b6d5752..6da7e9e411 100644
--- a/backend/windmill-api/Cargo.toml
+++ b/backend/windmill-api/Cargo.toml
@@ -32,7 +32,7 @@ http_trigger = ["dep:matchit", "dep:thiserror", "dep:sha1", "dep:constant_time_e
 static_frontend = ["dep:rust-embed"]
 postgres_trigger = ["dep:rust-postgres", "dep:pg_escape", "dep:byteorder", "dep:thiserror", "dep:rust_decimal", "dep:rust-postgres-native-tls"]
 mqtt_trigger = ["dep:thiserror", "dep:rumqttc"]
-sqs_trigger = ["dep:aws-sdk-sqs", "dep:thiserror", "dep:aws-config"]
+sqs_trigger = ["dep:aws-sdk-sqs", "dep:aws-sdk-sts", "dep:thiserror", "dep:aws-config", "dep:backon"]
 deno_core = ["dep:deno_core", "dep:deno_error"]
 gcp_trigger = ["dep:thiserror", "dep:google-cloud-pubsub", "dep:google-cloud-googleapis", "dep:tonic"]
 cloud = ["windmill-common/cloud"]
@@ -145,6 +145,6 @@ google-cloud-googleapis = { workspace = true , optional = true }
 tonic = { workspace = true, optional = true }
 deno_error = { workspace = true, optional = true }
 deno_core = { workspace = true, optional = true }
-
+backon = {workspace = true, optional = true}
 [build-dependencies]
 deno_core = { workspace = true, optional = true }
\ No newline at end of file
diff --git a/backend/windmill-common/src/auth.rs b/backend/windmill-common/src/auth.rs
index f4940793e4..a2de086e22 100644
--- a/backend/windmill-common/src/auth.rs
+++ b/backend/windmill-common/src/auth.rs
@@ -323,8 +323,6 @@ pub async fn create_token_for_owner(
 #[cfg(feature = "aws_auth")]
 pub mod aws {
 
-    use crate::error::to_anyhow;
-
     use super::*;
     use crate::utils::empty_as_none;
     use aws_config::{BehaviorVersion, Region};
@@ -332,7 +330,9 @@ pub mod aws {
         config::Credentials as AwsCredentials,
         operation::{
             assume_role_with_saml::AssumeRoleWithSamlOutput,
-            assume_role_with_web_identity::AssumeRoleWithWebIdentityOutput,
+            assume_role_with_web_identity::{
+                builders::AssumeRoleWithWebIdentityFluentBuilder, AssumeRoleWithWebIdentityOutput,
+            },
         },
         types::Credentials,
         Client,
@@ -396,33 +396,27 @@ pub mod aws {
         Oidc(OidcAuth),
     }
 
-    pub async fn get_oidc_authentication_data(
-        oidc_auth: OidcAuth,
-        role_session_name: Option<impl ToString>,
+    pub async fn get_assume_role_with_web_identity_fluent_builder(
+        oidc_auth: &OidcAuth,
         token: String,
-    ) -> Result<AssumeRoleWithWebIdentityOutput> {
-        let region = oidc_auth.region.unwrap_or_else(|| "us-east-1".to_string());
+        role_session_name: Option<impl ToString>,
+    ) -> Result<AssumeRoleWithWebIdentityFluentBuilder> {
+        let region = oidc_auth.region.as_deref().unwrap_or_else(|| "us-east-1");
 
         let credentials = AwsCredentials::new("", "", None, None, "UserInput");
 
         let config = aws_config::defaults(BehaviorVersion::latest())
             .credentials_provider(credentials)
-            .region(Region::new(region.clone()))
+            .region(Region::new(region.to_string()))
             .load()
             .await;
 
         let assume_role_with_web_identity_fluent_builder = Client::new(&config)
             .assume_role_with_web_identity()
-            .set_role_arn(Some(oidc_auth.role_arn))
+            .set_role_arn(Some(oidc_auth.role_arn.to_owned()))
             .set_role_session_name(role_session_name.map(|str| str.to_string()))
             .set_web_identity_token(Some(token));
 
-        let resp = assume_role_with_web_identity_fluent_builder
-            .clone()
-            .send()
-            .await
-            .map_err(to_anyhow)?;
-
-        Ok(resp)
+        Ok(assume_role_with_web_identity_fluent_builder)
     }
 }