Update request pipeline to use the AuthSchemeResolver to resolve the …

…identity and signer
aws · Nov 18, 2024 · 372f21e · 372f21e
1 parent 19c2d76
commit 372f21e
Show file tree

Hide file tree

Showing 28 changed files with 436 additions and 293 deletions.
diff --git a/extensions/src/Smithy.Identity.Abstractions/Auth/IAuthScheme.cs b/extensions/src/Smithy.Identity.Abstractions/Auth/IAuthScheme.cs
@@ -24,6 +24,9 @@ public interface IAuthScheme<out T> where T : BaseIdentity
         /// </summary>
         IIdentityResolver GetIdentityResolver(IIdentityResolverConfiguration configuration);
 
-        // TODO: Include methods to retrieve signer for the current scheme.
+        /// <summary>
+        /// Retrieves the signer associated with this authentication scheme.
+        /// </summary>
+        ISigner Signer();
     }
 }
diff --git a/generator/ServiceClientGeneratorLib/GeneratorHelpers.cs b/generator/ServiceClientGeneratorLib/GeneratorHelpers.cs
@@ -9,36 +9,16 @@ namespace ServiceClientGenerator
 {
     public static class GeneratorHelpers
     {
-        public static string DetermineSigner(string signatureVersion, string serviceBasename)
+        public static string DetermineSignerOverride(string signatureVersion)
         {
-            switch (serviceBasename)
-            {
-                case "EventBridge":
-                    // we should not continue to add new hardcoded service specific signers
-                    // and instead implement a solution based on a signer selection specification
-                    return "EventBridgeSigner";
-                case "CloudFrontKeyValueStore":
-                    return "AWS4aSignerCRTWrapper";
-            }
-
             switch (signatureVersion)
             {
-                case "v2":
-                    return "QueryStringSigner";
-                case "v3https":
-                    return "AWS3Signer";
-                case "v4":
-                    return "AWS4Signer";
                 case "s3":
                     return "Amazon.S3.Internal.S3Signer";
                 case "s3v4":
                     return "S3Signer";
-                case "bearer":
-                    return "BearerTokenSigner";
-                case "":
-                    return "NullSigner";
                 default:
-                    throw new Exception("Unknown signer: " + signatureVersion);
+                    return null;
             }
         }
 

diff --git a/generator/ServiceClientGeneratorLib/Generators/AuthResolvers/ModeledResolver.cs b/generator/ServiceClientGeneratorLib/Generators/AuthResolvers/ModeledResolver.cs
diff --git a/generator/ServiceClientGeneratorLib/Generators/AuthResolvers/ModeledResolver.tt b/generator/ServiceClientGeneratorLib/Generators/AuthResolvers/ModeledResolver.tt
@@ -11,6 +11,7 @@
 using Amazon.Runtime;
 using Amazon.Runtime.Credentials.Internal;
 using Amazon.Runtime.Internal;
+using Amazon.Runtime.Internal.Auth;
 using Smithy.Identity.Abstractions;
 using System.Collections.Generic;
 
@@ -35,6 +36,13 @@ namespace <#=this.Config.Namespace#>.Internal
     /// </summary>
     public class Amazon<#=this.Config.ClassName#>AuthSchemeHandler : BaseAuthResolverHandler
     {
+        /// <summary>
+        /// Construct an instance of Amazon<#=this.Config.ClassName#>AuthSchemeHandler
+        /// </summary>
+        public Amazon<#=this.Config.ClassName#>AuthSchemeHandler(BaseIdentity identity)
+            : base(identity)
+        { }
+
 <# if (IsServiceAllowListedForRulesBasedResolver()) { #>
         private readonly Amazon<#=this.Config.ClassName#>EndpointResolver _endpointResolver = new();
 
@@ -72,6 +80,14 @@ namespace <#=this.Config.Namespace#>.Internal
 
             return AuthSchemeResolver.ResolveAuthScheme(mappedParameters);
         }
+<# if (GeneratorHelpers.DetermineSignerOverride(this.Config.ServiceModel.SignatureVersion) != null) { #>
+
+        /// <inheritdoc/>
+        protected override AbstractAWSSigner GetSigner(IAuthScheme<BaseIdentity> scheme)
+        {
+            return new <#=GeneratorHelpers.DetermineSignerOverride(this.Config.ServiceModel.SignatureVersion)#>();
+        }
+<# } #>
     }
 
     /// <inheritdoc cref="IAuthSchemeResolver{T}" />

diff --git a/generator/ServiceClientGeneratorLib/Generators/SourceFiles/ServiceClientsNetFramework.cs b/generator/ServiceClientGeneratorLib/Generators/SourceFiles/ServiceClientsNetFramework.cs
diff --git a/generator/ServiceClientGeneratorLib/Generators/SourceFiles/ServiceClientsNetFramework.tt b/generator/ServiceClientGeneratorLib/Generators/SourceFiles/ServiceClientsNetFramework.tt
@@ -233,6 +233,12 @@ namespace <#=this.Config.Namespace#>
         /// <param name="pipeline"></param>
         protected override void CustomizeRuntimePipeline(RuntimePipeline pipeline)
         {
+<# 
+// TODO: Remove the if statement when we eventually generate auth resolvers for all services.
+if (this.Config.ClassName == "AutoScaling" || this.Config.ClassName == "S3") { 
+#>
+            pipeline.AddHandlerBefore<Amazon.Runtime.Internal.CredentialsRetriever>(new Amazon<#=this.Config.ClassName#>AuthSchemeHandler(this.DefaultAWSCredentials));
+<# } #>
 <#
     var pipelineOverrides = this.Config.ServiceModel.Customizations.PipelineOverride;
     if (pipelineOverrides != null)

diff --git a/generator/ServiceClientGeneratorLib/Generators/SourceFiles/ServiceClientsNetStandard.cs b/generator/ServiceClientGeneratorLib/Generators/SourceFiles/ServiceClientsNetStandard.cs
diff --git a/generator/ServiceClientGeneratorLib/Generators/SourceFiles/ServiceClientsNetStandard.tt b/generator/ServiceClientGeneratorLib/Generators/SourceFiles/ServiceClientsNetStandard.tt
@@ -237,6 +237,12 @@ namespace <#=this.Config.Namespace#>
         /// <param name="pipeline">Runtime pipeline for the current client.</param>
         protected override void CustomizeRuntimePipeline(RuntimePipeline pipeline)
         {
+<# 
+// TODO: Remove the if statement when we eventually generate auth resolvers for all services.
+if (this.Config.ClassName == "AutoScaling" || this.Config.ClassName == "S3") { 
+#>
+            pipeline.AddHandlerBefore<Amazon.Runtime.Internal.CredentialsRetriever>(new Amazon<#=this.Config.ClassName#>AuthSchemeHandler(this.DefaultAWSCredentials));
+<# } #>
 <#
     var pipelineOverrides = this.Config.ServiceModel.Customizations.PipelineOverride;
     if (pipelineOverrides != null)

diff --git a/generator/ServiceModels/cognito-sync/cognito-sync.customizations.json b/generator/ServiceModels/cognito-sync/cognito-sync.customizations.json
@@ -8,8 +8,7 @@
             {
                 "operation":"addBefore",
                 "newType":"Amazon.CognitoSync.Internal.CognitoCredentialsRetriever",
-                "targetType":"Amazon.Runtime.Internal.Marshaller",
-                "constructorInput":"this.DefaultAWSCredentials"
+                "targetType":"Amazon.Runtime.Internal.Marshaller"
             }
         ]
     },

diff --git a/sdk/src/Core/Amazon.Runtime/AmazonServiceClient.cs b/sdk/src/Core/Amazon.Runtime/AmazonServiceClient.cs
@@ -387,7 +387,7 @@ private void BuildRuntimePipeline()
                     // ChecksumHandler must come after EndpointsResolver because of an upcoming project.
                     new ChecksumHandler(),
                     // CredentialsRetriever must come after RetryHandler because of any credential related changes.
-                    new CredentialsRetriever(this.DefaultAWSCredentials),
+                    new CredentialsRetriever(),
                     new RetryHandler(retryPolicy),
                     new CompressionHandler(),
                     postMarshallHandler,

diff --git a/sdk/src/Core/Amazon.Runtime/Credentials/Internal/AnonymousAuthScheme.cs b/sdk/src/Core/Amazon.Runtime/Credentials/Internal/AnonymousAuthScheme.cs
@@ -14,6 +14,7 @@
  */
 
 using Smithy.Identity.Abstractions;
+using Amazon.Runtime.Internal.Auth;
 
 namespace Amazon.Runtime.Credentials.Internal
 {
@@ -30,5 +31,11 @@ public IIdentityResolver GetIdentityResolver(IIdentityResolverConfiguration conf
         {
             return configuration.GetIdentityResolver<AnonymousAWSCredentials>();       
         }
+
+        /// <inheritdoc/>
+        public ISigner Signer()
+        {
+            return new NullSigner();
+        }
     }
 }
diff --git a/sdk/src/Core/Amazon.Runtime/Credentials/Internal/AwsV4AuthScheme.cs b/sdk/src/Core/Amazon.Runtime/Credentials/Internal/AwsV4AuthScheme.cs
@@ -13,6 +13,7 @@
 * permissions and limitations under the License.
 */
 
+using Amazon.Runtime.Internal.Auth;
 using Smithy.Identity.Abstractions;
 
 namespace Amazon.Runtime.Credentials.Internal
@@ -31,5 +32,11 @@ public IIdentityResolver GetIdentityResolver(IIdentityResolverConfiguration conf
         {
             return configuration.GetIdentityResolver<AWSCredentials>();
         }
+
+        /// <inheritdoc/>
+        public ISigner Signer()
+        {
+            return new AWS4Signer();
+        }
     }
 }
diff --git a/sdk/src/Core/Amazon.Runtime/Credentials/Internal/AwsV4aAuthScheme.cs b/sdk/src/Core/Amazon.Runtime/Credentials/Internal/AwsV4aAuthScheme.cs
@@ -13,6 +13,7 @@
 * permissions and limitations under the License.
 */
 
+using Amazon.Runtime.Internal.Auth;
 using Smithy.Identity.Abstractions;
 
 namespace Amazon.Runtime.Credentials.Internal
@@ -31,5 +32,11 @@ public IIdentityResolver GetIdentityResolver(IIdentityResolverConfiguration conf
         {
             return configuration.GetIdentityResolver<AWSCredentials>();
         }
+
+        /// <inheritdoc/>
+        public ISigner Signer()
+        {
+            return new AWS4aSignerCRTWrapper();
+        }
     }
 }
diff --git a/sdk/src/Core/Amazon.Runtime/Internal/Auth/AbstractAWSSigner.cs b/sdk/src/Core/Amazon.Runtime/Internal/Auth/AbstractAWSSigner.cs
@@ -20,6 +20,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 #endif
+using Smithy.Identity.Abstractions;
 using Amazon.Internal;
 using Amazon.Util;
 using Amazon.Runtime.Internal.Util;
@@ -28,7 +29,7 @@ namespace Amazon.Runtime.Internal.Auth
 {
     public enum ClientProtocol { QueryStringProtocol, RestProtocol, Unknown }
 
-    public abstract class AbstractAWSSigner
+    public abstract class AbstractAWSSigner : ISigner
     {
         private readonly object _lock = new object();
         private AWS4Signer _aws4Signer;

diff --git a/sdk/src/Core/Amazon.Runtime/Pipeline/Contexts.cs b/sdk/src/Core/Amazon.Runtime/Pipeline/Contexts.cs
@@ -13,6 +13,7 @@
  * permissions and limitations under the License.
  */
 
+using Smithy.Identity.Abstractions;
 using Amazon.Runtime.Internal;
 using Amazon.Runtime.Internal.Auth;
 using Amazon.Runtime.Internal.Transform;
@@ -32,6 +33,7 @@ public interface IRequestContext
         InvokeOptionsBase Options { get; }        
         RequestMetrics Metrics { get; }
         AbstractAWSSigner Signer { get; set; }
+        BaseIdentity Identity { get; set; }
         IClientConfig ClientConfig { get; }
         ImmutableCredentials ImmutableCredentials { get; set; }
 
@@ -122,6 +124,7 @@ public RequestContext(bool enableMetrics, AbstractAWSSigner clientSigner)
         public InvokeOptionsBase Options { get; set; }        
         public ImmutableCredentials ImmutableCredentials { get; set; }
         public AbstractAWSSigner Signer { get; set; }
+        public BaseIdentity Identity { get; set; }
 
 #if AWS_ASYNC_API
         public System.Threading.CancellationToken CancellationToken { get; set; }

diff --git a/sdk/src/Core/Amazon.Runtime/Pipeline/Handlers/BaseAuthResolverHandler.cs b/sdk/src/Core/Amazon.Runtime/Pipeline/Handlers/BaseAuthResolverHandler.cs
@@ -15,6 +15,7 @@
 
 using Amazon.Runtime.Credentials.Internal;
 using Amazon.Runtime.Endpoints;
+using Amazon.Runtime.Internal.Auth;
 using Smithy.Identity.Abstractions;
 using System.Collections;
 using System.Collections.Generic;
@@ -32,13 +33,32 @@ public abstract class BaseAuthResolverHandler : PipelineHandler
             new AwsV4AuthScheme()
         };
 
-        // Actual pipeline handler implementation will include other methods (such as InvokeAsync) too.
+        public BaseAuthResolverHandler(BaseIdentity identity)
+        {
+            this.Identity = identity;
+        }
+
+        protected BaseIdentity Identity
+        {
+            get;
+            private set;
+        }
+
+
         public override void InvokeSync(IExecutionContext executionContext)
         {
             PreInvoke(executionContext);
             base.InvokeSync(executionContext);
         }
 
+#if AWS_ASYNC_API
+        public override System.Threading.Tasks.Task<T> InvokeAsync<T>(IExecutionContext executionContext)
+        {
+            PreInvoke(executionContext);
+            return base.InvokeAsync<T>(executionContext);
+        }
+#endif
+
         protected void PreInvoke(IExecutionContext executionContext)
         {
             var authOptions = ResolveAuthOptions(executionContext);
@@ -56,17 +76,22 @@ protected void PreInvoke(IExecutionContext executionContext)
                     continue;
                 }
 
-                // TODO: Retrieve identity resolver and signer for the current scheme and attach them to request context.
-                // Interfaces haven't been defined yet, but code will be similar to:
-                // var identityResolver = scheme.IdentityResolver();
-                // var identity = identityResolver.GetIdentity();
-                // var signer = scheme.Signer();
+                executionContext.RequestContext.Identity = this.Identity;
+                if (executionContext.RequestContext.Identity == null)
+                {
+                    var identityResolver = scheme.GetIdentityResolver(executionContext.RequestContext.ClientConfig.IdentityResolverConfiguration);
+                    executionContext.RequestContext.Identity = identityResolver.ResolveIdentity();
+                }
 
-                // This code would also need to handle existing scenarios, such as:
-                // - Customer explicitly set credentials on their service client (i.e. new AmazonS3Client(new BasicAWSCredentials))
+                executionContext.RequestContext.Signer = GetSigner(scheme);
             }
         }
 
+        protected virtual AbstractAWSSigner GetSigner(IAuthScheme<BaseIdentity> scheme)
+        {
+            return scheme.Signer() as AbstractAWSSigner;
+        }
+
         /// <summary>
         /// Certain allow-listed services contain auth schemes in their endpoint ruleset. For those, we'll attempt to
         /// retrieve the auth options from the attributes of the resolved endpoint.

diff --git a/sdk/src/Core/Amazon.Runtime/Pipeline/Handlers/CredentialsRetriever.cs b/sdk/src/Core/Amazon.Runtime/Pipeline/Handlers/CredentialsRetriever.cs
@@ -29,21 +29,6 @@ namespace Amazon.Runtime.Internal
     /// </summary>
     public class CredentialsRetriever : PipelineHandler
     {
-        /// <summary>
-        /// The constructor for CredentialsRetriever.
-        /// </summary>
-        /// <param name="credentials">An AWSCredentials instance.</param>
-        public CredentialsRetriever(AWSCredentials credentials)
-        {
-            this.Credentials = credentials;
-        }
-
-        protected AWSCredentials Credentials
-        {
-            get;
-            private set;
-        }
-
         /// <summary>
         /// Retrieves the credentials to be used for the current call before 
         /// invoking the next handler.
@@ -52,13 +37,15 @@ protected AWSCredentials Credentials
         protected virtual void PreInvoke(IExecutionContext executionContext)
         {
             ImmutableCredentials ic = null;
-            if (Credentials != null && !(Credentials is AnonymousAWSCredentials) && !(executionContext.RequestContext.Signer is BearerTokenSigner))
+            var identity = executionContext.RequestContext.Identity as AWSCredentials;
+
+            if (identity != null && !(identity is AnonymousAWSCredentials) && !(executionContext.RequestContext.Signer is BearerTokenSigner))
             {
                 using (TracingUtilities.CreateSpan(executionContext.RequestContext, TelemetryConstants.CredentialsRetrievalSpanName))
                 using (MetricsUtilities.MeasureDuration(executionContext.RequestContext, TelemetryConstants.ResolveIdentityDurationMetricName))
                 using (executionContext.RequestContext.Metrics.StartEvent(Metric.CredentialsRequestTime))
                 {
-                    ic = Credentials.GetCredentials();
+                    ic = identity.GetCredentials();
                 }
             }
 
@@ -89,13 +76,15 @@ public override void InvokeSync(IExecutionContext executionContext)
         public override async System.Threading.Tasks.Task<T> InvokeAsync<T>(IExecutionContext executionContext)
         {
             ImmutableCredentials ic = null;
-            if (Credentials != null && !(Credentials is AnonymousAWSCredentials))
+            var identity = executionContext.RequestContext.Identity as AWSCredentials;
+
+            if (identity != null && !(identity is AnonymousAWSCredentials))
             {
                 using (TracingUtilities.CreateSpan(executionContext.RequestContext, TelemetryConstants.CredentialsRetrievalSpanName))
                 using (MetricsUtilities.MeasureDuration(executionContext.RequestContext, TelemetryConstants.ResolveIdentityDurationMetricName))
                 using(executionContext.RequestContext.Metrics.StartEvent(Metric.CredentialsRequestTime))
                 {
-                    ic = await Credentials.GetCredentialsAsync().ConfigureAwait(false);
+                    ic = await identity.GetCredentialsAsync().ConfigureAwait(false);
                 }
             }
 

diff --git a/sdk/src/Services/AutoScaling/Generated/Internal/AmazonAutoScalingAuthResolver.cs b/sdk/src/Services/AutoScaling/Generated/Internal/AmazonAutoScalingAuthResolver.cs
@@ -20,6 +20,7 @@
 using Amazon.Runtime;
 using Amazon.Runtime.Credentials.Internal;
 using Amazon.Runtime.Internal;
+using Amazon.Runtime.Internal.Auth;
 using Smithy.Identity.Abstractions;
 using System.Collections.Generic;
 
@@ -42,6 +43,13 @@ public class AmazonAutoScalingAuthSchemeParameters : IAuthSchemeParameters
     /// </summary>
     public class AmazonAutoScalingAuthSchemeHandler : BaseAuthResolverHandler
     {
+        /// <summary>
+        /// Construct an instance of AmazonAutoScalingAuthSchemeHandler
+        /// </summary>
+        public AmazonAutoScalingAuthSchemeHandler(BaseIdentity identity)
+            : base(identity)
+        { }
+
         /// <summary>
         /// Modeled auth scheme resolver for AutoScaling.
         /// </summary>

diff --git a/sdk/src/Services/AutoScaling/Generated/_bcl/AmazonAutoScalingClient.cs b/sdk/src/Services/AutoScaling/Generated/_bcl/AmazonAutoScalingClient.cs
@@ -238,6 +238,7 @@ public AmazonAutoScalingClient(string awsAccessKeyId, string awsSecretAccessKey,
         /// <param name="pipeline"></param>
         protected override void CustomizeRuntimePipeline(RuntimePipeline pipeline)
         {
+            pipeline.AddHandlerBefore<Amazon.Runtime.Internal.CredentialsRetriever>(new AmazonAutoScalingAuthSchemeHandler(this.DefaultAWSCredentials));
             pipeline.RemoveHandler<Amazon.Runtime.Internal.EndpointResolver>();
             pipeline.AddHandlerAfter<Amazon.Runtime.Internal.Marshaller>(new AmazonAutoScalingEndpointResolver());
         }