feat: Add callback-based credential provider tracking system

kai-ion · kai-ion · commit 2e36fccd20e0 · 2025-08-14T14:44:22.000-04:00
Implement lightweight callback mechanism for tracking credential provider usage in User-Agent strings without breaking backward compatibility.

Changes:
- Add SetCredentialTrackingCallback() and NotifyCredentialUsage() to AWSCredentialsProvider base class
- Implement callback forwarding in AWSCredentialsProviderChain for provider chains
- Add TrackCredentialProviderUsage() method to AWSClient for setting up tracking callbacks
- Update EnvironmentAWSCredentialsProvider to call NotifyCredentialUsage() when credentials found
- Add unit test for environment credential tracking (currently failing due to GetCredentialsProvider() limitation)

The callback system allows individual credential providers to notify when they successfully retrieve credentials, enabling User-Agent tracking without requiring API changes or thread-local variables.

Note: Current implementation has limitation where GetCredentialsProvider() returns null when provider is embedded in signer, preventing callback setup in some scenarios.
diff --git a/src/aws-cpp-sdk-core/include/aws/core/auth/AWSCredentialsProvider.h b/src/aws-cpp-sdk-core/include/aws/core/auth/AWSCredentialsProvider.h
@@ -19,9 +19,12 @@
 #include <aws/core/config/AWSProfileConfigLoader.h>
 #include <aws/core/client/RetryStrategy.h>
 #include <memory>
+#include <functional>
 
 namespace Aws
 {
+    class AmazonWebServiceRequest;
+    
     namespace Client
     {
         struct ClientConfiguration;
@@ -32,17 +35,6 @@ namespace Aws
 
         constexpr int AWS_CREDENTIAL_PROVIDER_EXPIRATION_GRACE_PERIOD = 5 * 1000;
 
-        /**
-         * Enum to identify credential provider types for tracking purposes
-         */
-        enum class CredentialProviderType
-        {
-            DEFAULT,
-            ENVIRONMENT,
-            // ... add other types as needed
-
-        };
-
         /**
          * Returns the full path of the config file.
          */
@@ -73,22 +65,27 @@ namespace Aws
              * Initializes provider. Sets last Loaded time count to 0, forcing a refresh on the
              * first call to GetAWSCredentials.
              */
-            AWSCredentialsProvider(CredentialProviderType providerType = CredentialProviderType::DEFAULT)
-                : m_lastLoadedMs(0), m_providerType(providerType)
+            AWSCredentialsProvider() : m_lastLoadedMs(0)
             {
             }
 
-            /**
-             * Get the provider type for tracking purposes
-             */
-            CredentialProviderType GetProviderType() const { return m_providerType; }
-
             virtual ~AWSCredentialsProvider() = default;
 
             /**
              * The core of the credential provider interface. Override this method to control how credentials are retrieved.
              */
             virtual AWSCredentials GetAWSCredentials() = 0;
+            
+            /**
+             * Set callback for credential usage tracking
+             */
+            virtual void SetCredentialTrackingCallback(std::function<void()> callback) { m_trackingCallback = callback; }
+            
+        protected:
+            /**
+             * Call this when credentials are successfully retrieved for tracking
+             */
+            void NotifyCredentialUsage() { if (m_trackingCallback) m_trackingCallback(); }
 
         protected:
             /**
@@ -100,7 +97,7 @@ namespace Aws
             mutable Aws::Utils::Threading::ReaderWriterLock m_reloadLock;
         private:
             long long m_lastLoadedMs;
-            CredentialProviderType m_providerType;
+            std::function<void()> m_trackingCallback;
         };
 
         /**
@@ -114,7 +111,6 @@ namespace Aws
              * Returns empty credentials object.
              */
             inline AWSCredentials GetAWSCredentials() override { return AWSCredentials(); }
-
         };
 
         /**
@@ -161,7 +157,7 @@ namespace Aws
             /**
             * Initializes environment credentials provider
             */
-            EnvironmentAWSCredentialsProvider() : AWSCredentialsProvider(CredentialProviderType::ENVIRONMENT) {}
+            EnvironmentAWSCredentialsProvider() = default;
             
             /**
             * Reads AWS credentials from the Environment variables AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY and AWS_SESSION_TOKEN if they exist. If they
diff --git a/src/aws-cpp-sdk-core/include/aws/core/auth/AWSCredentialsProviderChain.h b/src/aws-cpp-sdk-core/include/aws/core/auth/AWSCredentialsProviderChain.h
@@ -27,7 +27,12 @@ namespace Aws
              * When a credentials provider in the chain returns empty credentials,
              * We go on to the next provider until we have either exhausted the installed providers in the chain or something returns non-empty credentials.
              */
-            virtual AWSCredentials GetAWSCredentials();
+            AWSCredentials GetAWSCredentials() override;
+            
+            /**
+             * Override to store and forward callback to providers in chain
+             */
+            void SetCredentialTrackingCallback(std::function<void()> callback) override { m_chainTrackingCallback = callback; }
 
             /**
              * Gets all providers stored in this chain.
@@ -50,6 +55,7 @@ namespace Aws
             Aws::Vector<std::shared_ptr<AWSCredentialsProvider> > m_providerChain;
             std::shared_ptr<AWSCredentialsProvider> m_cachedProvider;
             mutable Aws::Utils::Threading::ReaderWriterLock m_cachedProviderLock;
+            std::function<void()> m_chainTrackingCallback;
         };
 
         /**
diff --git a/src/aws-cpp-sdk-core/source/auth/AWSCredentialsProvider.cpp b/src/aws-cpp-sdk-core/source/auth/AWSCredentialsProvider.cpp
@@ -18,6 +18,8 @@
 #include <aws/core/client/AWSError.h>
 #include <aws/core/utils/StringUtils.h>
 #include <aws/core/utils/xml/XmlSerializer.h>
+#include <aws/core/AmazonWebServiceRequest.h>
+#include <aws/core/client/UserAgent.h>
 #include <cstdlib>
 #include <fstream>
 #include <string.h>
@@ -102,11 +104,12 @@ AWSCredentials EnvironmentAWSCredentialsProvider::GetAWSCredentials() //pass in
             AWS_LOGSTREAM_DEBUG(ENVIRONMENT_LOG_TAG, "Found accountId");
         }
     }
-
+    
     if (!credentials.IsEmpty()) {
-      // TODO: this will work
-      // TODO: how does request get here?????
-      // request.AddFeature(ENV_VAR)
+        // TODO: this will work
+        // TODO: how does request get here?????
+        // request.AddFeature(ENV_VAR)
+        NotifyCredentialUsage();
     }
 
     return credentials;
diff --git a/src/aws-cpp-sdk-core/source/auth/AWSCredentialsProviderChain.cpp b/src/aws-cpp-sdk-core/source/auth/AWSCredentialsProviderChain.cpp
@@ -22,6 +22,8 @@ AWSCredentials AWSCredentialsProviderChain::GetAWSCredentials()
 {
     ReaderLockGuard lock(m_cachedProviderLock);
     if (m_cachedProvider) {
+      // Forward callback to cached provider
+      m_cachedProvider->SetCredentialTrackingCallback(m_chainTrackingCallback);
       AWSCredentials credentials = m_cachedProvider->GetAWSCredentials();
       if (!credentials.GetAWSAccessKeyId().empty() && !credentials.GetAWSSecretKey().empty())
       {
@@ -31,6 +33,8 @@ AWSCredentials AWSCredentialsProviderChain::GetAWSCredentials()
     lock.UpgradeToWriterLock();
     for (auto&& credentialsProvider : m_providerChain)
     {
+        // Forward callback to each provider in chain
+        credentialsProvider->SetCredentialTrackingCallback(m_chainTrackingCallback);
         AWSCredentials credentials = credentialsProvider->GetAWSCredentials();
         if (!credentials.GetAWSAccessKeyId().empty() && !credentials.GetAWSSecretKey().empty())
         {
diff --git a/src/aws-cpp-sdk-core/source/client/AWSClient.cpp b/src/aws-cpp-sdk-core/source/client/AWSClient.cpp
@@ -1075,25 +1075,19 @@ void AWSClient::TrackCredentialProviderUsage(const Aws::AmazonWebServiceRequest&
         return;
     }
 
-    // Get the provider type
-    auto providerType = credentialsProvider->GetProviderType();
-
-    switch (providerType)
-    {
-        case Aws::Auth::CredentialProviderType::ENVIRONMENT:
-            // Environment credentials are being used
+    // Set up callback for credential tracking
+    if (credentialsProvider) {
+        credentialsProvider->SetCredentialTrackingCallback([&request]() {
             request.AddUserAgentFeature(Aws::Client::UserAgentFeature::CREDENTIALS_ENV_VARS);
             AWS_LOGSTREAM_DEBUG(AWS_CLIENT_LOG_TAG, "Added CREDENTIALS_ENV_VARS to User-Agent");
-            break;
+        });
+    }
+    
+    // Trigger credential retrieval to enable tracking
+    if (credentialsProvider) {
+        credentialsProvider->GetAWSCredentials();
         
-        // Add more provider types as needed
-        default:
-            // For provider chains or unknown types, check environment variables as fallback
-            if (!Aws::Environment::GetEnv("AWS_ACCESS_KEY_ID").empty())
-            {
-                request.AddUserAgentFeature(Aws::Client::UserAgentFeature::CREDENTIALS_ENV_VARS);
-                AWS_LOGSTREAM_DEBUG(AWS_CLIENT_LOG_TAG, "Added CREDENTIALS_ENV_VARS to User-Agent (fallback detection)");
-            }
-            break;
+        // Clear callback
+        credentialsProvider->SetCredentialTrackingCallback(nullptr);
     }
 }
diff --git a/tests/aws-cpp-sdk-core-tests/aws/auth/CredentialTrackingTest.cpp b/tests/aws-cpp-sdk-core-tests/aws/auth/CredentialTrackingTest.cpp
@@ -8,16 +8,50 @@
 #include <aws/testing/mocks/aws/client/MockAWSClient.h>
 #include <aws/testing/mocks/http/MockHttpClient.h>
 #include <aws/core/auth/AWSCredentialsProvider.h>
+#include <aws/core/auth/AWSCredentialsProviderChain.h>
 #include <aws/core/client/ClientConfiguration.h>
+#include <aws/core/client/AWSClient.h>
+#include <aws/core/auth/AWSAuthSigner.h>
 #include <aws/core/platform/Environment.h>
 #include <aws/core/utils/StringUtils.h>
+#include <iostream>
 
 using namespace Aws::Client;
 using namespace Aws::Auth;
 using namespace Aws::Http;
 
 static const char ALLOCATION_TAG[] = "CredentialTrackingTest";
 
+// Custom client that uses default credential provider chain for testing
+class CredentialTestingClient : public Aws::Client::AWSClient
+{
+public:
+    explicit CredentialTestingClient(const Aws::Client::ClientConfiguration& configuration)
+        : AWSClient(configuration,
+                   Aws::MakeShared<Aws::Client::AWSAuthV4Signer>(ALLOCATION_TAG,
+                       Aws::MakeShared<DefaultAWSCredentialsProviderChain>(ALLOCATION_TAG),
+                       "service", configuration.region),
+                   Aws::MakeShared<MockAWSErrorMarshaller>(ALLOCATION_TAG))
+    {
+        // Client created with DefaultAWSCredentialsProviderChain
+    }
+
+    Aws::Client::HttpResponseOutcome MakeRequest(const Aws::AmazonWebServiceRequest& request)
+    {
+        auto uri = Aws::Http::URI("https://test.com");
+        return AWSClient::AttemptExhaustively(uri, request, Aws::Http::HttpMethod::HTTP_POST, Aws::Auth::SIGV4_SIGNER);
+    }
+
+    const char* GetServiceClientName() const override { return "CredentialTestingClient"; }
+
+protected:
+    Aws::Client::AWSError<Aws::Client::CoreErrors> BuildAWSError(const std::shared_ptr<Aws::Http::HttpResponse>& response) const override
+    {
+        AWS_UNREFERENCED_PARAM(response);
+        return Aws::Client::AWSError<Aws::Client::CoreErrors>(Aws::Client::CoreErrors::UNKNOWN, false);
+    }
+};
+
 class CredentialTrackingTest : public Aws::Testing::AwsCppSdkGTestSuite
 {
 protected:
@@ -49,33 +83,35 @@ TEST_F(CredentialTrackingTest, TestEnvironmentCredentialsTracking)
     Aws::Environment::SetEnv("AWS_SECRET_ACCESS_KEY", "test-secret-key", 1);
 
     // Setup mock response
-    auto request = CreateHttpRequest(Aws::Http::URI("http://test.com"), 
-                                   Aws::Http::HttpMethod::HTTP_POST, 
-                                   Aws::Utils::Stream::DefaultResponseStreamFactoryMethod);
-    auto response = Aws::MakeShared<Standard::StandardHttpResponse>(ALLOCATION_TAG, request);
-    response->SetResponseCode(HttpResponseCode::OK);
-    response->GetResponseBody() << "{}";
-    mockHttpClient->AddResponseToReturn(response);
+    std::shared_ptr<HttpRequest> requestTmp =
+        CreateHttpRequest(Aws::Http::URI("dummy"), Aws::Http::HttpMethod::HTTP_POST, 
+                        Aws::Utils::Stream::DefaultResponseStreamFactoryMethod);
+    auto successResponse = Aws::MakeShared<Standard::StandardHttpResponse>(ALLOCATION_TAG, requestTmp);
+    successResponse->SetResponseCode(HttpResponseCode::OK);
+    successResponse->GetResponseBody() << "{}";
+    mockHttpClient->AddResponseToReturn(successResponse);
 
     // Create client configuration
-    ClientConfiguration config;
-    config.region = Aws::Region::US_EAST_1;
+    Aws::Client::ClientConfigurationInitValues cfgInit;
+    cfgInit.shouldDisableIMDS = true;
+    Aws::Client::ClientConfiguration clientConfig(cfgInit);
+    clientConfig.region = Aws::Region::US_EAST_1;
 
-    // Create mock client
-    MockAWSClient client(config);
+    // Create credential testing client that uses default provider chain
+    CredentialTestingClient client(clientConfig);
 
-    // Make a request
+    // Create mock request
     AmazonWebServiceRequestMock mockRequest;
-    auto outcome = client.MakeRequest(mockRequest);
 
-    // Verify request succeeded
-    AWS_ASSERT_SUCCESS(outcome);
+    // Make request
+    auto outcome = client.MakeRequest(mockRequest);
+    ASSERT_TRUE(outcome.IsSuccess());
 
     // Verify User-Agent contains environment credentials tracking
     auto lastRequest = mockHttpClient->GetMostRecentHttpRequest();
-    EXPECT_TRUE(lastRequest.HasUserAgent());
-    const auto& userAgent = lastRequest.GetUserAgent();
-    EXPECT_TRUE(!userAgent.empty());
+    EXPECT_TRUE(lastRequest.HasHeader(Aws::Http::USER_AGENT_HEADER));
+    const auto& userAgent = lastRequest.GetHeaderValue(Aws::Http::USER_AGENT_HEADER);
+    EXPECT_FALSE(userAgent.empty());
 
     const auto userAgentParsed = Aws::Utils::StringUtils::Split(userAgent, ' ');
 
@@ -89,12 +125,3 @@ TEST_F(CredentialTrackingTest, TestEnvironmentCredentialsTracking)
     Aws::Environment::UnSetEnv("AWS_ACCESS_KEY_ID");
     Aws::Environment::UnSetEnv("AWS_SECRET_ACCESS_KEY");
 }
-
-TEST_F(CredentialTrackingTest, TestEnvironmentProviderType)
-{
-    // Test that EnvironmentAWSCredentialsProvider has correct provider type
-    auto envProvider = Aws::MakeShared<EnvironmentAWSCredentialsProvider>(ALLOCATION_TAG);
-    
-    // Verify the provider type is set correctly
-    EXPECT_EQ(envProvider->GetProviderType(), Aws::Auth::CredentialProviderType::ENVIRONMENT);
-}

Original file line number	Diff line number	Diff line change
`@@ -22,6 +22,8 @@ AWSCredentials AWSCredentialsProviderChain::GetAWSCredentials()`
`22`	`22`	`{`
`23`	`23`	`ReaderLockGuard lock(m_cachedProviderLock);`
`24`	`24`	`if (m_cachedProvider) {`
	`25`	`+ // Forward callback to cached provider`
	`26`	`+ m_cachedProvider->SetCredentialTrackingCallback(m_chainTrackingCallback);`
`25`	`27`	`AWSCredentials credentials = m_cachedProvider->GetAWSCredentials();`
`26`	`28`	`if (!credentials.GetAWSAccessKeyId().empty() && !credentials.GetAWSSecretKey().empty())`
`27`	`29`	`{`
`@@ -31,6 +33,8 @@ AWSCredentials AWSCredentialsProviderChain::GetAWSCredentials()`
`31`	`33`	`lock.UpgradeToWriterLock();`
`32`	`34`	`for (auto&& credentialsProvider : m_providerChain)`
`33`	`35`	`{`
	`36`	`+ // Forward callback to each provider in chain`
	`37`	`+ credentialsProvider->SetCredentialTrackingCallback(m_chainTrackingCallback);`
`34`	`38`	`AWSCredentials credentials = credentialsProvider->GetAWSCredentials();`
`35`	`39`	`if (!credentials.GetAWSAccessKeyId().empty() && !credentials.GetAWSSecretKey().empty())`
`36`	`40`	`{`