Enable request-level token provider override

.changes/next-release/feature-AWSSDKforJavav2-bc8f1f1.json

@@ -0,0 +1,6 @@
+{
+    "category": "AWS SDK for Java v2", 
+    "contributor": "chenying-wang", 
+    "type": "feature", 
+    "description": "Enable request-level token provider override"
+}

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

@@ -19,6 +19,7 @@
 import java.util.Optional;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
+import software.amazon.awssdk.auth.token.credentials.SdkTokenProvider;
 import software.amazon.awssdk.core.RequestOverrideConfiguration;
 import software.amazon.awssdk.utils.builder.SdkBuilder;
 
@@ -28,10 +29,12 @@
 @SdkPublicApi
 public final class AwsRequestOverrideConfiguration extends RequestOverrideConfiguration {
     private final AwsCredentialsProvider credentialsProvider;
+    private final SdkTokenProvider tokenProvider;
 
     private AwsRequestOverrideConfiguration(Builder builder) {
         super(builder);
         this.credentialsProvider = builder.credentialsProvider();
+        this.tokenProvider = builder.tokenProvider();
     }
 
     /**
@@ -62,6 +65,16 @@ public Optional<AwsCredentialsProvider> credentialsProvider() {
         return Optional.ofNullable(credentialsProvider);
     }
 
+    /**
+     * The optional {@link SdkTokenProvider} that will provide a token to be used to authorize this request. This will
+     * be used only if the requested operation uses bearer token authorization.
+     *
+     * @return The optional {@link SdkTokenProvider}.
+     */
+    public Optional<SdkTokenProvider> tokenProvider() {
+        return Optional.ofNullable(tokenProvider);
+    }
+
     @Override
     public Builder toBuilder() {
         return new BuilderImpl(this);
@@ -84,14 +97,16 @@ public boolean equals(Object o) {
             return false;
         }
         AwsRequestOverrideConfiguration that = (AwsRequestOverrideConfiguration) o;
-        return Objects.equals(credentialsProvider, that.credentialsProvider);
+        return Objects.equals(credentialsProvider, that.credentialsProvider)
+               && Objects.equals(tokenProvider, that.tokenProvider);
     }
 
     @Override
     public int hashCode() {
         int hashCode = 1;
         hashCode = 31 * hashCode + super.hashCode();
         hashCode = 31 * hashCode + Objects.hashCode(credentialsProvider);
+        hashCode = 31 * hashCode + Objects.hashCode(tokenProvider);
         return hashCode;
     }
 
@@ -113,14 +128,31 @@ public interface Builder extends RequestOverrideConfiguration.Builder<Builder>,
          */
         AwsCredentialsProvider credentialsProvider();
 
+        /**
+         * Set the optional {@link SdkTokenProvider} that will provide a token to be used to authorize this request.
+         * This will be used only if the requested operation uses bearer token authorization.
+         *
+         * @param tokenProvider The {@link SdkTokenProvider}.
+         * @return This object for chaining.
+         */
+        Builder tokenProvider(SdkTokenProvider tokenProvider);
+
+        /**
+         * Return the optional {@link SdkTokenProvider} that will provide a token to be used to authorize this request.
+         * This will be used only if the requested operation uses bearer token authorization.
+         *
+         * @return The optional {@link AwsCredentialsProvider}.
+         */
+        SdkTokenProvider tokenProvider();
+
         @Override
         AwsRequestOverrideConfiguration build();
     }
 
     private static final class BuilderImpl extends RequestOverrideConfiguration.BuilderImpl<Builder> implements Builder {
 
         private AwsCredentialsProvider awsCredentialsProvider;
-
+        private SdkTokenProvider sdkTokenProvider;
 
         private BuilderImpl() {
         }
@@ -132,6 +164,7 @@ private BuilderImpl(RequestOverrideConfiguration requestOverrideConfiguration) {
         private BuilderImpl(AwsRequestOverrideConfiguration awsRequestOverrideConfig) {
             super(awsRequestOverrideConfig);
             this.awsCredentialsProvider = awsRequestOverrideConfig.credentialsProvider;
+            this.sdkTokenProvider = awsRequestOverrideConfig.tokenProvider;
         }
 
         @Override
@@ -145,6 +178,17 @@ public AwsCredentialsProvider credentialsProvider() {
             return awsCredentialsProvider;
         }
 
+        @Override
+        public Builder tokenProvider(SdkTokenProvider tokenProvider) {
+            this.sdkTokenProvider = tokenProvider;
+            return this;
+        }
+
+        @Override
+        public SdkTokenProvider tokenProvider() {
+            return sdkTokenProvider;
+        }
+
         @Override
         public AwsRequestOverrideConfiguration build() {
             return new AwsRequestOverrideConfiguration(this);

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

@@ -20,6 +20,7 @@
 import software.amazon.awssdk.auth.token.credentials.SdkToken;
 import software.amazon.awssdk.auth.token.credentials.SdkTokenProvider;
 import software.amazon.awssdk.auth.token.signer.SdkTokenExecutionAttribute;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.core.RequestOverrideConfiguration;
 import software.amazon.awssdk.core.SdkRequest;
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
@@ -71,10 +72,20 @@ public Signer resolveSigner() {
      */
     @Override
     public void addCredentialsToExecutionAttributes(ExecutionAttributes executionAttributes) {
-        SdkToken credentials = resolveToken(defaultTokenProvider, metricCollector);
+        SdkTokenProvider tokenProvider = resolveTokenProvider(request, defaultTokenProvider);
+        SdkToken credentials = resolveToken(tokenProvider, metricCollector);
         executionAttributes.putAttribute(SdkTokenExecutionAttribute.SDK_TOKEN, credentials);
     }
 
+    private static SdkTokenProvider resolveTokenProvider(SdkRequest originalRequest,
+                                                         SdkTokenProvider defaultProvider) {
+        return originalRequest.overrideConfiguration()
+                              .filter(c -> c instanceof AwsRequestOverrideConfiguration)
+                              .map(c -> (AwsRequestOverrideConfiguration) c)
+                              .flatMap(AwsRequestOverrideConfiguration::tokenProvider)
+                              .orElse(defaultProvider);
+    }
+
     private static SdkToken resolveToken(SdkTokenProvider tokenProvider, MetricCollector metricCollector) {
         Validate.notNull(tokenProvider, "No token provider exists to resolve a token from.");
 

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

@@ -30,6 +30,7 @@
 import software.amazon.awssdk.auth.token.signer.SdkTokenExecutionAttribute;
 import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.awscore.internal.token.TestToken;
+import software.amazon.awssdk.core.RequestOverrideConfiguration;
 import software.amazon.awssdk.core.SdkRequest;
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 import software.amazon.awssdk.core.signer.Signer;
@@ -39,19 +40,25 @@
 public class TokenAuthorizationStrategyTest {
 
     private static final String TOKEN_VALUE = "token_value";
+    private static final String REQUEST_OVERRIDE_TOKEN_VALUE = "request_override_token_value";
+
     private SdkToken token;
+    private SdkToken requestOverrideToken;
 
     @Mock SdkRequest sdkRequest;
     @Mock Signer defaultSigner;
     @Mock Signer requestOverrideSigner;
     @Mock SdkTokenProvider tokenProvider;
+    @Mock SdkTokenProvider requestOverrideTokenProvider;
     @Mock MetricCollector metricCollector;
 
     @Before
     public void setUp() throws Exception {
         token = TestToken.builder().token(TOKEN_VALUE).build();
+        requestOverrideToken = TestToken.builder().token(REQUEST_OVERRIDE_TOKEN_VALUE).build();
         when(sdkRequest.overrideConfiguration()).thenReturn(Optional.empty());
         when(tokenProvider.resolveToken()).thenReturn(token);
+        when(requestOverrideTokenProvider.resolveToken()).thenReturn(requestOverrideToken);
     }
 
     @Test
@@ -68,7 +75,7 @@ public void noOverrideSigner_returnsDefaultSigner() {
 
     @Test
     public void overrideSigner_returnsOverrideSigner() {
-        Optional cfg = Optional.of(requestOverrideConfiguration());
+        Optional cfg = Optional.of(requestOverrideSignerConfiguration());
         when(sdkRequest.overrideConfiguration()).thenReturn(cfg);
         TokenAuthorizationStrategy authorizationContext = TokenAuthorizationStrategy.builder()
                                                                                     .request(sdkRequest)
@@ -93,7 +100,7 @@ public void noDefaultSignerNoOverride_returnsNull() {
     }
 
     @Test
-    public void providerExists_credentialsAddedToExecutionAttributes() {
+    public void noOverrideProvider_tokenAddedToExecutionAttributes() {
         TokenAuthorizationStrategy authorizationContext = TokenAuthorizationStrategy.builder()
                                                                                     .request(sdkRequest)
                                                                                     .defaultSigner(defaultSigner)
@@ -105,6 +112,21 @@ public void providerExists_credentialsAddedToExecutionAttributes() {
         assertThat(executionAttributes.getAttribute(SdkTokenExecutionAttribute.SDK_TOKEN)).isEqualTo(token);
     }
 
+    @Test
+    public void overrideProvider_overrideTokenAddedToExecutionAttributes() {
+        Optional cfg = Optional.of(requestOverrideTokenProviderConfiguration());
+        when(sdkRequest.overrideConfiguration()).thenReturn(cfg);
+        TokenAuthorizationStrategy authorizationContext = TokenAuthorizationStrategy.builder()
+                                                                                    .request(sdkRequest)
+                                                                                    .defaultSigner(defaultSigner)
+                                                                                    .defaultTokenProvider(tokenProvider)
+                                                                                    .metricCollector(metricCollector)
+                                                                                    .build();
+        ExecutionAttributes executionAttributes = new ExecutionAttributes();
+        authorizationContext.addCredentialsToExecutionAttributes(executionAttributes);
+        assertThat(executionAttributes.getAttribute(SdkTokenExecutionAttribute.SDK_TOKEN)).isEqualTo(requestOverrideToken);
+    }
+
     @Test
     public void noProvider_throwsError() {
         TokenAuthorizationStrategy authorizationContext = TokenAuthorizationStrategy.builder()
@@ -119,9 +141,15 @@ public void noProvider_throwsError() {
             .hasMessageContaining("No token provider exists to resolve a token from.");
     }
 
-    private AwsRequestOverrideConfiguration requestOverrideConfiguration() {
+    private AwsRequestOverrideConfiguration requestOverrideSignerConfiguration() {
         return AwsRequestOverrideConfiguration.builder()
                                               .signer(requestOverrideSigner)
                                               .build();
     }
+
+    private AwsRequestOverrideConfiguration requestOverrideTokenProviderConfiguration() {
+        return AwsRequestOverrideConfiguration.builder()
+                                              .tokenProvider(requestOverrideTokenProvider)
+                                              .build();
+    }
 }