fix cache and threading concerns with STS Webidentity provider

Author: sbiscigl

src/aws-cpp-sdk-core/include/aws/core/auth/STSCredentialsProvider.h

@@ -8,11 +8,14 @@
 
 #include <aws/core/Core_EXPORTS.h>
 #include <aws/core/auth/AWSCredentialsProvider.h>
+#include <atomic>
+#include <memory>
 
 namespace Aws {
 namespace Crt {
 namespace Auth {
 class ICredentialsProvider;
+class Credentials;
 }
 }
 }
@@ -46,10 +49,19 @@ namespace Aws
               INITIALIZED,
               SHUT_DOWN,
             } m_state{STATE::SHUT_DOWN};
-            std::mutex m_refreshMutex;
-            std::condition_variable m_refreshSignal;
+            mutable std::mutex m_refreshMutex;
+            mutable std::condition_variable m_refreshSignal;
             std::shared_ptr<Aws::Crt::Auth::ICredentialsProvider> m_credentialsProvider;
             std::chrono::milliseconds m_providerFuturesTimeoutMs;
+
+            // Thread-safe credential fetch coordination
+            mutable std::atomic<bool> m_refreshInProgress{false};
+            mutable std::shared_ptr<AWSCredentials> m_pendingCredentials;
+
+            // Helper methods for credential retrieval
+            AWSCredentials waitForSharedCredentials() const;
+            AWSCredentials extractCredentialsFromCrt(const Aws::Crt::Auth::Credentials& crtCredentials) const;
+            AWSCredentials fetchCredentialsAsync();
         };
     } // namespace Auth
 } // namespace Aws

src/aws-cpp-sdk-core/include/aws/core/client/ClientConfiguration.h

@@ -544,6 +544,11 @@ namespace Aws
                * Time out for the credentials future call.
                */
               std::chrono::milliseconds retrieveCredentialsFutureTimeout = std::chrono::seconds(10);
+
+              /**
+               * How long a cached credential set will be used for
+               */
+              std::chrono::milliseconds credentialCacheCacheTTL = std::chrono::minutes(50);
             } stsCredentialsProviderConfig;
           } credentialProviderConfig;
         };

src/aws-cpp-sdk-core/source/auth/STSCredentialsProvider.cpp

@@ -37,11 +37,27 @@ STSAssumeRoleWebIdentityCredentialsProvider::STSAssumeRoleWebIdentityCredentials
     }
     return UUID::RandomUUID();
   }().c_str();
-  m_credentialsProvider = Aws::Crt::Auth::CredentialsProvider::CreateCredentialsProviderSTSWebIdentity(stsConfig);
+
+  // Create underlying STS provider
+  auto stsProvider = Aws::Crt::Auth::CredentialsProvider::CreateCredentialsProviderSTSWebIdentity(stsConfig);
+  if (!stsProvider || !stsProvider->IsValid()) {
+    AWS_LOGSTREAM_WARN(STS_LOG_TAG, "Failed to create underlying STS credentials provider");
+    return;
+  }
+
+  // Wrap with caching provider
+  Aws::Crt::Auth::CredentialsProviderCachedConfig cachedConfig;
+  cachedConfig.Provider = stsProvider;
+  cachedConfig.CachedCredentialTTL = credentialsConfig.stsCredentialsProviderConfig.credentialCacheCacheTTL;
+
+  m_credentialsProvider = Aws::Crt::Auth::CredentialsProvider::CreateCredentialsProviderCached(cachedConfig);
   if (m_credentialsProvider && m_credentialsProvider->IsValid()) {
     m_state = STATE::INITIALIZED;
+    AWS_LOGSTREAM_INFO(STS_LOG_TAG, "STS credentials provider initialized with cache TTL "
+      << cachedConfig.CachedCredentialTTL.count()
+      << " ms");
   } else {
-    AWS_LOGSTREAM_WARN(STS_LOG_TAG, "Failed to create STS credentials provider");
+    AWS_LOGSTREAM_WARN(STS_LOG_TAG, "Failed to create cached STS credentials provider");
   }
 }
 
@@ -77,34 +93,18 @@ STSAssumeRoleWebIdentityCredentialsProvider::~STSAssumeRoleWebIdentityCredential
 
 AWSCredentials STSAssumeRoleWebIdentityCredentialsProvider::GetAWSCredentials() {
   if (m_state != STATE::INITIALIZED) {
-    AWS_LOGSTREAM_DEBUG(STS_LOG_TAG, "STSCredentialsProvider is not initialized, returning empty credentials");
+    AWS_LOGSTREAM_WARN(STS_LOG_TAG, "STSCredentialsProvider is not initialized, returning empty credentials");
     return AWSCredentials{};
   }
-  AWSCredentials credentials{};
-  auto refreshDone = false;
-  m_credentialsProvider->GetCredentials(
-      [this, &credentials, &refreshDone](std::shared_ptr<Aws::Crt::Auth::Credentials> crtCredentials, int errorCode) -> void {
-        {
-          const std::unique_lock<std::mutex> lock{m_refreshMutex};
-          if (errorCode != AWS_ERROR_SUCCESS) {
-            AWS_LOGSTREAM_ERROR(STS_LOG_TAG, "Failed to get credentials from STS: " << errorCode);
-          } else {
-            const auto accountIdCursor = crtCredentials->GetAccessKeyId();
-            credentials.SetAWSAccessKeyId({reinterpret_cast<char*>(accountIdCursor.ptr), accountIdCursor.len});
-            const auto secretKeuCursor = crtCredentials->GetSecretAccessKey();
-            credentials.SetAWSSecretKey({reinterpret_cast<char*>(secretKeuCursor.ptr), secretKeuCursor.len});
-            const auto expiration = crtCredentials->GetExpirationTimepointInSeconds();
-            credentials.SetExpiration(DateTime{static_cast<double>(expiration)});
-            const auto sessionTokenCursor = crtCredentials->GetSessionToken();
-            credentials.SetSessionToken({reinterpret_cast<char*>(sessionTokenCursor.ptr), sessionTokenCursor.len});
-          }
-          refreshDone = true;
-        }
-        m_refreshSignal.notify_one();
-      });
 
-  std::unique_lock<std::mutex> lock{m_refreshMutex};
-  m_refreshSignal.wait_for(lock, m_providerFuturesTimeoutMs, [&refreshDone]() -> bool { return refreshDone; });
+  // Thread-safe check: If another thread is already fetching, wait for its result
+  auto expected = false;
+  if (!m_refreshInProgress.compare_exchange_strong(expected, true)) {
+    return waitForSharedCredentials();
+  }
+
+  // This thread will fetch the credentials
+  auto credentials = fetchCredentialsAsync();
 
   if (!credentials.IsEmpty()) {
     credentials.AddUserAgentFeature(Aws::Client::UserAgentFeature::CREDENTIALS_STS_WEB_IDENTITY_TOKEN);
@@ -116,3 +116,71 @@ AWSCredentials STSAssumeRoleWebIdentityCredentialsProvider::GetAWSCredentials()
 void STSAssumeRoleWebIdentityCredentialsProvider::Reload() {
   AWS_LOGSTREAM_DEBUG(STS_LOG_TAG, "Calling reload on STSCredentialsProvider is a no-op and no longer in the call path");
 }
+
+AWSCredentials STSAssumeRoleWebIdentityCredentialsProvider::waitForSharedCredentials() const {
+  AWS_LOGSTREAM_DEBUG(STS_LOG_TAG, "Another thread is fetching credentials, waiting for result");
+  std::unique_lock<std::mutex> lock{m_refreshMutex};
+  m_refreshSignal.wait_for(lock, m_providerFuturesTimeoutMs, [this]() -> bool {
+    return !m_refreshInProgress.load();
+  });
+
+  if (m_pendingCredentials) {
+    return *m_pendingCredentials;
+  }
+
+  AWS_LOGSTREAM_WARN(STS_LOG_TAG, "Failed to get shared credentials after timeout");
+  return AWSCredentials{};
+}
+
+AWSCredentials STSAssumeRoleWebIdentityCredentialsProvider::extractCredentialsFromCrt(
+    const Aws::Crt::Auth::Credentials& crtCredentials) const {
+  AWSCredentials credentials{};
+  const auto accountIdCursor = crtCredentials.GetAccessKeyId();
+  credentials.SetAWSAccessKeyId({reinterpret_cast<char*>(accountIdCursor.ptr), accountIdCursor.len});
+  const auto secretKeyCursor = crtCredentials.GetSecretAccessKey();
+  credentials.SetAWSSecretKey({reinterpret_cast<char*>(secretKeyCursor.ptr), secretKeyCursor.len});
+  const auto expiration = crtCredentials.GetExpirationTimepointInSeconds();
+  credentials.SetExpiration(DateTime{static_cast<double>(expiration)});
+  const auto sessionTokenCursor = crtCredentials.GetSessionToken();
+  credentials.SetSessionToken({reinterpret_cast<char*>(sessionTokenCursor.ptr), sessionTokenCursor.len});
+  return credentials;
+}
+
+AWSCredentials STSAssumeRoleWebIdentityCredentialsProvider::fetchCredentialsAsync() {
+  AWS_LOGSTREAM_DEBUG(STS_LOG_TAG, "Initiating credential fetch from STS/cache");
+
+  AWSCredentials credentials{};
+  std::atomic<bool> refreshDone{false};
+
+  m_credentialsProvider->GetCredentials(
+      [this, &credentials, &refreshDone](
+          std::shared_ptr<Aws::Crt::Auth::Credentials> crtCredentials, int errorCode) -> void {
+
+        std::unique_lock<std::mutex> lock{m_refreshMutex};
+        if (errorCode != AWS_ERROR_SUCCESS) {
+          m_pendingCredentials.reset();
+        } else {
+          credentials = extractCredentialsFromCrt(*crtCredentials);
+
+          // Store for other waiting threads
+          m_pendingCredentials = std::make_shared<AWSCredentials>(credentials);
+        }
+        refreshDone.store(true);
+        m_refreshInProgress.store(false);
+        m_refreshSignal.notify_all();
+      });
+
+      // Wait for completion
+      std::unique_lock<std::mutex> lock{m_refreshMutex};
+      auto completed = m_refreshSignal.wait_for(lock, m_providerFuturesTimeoutMs, [&refreshDone]() -> bool {
+        return refreshDone.load();
+      });
+
+  if (!completed) {
+    AWS_LOGSTREAM_ERROR(STS_LOG_TAG, "Credential fetch timed out after " << m_providerFuturesTimeoutMs.count() << "ms");
+    m_refreshInProgress.store(false);
+    m_refreshSignal.notify_all();
+  }
+
+  return credentials;
+}