wire smithy client to use account id resolved from identity for endpoint paramters

sbiscigl · sbiscigl · commit a241c97925fb · 2025-03-17T10:12:04.000-04:00
diff --git a/src/aws-cpp-sdk-core/include/smithy/client/AwsSmithyClient.h b/src/aws-cpp-sdk-core/include/smithy/client/AwsSmithyClient.h
@@ -20,11 +20,9 @@
 #include <aws/core/utils/FutureOutcome.h>
 #include <aws/core/utils/Outcome.h>
 #include <aws/core/utils/threading/Executor.h>
-#include <aws/core/utils/threading/SameThreadExecutor.h>
 #include <aws/core/http/HttpResponse.h>
 #include <aws/core/http/HttpClientFactory.h>
 #include <smithy/identity/signer/built-in/SignerProperties.h>
-#include <smithy/identity/auth/built-in/SigV4AuthSchemeOption.h>
 #include <smithy/client/AwsLegacyClient.h>
 
 namespace smithy {
@@ -204,16 +202,24 @@ namespace client
                                  false/*retryable*/);
         }
 
-        SigningOutcome SignHttpRequest(std::shared_ptr<HttpRequest> httpRequest, const AuthSchemeOption& targetAuthSchemeOption) const override
+        SigningOutcome SignHttpRequest(std::shared_ptr<HttpRequest> httpRequest, const AwsSmithyClientAsyncRequestContext& ctx) const override
         {
-            return AwsClientRequestSigning<AuthSchemesVariantT>::SignRequest(httpRequest, targetAuthSchemeOption, m_authSchemes);
+            return AwsClientRequestSigning<AuthSchemesVariantT>::SignRequest(httpRequest, ctx, m_authSchemes);
         }
 
         bool AdjustClockSkew(HttpResponseOutcome& outcome, const AuthSchemeOption& authSchemeOption) const override
         {
             return AwsClientRequestSigning<AuthSchemesVariantT>::AdjustClockSkew(outcome, authSchemeOption, m_authSchemes);
         }
 
+        IdentityOutcome ResolveIdentity(const AwsSmithyClientAsyncRequestContext& ctx) const override {
+          return AwsClientRequestSigning<AuthSchemesVariantT>::ResolveIdentity(ctx, m_authSchemes);
+        }
+
+        GetContextEndpointParametersOutcome GetContextEndpointParameters(const AwsSmithyClientAsyncRequestContext& ctx) const override {
+          return GetContextEndpointParametersImpl(ctx);
+        }
+
         ResponseT MakeRequestDeserialize(Aws::AmazonWebServiceRequest const * const request,
                                      const char* requestName,
                                      Aws::Http::HttpMethod method,
@@ -354,6 +360,33 @@ namespace client
             SerializerT,
             ResponseT,
             ErrorMarshallerT>>;
+
+        /**
+         * SFINAE implementation for refreshing client context params enabled if the client configuration
+         * type has a AccountId member. If there is a corresponding member we want to use that in the endpoint
+         * calculation and set it as a client context parameter.
+         */
+        template <typename T, typename = void>
+        struct HasAccountId : std::false_type {};
+
+        template<typename T>
+        struct HasAccountId<T, decltype(void(std::declval<T>().accountId))> : std::true_type {};
+
+        template<typename ConfigT = ServiceClientConfigurationT, typename std::enable_if<HasAccountId<ConfigT>::value, int>::type = 0>
+        GetContextEndpointParametersOutcome GetContextEndpointParametersImpl(const AwsSmithyClientAsyncRequestContext& ctx) const {
+          Aws::Vector<Aws::Endpoint::EndpointParameter> endpointParameters;
+          const auto resolvedAccountId = ctx.m_awsIdentity->accountId();
+          if (resolvedAccountId.has_value() && !resolvedAccountId.value().empty() && m_clientConfiguration.accountId.empty()) {
+            endpointParameters.emplace_back("AccountId", resolvedAccountId.value(), Aws::Endpoint::EndpointParameter::ParameterOrigin::OPERATION_CONTEXT);
+          }
+          return endpointParameters;
+        }
+
+        template<typename ConfigT = ServiceClientConfigurationT, typename std::enable_if<!HasAccountId<ConfigT>::value, int>::type = 0>
+        GetContextEndpointParametersOutcome GetContextEndpointParametersImpl(const AwsSmithyClientAsyncRequestContext& ctx) const {
+          AWS_UNREFERENCED_PARAM(ctx);
+          return Aws::Vector<Aws::Endpoint::EndpointParameter>{};
+        }
     };
 
 } // namespace client
diff --git a/src/aws-cpp-sdk-core/include/smithy/client/AwsSmithyClientAsyncRequestContext.h b/src/aws-cpp-sdk-core/include/smithy/client/AwsSmithyClientAsyncRequestContext.h
@@ -70,6 +70,7 @@ namespace smithy
             ResponseHandlerFunc m_responseHandler;
             std::shared_ptr<Aws::Utils::Threading::Executor> m_pExecutor;
             std::shared_ptr<interceptor::InterceptorContext> m_interceptorContext;
+            std::shared_ptr<smithy::AwsIdentity> m_awsIdentity;
         };
     } // namespace client
 } // namespace smithy
diff --git a/src/aws-cpp-sdk-core/include/smithy/client/AwsSmithyClientBase.h b/src/aws-cpp-sdk-core/include/smithy/client/AwsSmithyClientBase.h
@@ -18,6 +18,7 @@
 #include <aws/core/utils/FutureOutcome.h>
 #include <aws/core/utils/memory/stl/AWSMap.h>
 #include <aws/core/utils/Outcome.h>
+#include <aws/core/NoResult.h>
 #include <aws/core/http/HttpClientFactory.h>
 #include <aws/core/client/AWSErrorMarshaller.h>
 #include <aws/core/AmazonWebServiceResult.h>
@@ -84,6 +85,8 @@ namespace client
         using SelectAuthSchemeOptionOutcome = Aws::Utils::Outcome<AuthSchemeOption, AWSCoreError>;
         using ResolveEndpointOutcome = Aws::Utils::Outcome<Aws::Endpoint::AWSEndpoint, AWSCoreError>;
         using StreamOutcome = Aws::Utils::Outcome<Aws::AmazonWebServiceResult<Aws::Utils::Stream::ResponseStream>, AWSCoreError >;
+        using IdentityOutcome = Aws::Utils::Outcome<std::shared_ptr<smithy::AwsIdentity>, AWSCoreError>;
+        using GetContextEndpointParametersOutcome = Aws::Utils::Outcome<Aws::Vector<Aws::Endpoint::EndpointParameter>, AWSCoreError>;
 
         /* primary constructor */
         AwsSmithyClientBase(Aws::UniquePtr<Aws::Client::ClientConfiguration>&& clientConfig,
@@ -199,8 +202,10 @@ namespace client
 
         virtual ResolveEndpointOutcome ResolveEndpoint(const Aws::Endpoint::EndpointParameters& endpointParameters, EndpointUpdateCallback&& epCallback) const = 0;
         virtual SelectAuthSchemeOptionOutcome SelectAuthSchemeOption(const AwsSmithyClientAsyncRequestContext& ctx) const = 0;
-        virtual SigningOutcome SignHttpRequest(std::shared_ptr<HttpRequest> httpRequest, const AuthSchemeOption& targetAuthSchemeOption) const = 0;
+        virtual SigningOutcome SignHttpRequest(std::shared_ptr<HttpRequest> httpRequest, const AwsSmithyClientAsyncRequestContext& ctx) const = 0;
         virtual bool AdjustClockSkew(HttpResponseOutcome& outcome, const AuthSchemeOption& authSchemeOption) const = 0;
+        virtual IdentityOutcome ResolveIdentity(const AwsSmithyClientAsyncRequestContext& ctx) const = 0;
+        virtual GetContextEndpointParametersOutcome GetContextEndpointParameters(const AwsSmithyClientAsyncRequestContext& ctx) const = 0;
 
         /* AwsSmithyClientT class binds its config reference to this pointer, so don't remove const and don't re-allocate it.
          * This is done to avoid duplication of config object between this base and actual service template classes.
diff --git a/src/aws-cpp-sdk-core/include/smithy/client/common/AwsSmithyRequestSigning.h b/src/aws-cpp-sdk-core/include/smithy/client/common/AwsSmithyRequestSigning.h
@@ -38,12 +38,42 @@ namespace smithy
         using SigningError = Aws::Client::AWSError<Aws::Client::CoreErrors>;
         using SigningOutcome = Aws::Utils::FutureOutcome<std::shared_ptr<HttpRequest>, SigningError>;
         using HttpResponseOutcome = Aws::Utils::Outcome<std::shared_ptr<Aws::Http::HttpResponse>, Aws::Client::AWSError<Aws::Client::CoreErrors>>;
+        using IdentityOutcome = Aws::Utils::Outcome<std::shared_ptr<smithy::AwsIdentity>, Aws::Client::AWSError<Aws::Client::CoreErrors>>;
 
-        static SigningOutcome SignRequest(std::shared_ptr<HttpRequest> HTTPRequest, const AuthSchemeOption& authSchemeOption,
-                                          const Aws::UnorderedMap<Aws::String, AuthSchemesVariantT>& authSchemes)
+        static IdentityOutcome ResolveIdentity(const client::AwsSmithyClientAsyncRequestContext& ctx,
+          const Aws::UnorderedMap<Aws::String, AuthSchemesVariantT>& authSchemes)
         {
+          auto authSchemeIt = authSchemes.find(ctx.m_authSchemeOption.schemeId);
+          if (authSchemeIt == authSchemes.end())
+          {
+            assert(!"Auth scheme has not been found for a given auth option!");
+            return (SigningError(Aws::Client::CoreErrors::CLIENT_SIGNING_FAILURE,
+                                 "",
+                                 "Requested AuthSchemeOption was not found within client Auth Schemes",
+                                 false/*retryable*/));
+          }
 
-            auto authSchemeIt = authSchemes.find(authSchemeOption.schemeId);
+          const AuthSchemesVariantT& authScheme = authSchemeIt->second;
+          IdentityVisitor visitor(ctx);
+          AuthSchemesVariantT authSchemesVariantCopy(authScheme); // TODO: allow const visiting
+          authSchemesVariantCopy.Visit(visitor);
+
+          if (!visitor.result)
+          {
+            return (SigningError(Aws::Client::CoreErrors::CLIENT_SIGNING_FAILURE,
+                                 "",
+                                 "Failed to sign with an unknown error",
+                                 false/*retryable*/));
+          }
+
+          return std::move(*visitor.result);
+        }
+
+        static SigningOutcome SignRequest(std::shared_ptr<HttpRequest> HTTPRequest,
+          const client::AwsSmithyClientAsyncRequestContext& ctx,
+          const Aws::UnorderedMap<Aws::String, AuthSchemesVariantT>& authSchemes)
+        {
+            auto authSchemeIt = authSchemes.find(ctx.m_authSchemeOption.schemeId);
             if (authSchemeIt == authSchemes.end())
             {
                 assert(!"Auth scheme has not been found for a given auth option!");
@@ -55,8 +85,9 @@ namespace smithy
 
             const AuthSchemesVariantT& authScheme = authSchemeIt->second;
 
-            return SignWithAuthScheme(std::move(HTTPRequest), authScheme, authSchemeOption);
+            return SignWithAuthScheme(std::move(HTTPRequest), authScheme, ctx);
         }
+
         static SigningOutcome PreSignRequest(std::shared_ptr<HttpRequest> httpRequest,
                                   const AuthSchemeOption& authSchemeOption,
                                   const Aws::UnorderedMap<Aws::String, AuthSchemesVariantT>& authSchemes,
@@ -113,59 +144,74 @@ namespace smithy
 
 
     protected:
+        struct IdentityVisitor
+        {
+            IdentityVisitor(const client::AwsSmithyClientAsyncRequestContext& ctx): m_requestContext(ctx)
+            {
+            }
+
+            const client::AwsSmithyClientAsyncRequestContext& m_requestContext;
+            Aws::Crt::Optional<IdentityOutcome> result;
+
+            template <typename AuthSchemeAlternativeT>
+            void operator()(AuthSchemeAlternativeT& authScheme)
+            {
+              using IdentityT = typename std::remove_reference<decltype(authScheme)>::type::IdentityT;
+              using IdentityResolver = IdentityResolverBase<IdentityT>;
+
+              std::shared_ptr<IdentityResolver> identityResolver = authScheme.identityResolver();
+              if (!identityResolver)
+              {
+                result.emplace(SigningError(Aws::Client::CoreErrors::CLIENT_SIGNING_FAILURE,
+                                            "",
+                                            "Auth scheme provided a nullptr identityResolver",
+                                            false/*retryable*/));
+                return;
+              }
+
+              //relay service params in additional properties which will be relevant in credential resolution
+              // example: bucket Name
+              Aws::UnorderedMap<Aws::String, Aws::Crt::Variant<Aws::String, bool>> additionalIdentityProperties;
+              const auto& serviceSpecificParameters = m_requestContext.m_pRequest->GetServiceSpecificParameters();
+              if(serviceSpecificParameters)
+              {
+                for(const auto& propPair : serviceSpecificParameters->parameterMap)
+                {
+                  additionalIdentityProperties.emplace(propPair.first,Aws::Crt::Variant<Aws::String, bool>{propPair.second} );
+                }
+              }
+
+              auto identityResult = identityResolver->getIdentity(m_requestContext.m_authSchemeOption.identityProperties(), additionalIdentityProperties);
+              if (!identityResult.IsSuccess())
+              {
+                result.emplace(identityResult.GetError());
+                return;
+              }
+              result.emplace(std::move(identityResult.GetResultWithOwnership()));
+            }
+        };
+
         struct SignerVisitor
         {
-            SignerVisitor(std::shared_ptr<HttpRequest> httpRequest, const AuthSchemeOption& targetAuthSchemeOption)
-                : m_httpRequest(std::move(httpRequest)), m_targetAuthSchemeOption(targetAuthSchemeOption)
+            SignerVisitor(std::shared_ptr<HttpRequest> httpRequest, const client::AwsSmithyClientAsyncRequestContext& ctx)
+                : m_httpRequest(std::move(httpRequest)), m_requestContext(ctx)
             {
             }
 
             const std::shared_ptr<HttpRequest> m_httpRequest;
-            const AuthSchemeOption& m_targetAuthSchemeOption;
+            const client::AwsSmithyClientAsyncRequestContext& m_requestContext;
 
             Aws::Crt::Optional<SigningOutcome> result;
 
             template <typename AuthSchemeAlternativeT>
             void operator()(AuthSchemeAlternativeT& authScheme)
             {
                 // Auth Scheme Variant alternative contains the requested auth option
-                assert(strcmp(authScheme.schemeId, m_targetAuthSchemeOption.schemeId) == 0);
+                assert(strcmp(authScheme.schemeId, m_requestContext.m_authSchemeOption.schemeId) == 0);
 
                 using IdentityT = typename std::remove_reference<decltype(authScheme)>::type::IdentityT;
-                using IdentityResolver = IdentityResolverBase<IdentityT>;
                 using Signer = AwsSignerBase<IdentityT>;
 
-                std::shared_ptr<IdentityResolver> identityResolver = authScheme.identityResolver();
-                if (!identityResolver)
-                {
-                    result.emplace(SigningError(Aws::Client::CoreErrors::CLIENT_SIGNING_FAILURE,
-                                                "",
-                                                "Auth scheme provided a nullptr identityResolver",
-                                                false/*retryable*/));
-                    return;
-                }
-
-                //relay service params in additional properties which will be relevant in credential resolution
-                // example: bucket Name
-                Aws::UnorderedMap<Aws::String, Aws::Crt::Variant<Aws::String, bool>> additionalIdentityProperties;
-                const auto& serviceSpecificParameters = m_httpRequest->GetServiceSpecificParameters();
-                if(serviceSpecificParameters)
-                {
-                    for(const auto& propPair : serviceSpecificParameters->parameterMap)
-                    {
-                        additionalIdentityProperties.emplace(propPair.first,Aws::Crt::Variant<Aws::String, bool>{propPair.second} );
-                    }
-                }
-
-                auto identityResult = identityResolver->getIdentity(m_targetAuthSchemeOption.identityProperties(), additionalIdentityProperties);
-
-                if (!identityResult.IsSuccess())
-                {
-                    result.emplace(identityResult.GetError());
-                    return;
-                }
-                auto identity = std::move(identityResult.GetResultWithOwnership());
-
                 std::shared_ptr<Signer> signer = authScheme.signer();
                 if (!signer)
                 {
@@ -176,7 +222,9 @@ namespace smithy
                     return;
                 }
 
-                result.emplace(signer->sign(m_httpRequest, *identity, m_targetAuthSchemeOption.signerProperties()));
+                result.emplace(signer->sign(m_httpRequest,
+                  *static_cast<IdentityT*>(m_requestContext.m_awsIdentity.get()),
+                  m_requestContext.m_authSchemeOption.signerProperties()));
             }
         };
 
@@ -236,11 +284,11 @@ namespace smithy
           }
         };
 
-        static
-        SigningOutcome SignWithAuthScheme(std::shared_ptr<HttpRequest> httpRequest, const AuthSchemesVariantT& authSchemesVariant,
-                                          const AuthSchemeOption& targetAuthSchemeOption)
+        static SigningOutcome SignWithAuthScheme(std::shared_ptr<HttpRequest> httpRequest,
+          const AuthSchemesVariantT& authSchemesVariant,
+          const client::AwsSmithyClientAsyncRequestContext& ctx)
         {
-            SignerVisitor visitor(httpRequest, targetAuthSchemeOption);
+            SignerVisitor visitor(httpRequest, ctx);
             AuthSchemesVariantT authSchemesVariantCopy(authSchemesVariant); // TODO: allow const visiting
             authSchemesVariantCopy.Visit(visitor);
 
diff --git a/src/aws-cpp-sdk-core/include/smithy/identity/identity/impl/AwsCredentialIdentityImpl.h b/src/aws-cpp-sdk-core/include/smithy/identity/identity/impl/AwsCredentialIdentityImpl.h
@@ -25,6 +25,6 @@ namespace smithy {
     }
 
     inline Aws::Crt::Optional<Aws::String> AwsCredentialIdentity::accountId() const {
-        return m_sessionToken;
+        return m_accountId;
     }
 }
diff --git a/src/aws-cpp-sdk-core/include/smithy/identity/resolver/built-in/DefaultAwsCredentialIdentityResolver.h b/src/aws-cpp-sdk-core/include/smithy/identity/resolver/built-in/DefaultAwsCredentialIdentityResolver.h
@@ -50,7 +50,7 @@ namespace smithy {
                                                                      legacyCreds.GetAWSSecretKey(),
                                                                      legacyCreds.GetSessionToken().empty()? Aws::Crt::Optional<Aws::String>() : legacyCreds.GetSessionToken(),
                                                                      legacyCreds.GetExpiration(),
-                                                                     legacyCreds.GetAccountId().empty()? Aws::Crt::Optional<Aws::String>() : legacyCreds.GetSessionToken());
+                                                                     legacyCreds.GetAccountId().empty()? Aws::Crt::Optional<Aws::String>() : legacyCreds.GetAccountId());
 
             return ResolveIdentityFutureOutcome(std::move(smithyCreds));
         }
diff --git a/src/aws-cpp-sdk-core/source/smithy/client/AwsSmithyClientBase.cpp b/src/aws-cpp-sdk-core/source/smithy/client/AwsSmithyClientBase.cpp
@@ -223,9 +223,35 @@ void AwsSmithyClientBase::MakeRequestAsync(Aws::AmazonWebServiceRequest const* c
     }
     pRequestCtx->m_authSchemeOption = std::move(authSchemeOptionOutcome.GetResultWithOwnership());
     assert(pRequestCtx->m_authSchemeOption.schemeId);
+
+    // resolve identity
+    auto identityOutcome = this->ResolveIdentity(*pRequestCtx);
+    if (!identityOutcome.IsSuccess())
+    {
+      pExecutor->Submit([identityOutcome, responseHandler]() mutable
+        {
+            responseHandler(std::move(identityOutcome));
+        });
+      return;
+    }
+    pRequestCtx->m_awsIdentity = std::move(identityOutcome.GetResultWithOwnership());
+
+    // get endpoint params from operation context
+    const auto contextEndpointParameters = this->GetContextEndpointParameters(*pRequestCtx);
+    if (!contextEndpointParameters.IsSuccess())
+    {
+      pExecutor->Submit([contextEndpointParameters, responseHandler]() mutable
+        {
+            responseHandler(std::move(contextEndpointParameters.GetError()));
+        });
+      return;
+    }
+
     Aws::Endpoint::EndpointParameters epParams = request ? request->GetEndpointContextParams() : Aws::Endpoint::EndpointParameters();
     const auto authSchemeEpParams = pRequestCtx->m_authSchemeOption.endpointParameters();
     epParams.insert(epParams.end(), authSchemeEpParams.begin(), authSchemeEpParams.end());
+    const auto contextParams = contextEndpointParameters.GetResult();
+    epParams.insert(epParams.end(), contextParams.begin(), contextParams.end());
     auto epResolutionOutcome = this->ResolveEndpoint(std::move(epParams), std::move(endpointCallback));
     if (!epResolutionOutcome.IsSuccess())
     {
@@ -348,7 +374,7 @@ void AwsSmithyClientBase::AttemptOneRequestAsync(std::shared_ptr<AwsSmithyClient
     };
 
     SigningOutcome signingOutcome = TracingUtils::MakeCallWithTiming<SigningOutcome>([&]() -> SigningOutcome {
-            return this->SignHttpRequest(pRequestCtx->m_httpRequest, pRequestCtx->m_authSchemeOption);
+            return this->SignHttpRequest(pRequestCtx->m_httpRequest, *pRequestCtx);
         },
         TracingUtils::SMITHY_CLIENT_SIGNING_METRIC,
         *m_clientConfig->telemetryProvider->getMeter(this->GetServiceClientName(), {}),
diff --git a/tests/aws-cpp-sdk-core-tests/smithy/client/SmithyClientTest.cpp b/tests/aws-cpp-sdk-core-tests/smithy/client/SmithyClientTest.cpp
diff --git a/tests/aws-cpp-sdk-dynamodb-unit-tests/DynamoDBUnitTests.cpp b/tests/aws-cpp-sdk-dynamodb-unit-tests/DynamoDBUnitTests.cpp
diff --git a/tests/testing-resources/include/aws/testing/platform/PlatformTesting.h b/tests/testing-resources/include/aws/testing/platform/PlatformTesting.h

Original file line number	Diff line number	Diff line change
`@@ -25,6 +25,6 @@ namespace smithy {`
`25`	`25`	`}`
`26`	`26`
`27`	`27`	`inline Aws::Crt::Optional<Aws::String> AwsCredentialIdentity::accountId() const {`
`28`		`- return m_sessionToken;`
	`28`	`+ return m_accountId;`
`29`	`29`	`}`
`30`	`30`	`}`
Original file line number	Diff line number	Diff line change
`@@ -50,7 +50,7 @@ namespace smithy {`
`50`	`50`	`legacyCreds.GetAWSSecretKey(),`
`51`	`51`	`legacyCreds.GetSessionToken().empty()? Aws::Crt::Optional<Aws::String>() : legacyCreds.GetSessionToken(),`
`52`	`52`	`legacyCreds.GetExpiration(),`
`53`		`- legacyCreds.GetAccountId().empty()? Aws::Crt::Optional<Aws::String>() : legacyCreds.GetSessionToken());`
	`53`	`+ legacyCreds.GetAccountId().empty()? Aws::Crt::Optional<Aws::String>() : legacyCreds.GetAccountId());`
`54`	`54`
`55`	`55`	`return ResolveIdentityFutureOutcome(std::move(smithyCreds));`
`56`	`56`	`}`