Add support for overriding authSchemeProvider per request

.changes/next-release/feature-AWSSDKforJavav2-1f3cedc.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Add support for overriding authSchemeProvider per request."
+}

core/aws-core/src/main/java/software/amazon/awssdk/awscore/internal/AwsExecutionContextBuilder.java

@@ -269,9 +269,9 @@ private static void putAuthSchemeResolutionAttributes(ExecutionAttributes execut
                                                           SdkClientConfiguration clientConfig,
                                                           SdkRequest originalRequest) {
 
-        // TODO(request-override auth scheme feature): When request-level auth scheme provider is added, use the request-level
-        //  auth scheme provider if the customer specified an override, otherwise fall back to the one on the client.
-        AuthSchemeProvider authSchemeProvider = clientConfig.option(SdkClientOption.AUTH_SCHEME_PROVIDER);
+        // Use the request-level auth scheme provider if the customer specified an override, otherwise fall back to the one
+        // on the client.
+        AuthSchemeProvider authSchemeProvider = resolveAuthSchemeProvider(originalRequest, clientConfig);
 
         // Use auth schemes that the user specified at the request level with
         // preference over those on the client.
@@ -360,6 +360,19 @@ private static EndpointProvider resolveEndpointProvider(SdkRequest request,
                       .orElse(clientConfig.option(SdkClientOption.ENDPOINT_PROVIDER));
     }
 
+    /**
+     * Resolves the auth scheme provider, with the request override configuration taking precedence over the provided client
+     * configuration.
+     *
+     * @return The auth scheme provider that will be used by the SDK to resolve auth schemes.
+     */
+    private static AuthSchemeProvider resolveAuthSchemeProvider(SdkRequest request,
+                                                               SdkClientConfiguration clientConfig) {
+        return request.overrideConfiguration()
+                      .flatMap(RequestOverrideConfiguration::authSchemeProvider)
+                      .orElse(clientConfig.option(SdkClientOption.AUTH_SCHEME_PROVIDER));
+    }
+
     private static <InputT extends SdkRequest, OutputT extends SdkResponse> BusinessMetricCollection
         resolveUserAgentBusinessMetrics(SdkClientConfiguration clientConfig, 
                                         ClientExecutionParams<InputT, OutputT> executionParams) {

core/aws-core/src/test/java/software/amazon/awssdk/awscore/internal/AwsExecutionContextBuilderTest.java

@@ -72,6 +72,7 @@
 import software.amazon.awssdk.http.auth.scheme.NoAuthAuthScheme;
 import software.amazon.awssdk.http.auth.spi.scheme.AuthScheme;
 import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
+import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeProvider;
 import software.amazon.awssdk.http.auth.spi.signer.HttpSigner;
 import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
 import software.amazon.awssdk.identity.spi.IdentityProvider;
@@ -557,6 +558,51 @@ public void invokeInterceptorsAndCreateExecutionContext_withExplicitHttpClientFa
         assertThat(businessMetrics.recordedMetrics()).contains("AL");
     }
 
+    @Test
+    public void invokeInterceptorsAndCreateExecutionContext_authSchemeProviderRequestOverride_usesRequestOverride() {
+        AuthSchemeProvider clientAuthSchemeProvider = mock(AuthSchemeProvider.class);
+        AuthSchemeProvider requestAuthSchemeProvider = mock(AuthSchemeProvider.class);
+
+        SdkClientConfiguration clientConfig = testClientConfiguration()
+            .option(SdkClientOption.AUTH_SCHEME_PROVIDER, clientAuthSchemeProvider)
+            .build();
+
+        Optional overrideConfiguration =
+            Optional.of(AwsRequestOverrideConfiguration.builder()
+                                                       .authSchemeProvider(requestAuthSchemeProvider)
+                                                       .build());
+        when(sdkRequest.overrideConfiguration()).thenReturn(overrideConfiguration);
+
+        ClientExecutionParams<SdkRequest, SdkResponse> executionParams = clientExecutionParams();
+
+        ExecutionContext executionContext =
+            AwsExecutionContextBuilder.invokeInterceptorsAndCreateExecutionContext(executionParams, clientConfig);
+
+        AuthSchemeProvider actualProvider =
+            executionContext.executionAttributes().getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER);
+
+        assertThat(actualProvider).isSameAs(requestAuthSchemeProvider);
+    }
+
+    @Test
+    public void invokeInterceptorsAndCreateExecutionContext_noAuthSchemeProviderRequestOverride_usesClientProvider() {
+        AuthSchemeProvider clientAuthSchemeProvider = mock(AuthSchemeProvider.class);
+
+        SdkClientConfiguration clientConfig = testClientConfiguration()
+            .option(SdkClientOption.AUTH_SCHEME_PROVIDER, clientAuthSchemeProvider)
+            .build();
+
+        ClientExecutionParams<SdkRequest, SdkResponse> executionParams = clientExecutionParams();
+
+        ExecutionContext executionContext =
+            AwsExecutionContextBuilder.invokeInterceptorsAndCreateExecutionContext(executionParams, clientConfig);
+
+        AuthSchemeProvider actualProvider =
+            executionContext.executionAttributes().getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEME_RESOLVER);
+
+        assertThat(actualProvider).isSameAs(clientAuthSchemeProvider);
+    }
+
     private ClientExecutionParams<SdkRequest, SdkResponse> clientExecutionParams() {
         return clientExecutionParams(sdkRequest);
     }

core/sdk-core/src/main/java/software/amazon/awssdk/core/RequestOverrideConfiguration.java

@@ -33,6 +33,7 @@
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 import software.amazon.awssdk.core.signer.Signer;
 import software.amazon.awssdk.endpoints.EndpointProvider;
+import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeProvider;
 import software.amazon.awssdk.metrics.MetricPublisher;
 import software.amazon.awssdk.utils.CollectionUtils;
 import software.amazon.awssdk.utils.Validate;
@@ -53,6 +54,7 @@ public abstract class RequestOverrideConfiguration {
     private final List<MetricPublisher> metricPublishers;
     private final ExecutionAttributes executionAttributes;
     private final EndpointProvider endpointProvider;
+    private final AuthSchemeProvider authSchemeProvider;
     private final CompressionConfiguration compressionConfiguration;
     private final List<SdkPlugin> plugins;
 
@@ -66,6 +68,7 @@ protected RequestOverrideConfiguration(Builder<?> builder) {
         this.metricPublishers = Collections.unmodifiableList(new ArrayList<>(builder.metricPublishers()));
         this.executionAttributes = ExecutionAttributes.unmodifiableExecutionAttributes(builder.executionAttributes());
         this.endpointProvider = builder.endpointProvider();
+        this.authSchemeProvider = builder.authSchemeProvider();
         this.compressionConfiguration = builder.compressionConfiguration();
         this.plugins = Collections.unmodifiableList(new ArrayList<>(builder.plugins()));
     }
@@ -177,6 +180,14 @@ public Optional<EndpointProvider> endpointProvider() {
         return Optional.ofNullable(endpointProvider);
     }
 
+    /**
+     * Returns the auth scheme provider for resolving the auth scheme for this request. This supersedes the
+     * auth scheme provider set on the client.
+     */
+    public Optional<AuthSchemeProvider> authSchemeProvider() {
+        return Optional.ofNullable(authSchemeProvider);
+    }
+
     /**
      * Returns the compression configuration object, if present, which includes options to enable/disable compression and set
      * the minimum compression threshold. This compression config object supersedes the compression config object set on the
@@ -204,6 +215,7 @@ public boolean equals(Object o) {
                Objects.equals(metricPublishers, that.metricPublishers) &&
                Objects.equals(executionAttributes, that.executionAttributes) &&
                Objects.equals(endpointProvider, that.endpointProvider) &&
+               Objects.equals(authSchemeProvider, that.authSchemeProvider) &&
                Objects.equals(compressionConfiguration, that.compressionConfiguration) &&
                Objects.equals(plugins, that.plugins);
     }
@@ -220,6 +232,7 @@ public int hashCode() {
         hashCode = 31 * hashCode + Objects.hashCode(metricPublishers);
         hashCode = 31 * hashCode + Objects.hashCode(executionAttributes);
         hashCode = 31 * hashCode + Objects.hashCode(endpointProvider);
+        hashCode = 31 * hashCode + Objects.hashCode(authSchemeProvider);
         hashCode = 31 * hashCode + Objects.hashCode(compressionConfiguration);
         hashCode = 31 * hashCode + Objects.hashCode(plugins);
         return hashCode;
@@ -456,13 +469,57 @@ default B putRawQueryParameter(String name, String value) {
          * over the endpointProvider set on the client while resolving the endpoint for  the requests.
          * If this value is null, then the client level endpointProvider is used for resolving the endpoint.
          *
-         * @param endpointProvider Endpoint Provider that will override the resolving the endpoint for the request.
+         * @param endpointProvider Endpoint Provider that will override client's configured or default EndpointProvider
+         *                         for the request.
          * @return This object for method chaining
          */
         B endpointProvider(EndpointProvider endpointProvider);
 
         EndpointProvider endpointProvider();
 
+        /**
+         * Sets the auth scheme provider to use for resolving the auth scheme of the request. This auth scheme provider gets
+         * priority over the auth scheme provider set on the client while resolving the auth scheme for the requests.
+         * If this value is null, then the client level auth scheme provider is used for resolving the auth scheme.
+         *
+         * <p>Example — overriding the signing region for a single request:
+         * {@snippet :
+         *   // Create a custom auth scheme provider that overrides the signing region.
+         *   // This wraps the default provider and modifies the resolved auth scheme options.
+         *   S3AuthSchemeProvider defaultProvider = S3AuthSchemeProvider.defaultProvider();
+         *
+         *   S3AuthSchemeProvider customRegionProvider = params -> {
+         *       List<AuthSchemeOption> options = defaultProvider.resolveAuthScheme(params);
+         *       return options.stream()
+         *                     .map(option -> option.toBuilder()
+         *                                          .putSignerProperty(AwsV4HttpSigner.REGION_NAME, "eu-west-1")
+         *                                          .build())
+         *                     .collect(Collectors.toList());
+         *   };
+         *
+         *   S3Client s3 = S3Client.builder()
+         *                         .region(Region.US_EAST_1)
+         *                         .build();
+         *
+         *   // Override the auth scheme provider for a single request
+         *   GetObjectResponse response = s3.getObject(
+         *       GetObjectRequest.builder()
+         *           .bucket("my-bucket")
+         *           .key("my-key")
+         *           .overrideConfiguration(c -> c.authSchemeProvider(customRegionProvider))
+         *           .build(),
+         *       ResponseTransformer.toBytes()
+         *   );
+         * }
+         *
+         * @param authSchemeProvider Auth scheme provider that will override the client's configured or default
+         *                           AuthSchemeProvider for the request.
+         * @return This object for method chaining
+         */
+        B authSchemeProvider(AuthSchemeProvider authSchemeProvider);
+
+        AuthSchemeProvider authSchemeProvider();
+
         /**
          * Sets the {@link CompressionConfiguration} for this request. The order of precedence, from highest to lowest,
          * for this setting is: 1) Per request configuration 2) Client configuration 3) Environment variables 4) Profile setting.
@@ -546,6 +603,7 @@ protected abstract static class BuilderImpl<B extends Builder> implements Builde
         private List<MetricPublisher> metricPublishers = new ArrayList<>();
         private ExecutionAttributes.Builder executionAttributesBuilder = ExecutionAttributes.builder();
         private EndpointProvider endpointProvider;
+        private AuthSchemeProvider authSchemeProvider;
         private CompressionConfiguration compressionConfiguration;
         private List<SdkPlugin> plugins = new ArrayList<>();
 
@@ -563,6 +621,7 @@ protected BuilderImpl(RequestOverrideConfiguration sdkRequestOverrideConfig) {
             metricPublishers(sdkRequestOverrideConfig.metricPublishers());
             executionAttributes(sdkRequestOverrideConfig.executionAttributes());
             endpointProvider(sdkRequestOverrideConfig.endpointProvider);
+            authSchemeProvider(sdkRequestOverrideConfig.authSchemeProvider);
             compressionConfiguration(sdkRequestOverrideConfig.compressionConfiguration);
             plugins(sdkRequestOverrideConfig.plugins);
         }
@@ -734,6 +793,21 @@ public EndpointProvider endpointProvider() {
             return endpointProvider;
         }
 
+        @Override
+        public B authSchemeProvider(AuthSchemeProvider authSchemeProvider) {
+            this.authSchemeProvider = authSchemeProvider;
+            return (B) this;
+        }
+
+        public void setAuthSchemeProvider(AuthSchemeProvider authSchemeProvider) {
+            authSchemeProvider(authSchemeProvider);
+        }
+
+        @Override
+        public AuthSchemeProvider authSchemeProvider() {
+            return authSchemeProvider;
+        }
+
         @Override
         public B compressionConfiguration(CompressionConfiguration compressionConfiguration) {
             this.compressionConfiguration = compressionConfiguration;

core/sdk-core/src/test/java/software/amazon/awssdk/core/RequestOverrideConfigurationTest.java

@@ -34,6 +34,7 @@
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 import software.amazon.awssdk.core.signer.Signer;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
+import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeProvider;
 import software.amazon.awssdk.metrics.MetricPublisher;
 import software.amazon.awssdk.utils.ImmutableMap;
 
@@ -75,6 +76,7 @@ public void toBuilder_maximal() {
                 .apiCallTimeout(Duration.ofSeconds(1))
                 .apiCallAttemptTimeout(Duration.ofSeconds(1))
                 .signer(new NoOpSigner())
+                .authSchemeProvider(mock(AuthSchemeProvider.class))
                 .executionAttributes(ExecutionAttributes.builder().put(testAttribute, expectedValue).build())
                 .addMetricPublisher(mock(MetricPublisher.class))
                 .build();
@@ -330,6 +332,46 @@ public void testConfigurationEquals() {
         assertThat(request1Override).isNotEqualTo(null);
     }
 
+    @Test
+    public void authSchemeProvider_setOnBuilder_isPresent() {
+        AuthSchemeProvider provider = mock(AuthSchemeProvider.class);
+        RequestOverrideConfiguration configuration = SdkRequestOverrideConfiguration.builder()
+                .authSchemeProvider(provider)
+                .build();
+
+        assertThat(configuration.authSchemeProvider()).isPresent();
+        assertThat(configuration.authSchemeProvider().get()).isSameAs(provider);
+    }
+
+    @Test
+    public void authSchemeProvider_notSet_isEmpty() {
+        RequestOverrideConfiguration configuration = SdkRequestOverrideConfiguration.builder()
+                .build();
+
+        assertThat(configuration.authSchemeProvider()).isEmpty();
+    }
+
+    @Test
+    public void authSchemeProvider_setToNull_isEmpty() {
+        RequestOverrideConfiguration configuration = SdkRequestOverrideConfiguration.builder()
+                .authSchemeProvider(null)
+                .build();
+
+        assertThat(configuration.authSchemeProvider()).isEmpty();
+    }
+
+    @Test
+    public void authSchemeProvider_toBuilder_roundTrips() {
+        AuthSchemeProvider provider = mock(AuthSchemeProvider.class);
+        RequestOverrideConfiguration configuration = SdkRequestOverrideConfiguration.builder()
+                .authSchemeProvider(provider)
+                .build();
+
+        RequestOverrideConfiguration rebuilt = configuration.toBuilder().build();
+        assertThat(rebuilt.authSchemeProvider()).isPresent();
+        assertThat(rebuilt.authSchemeProvider().get()).isSameAs(provider);
+    }
+
     private static class NoOpSigner implements Signer {
 
         @Override