diff --git a/services/sts/pom.xml b/services/sts/pom.xml index aefd6760fbc6..cc58373c4151 100644 --- a/services/sts/pom.xml +++ b/services/sts/pom.xml @@ -63,6 +63,12 @@ + + software.amazon.awssdk + retries + ${awsjavasdk.version} + test + iam software.amazon.awssdk diff --git a/services/sts/src/main/java/software/amazon/awssdk/services/sts/internal/StsRetryStrategy.java b/services/sts/src/main/java/software/amazon/awssdk/services/sts/internal/StsRetryStrategy.java new file mode 100644 index 000000000000..f531cb4c94f6 --- /dev/null +++ b/services/sts/src/main/java/software/amazon/awssdk/services/sts/internal/StsRetryStrategy.java @@ -0,0 +1,68 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.services.sts.internal; + +import software.amazon.awssdk.annotations.SdkInternalApi; +import software.amazon.awssdk.awscore.retry.AwsRetryStrategy; +import software.amazon.awssdk.core.client.config.SdkClientConfiguration; +import software.amazon.awssdk.core.client.config.SdkClientOption; +import software.amazon.awssdk.core.retry.NewRetries2026Resolver; +import software.amazon.awssdk.core.retry.RetryMode; +import software.amazon.awssdk.retries.api.RetryStrategy; +import software.amazon.awssdk.services.sts.model.IdpCommunicationErrorException; + +/** + * Specialized retry strategy resolution for STS to enable retrying for {@link IdpCommunicationErrorException}. + */ +@SdkInternalApi +public final class StsRetryStrategy { + + private StsRetryStrategy() { + } + + public static RetryStrategy resolveRetryStrategy(SdkClientConfiguration config) { + RetryStrategy configuredRetryStrategy = config.option(SdkClientOption.RETRY_STRATEGY); + if (configuredRetryStrategy != null) { + return configuredRetryStrategy; + } + + // Just return null and let the normal retry strategy resolution occur + if (!isNewRetries2026Enabled(config)) { + return null; + } + + RetryMode retryMode = resolveRetryMode(config); + + return AwsRetryStrategy.forRetryMode(retryMode, true) + .toBuilder() + .retryOnException(IdpCommunicationErrorException.class) + .build(); + } + + private static RetryMode resolveRetryMode(SdkClientConfiguration config) { + return RetryMode.resolver() + .profileFile(config.option(SdkClientOption.PROFILE_FILE_SUPPLIER)) + .profileName(config.option(SdkClientOption.PROFILE_NAME)) + .defaultRetryMode(config.option(SdkClientOption.DEFAULT_RETRY_MODE)) + .defaultNewRetries2026(config.option(SdkClientOption.DEFAULT_NEW_RETRIES_2026)) + .resolve(); + } + + private static boolean isNewRetries2026Enabled(SdkClientConfiguration config) { + Boolean defaultNewRetries2026 = config.option(SdkClientOption.DEFAULT_NEW_RETRIES_2026); + return new NewRetries2026Resolver().defaultNewRetries2026(defaultNewRetries2026).resolve(); + } +} \ No newline at end of file diff --git a/services/sts/src/main/resources/codegen-resources/customization.config b/services/sts/src/main/resources/codegen-resources/customization.config index a48f004b088a..26879092eaa6 100644 --- a/services/sts/src/main/resources/codegen-resources/customization.config +++ b/services/sts/src/main/resources/codegen-resources/customization.config @@ -21,6 +21,6 @@ "UseGlobalEndpoint with legacy region `us-west-2`": "V2 does not support setting UseGlobalEndpoint. It's regional only/by default" }, - "enableGenerateCompiledEndpointRules": true - + "enableGenerateCompiledEndpointRules": true, + "customRetryStrategy" : "software.amazon.awssdk.services.sts.internal.StsRetryStrategy" } diff --git a/services/sts/src/test/java/software/amazon/awssdk/services/sts/internal/StsRetryStrategyTest.java b/services/sts/src/test/java/software/amazon/awssdk/services/sts/internal/StsRetryStrategyTest.java new file mode 100644 index 000000000000..8535205bf933 --- /dev/null +++ b/services/sts/src/test/java/software/amazon/awssdk/services/sts/internal/StsRetryStrategyTest.java @@ -0,0 +1,111 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.services.sts.internal; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.mock; + +import java.time.Duration; +import java.util.stream.Stream; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import software.amazon.awssdk.awscore.exception.AwsErrorDetails; +import software.amazon.awssdk.core.SdkSystemSetting; +import software.amazon.awssdk.core.client.config.SdkClientConfiguration; +import software.amazon.awssdk.core.client.config.SdkClientOption; +import software.amazon.awssdk.retries.AdaptiveRetryStrategy; +import software.amazon.awssdk.retries.LegacyRetryStrategy; +import software.amazon.awssdk.retries.StandardRetryStrategy; +import software.amazon.awssdk.retries.api.AcquireInitialTokenRequest; +import software.amazon.awssdk.retries.api.RefreshRetryTokenRequest; +import software.amazon.awssdk.retries.api.RetryStrategy; +import software.amazon.awssdk.retries.api.RetryToken; +import software.amazon.awssdk.services.sts.model.IdpCommunicationErrorException; + +public class StsRetryStrategyTest { + @Test + void resolveRetryStrategy_preexistingStrategy_returnsPreexisting() { + RetryStrategy strategy = mock(RetryStrategy.class); + SdkClientConfiguration config = SdkClientConfiguration.builder() + .option(SdkClientOption.RETRY_STRATEGY, strategy) + .build(); + + assertThat(StsRetryStrategy.resolveRetryStrategy(config)).isSameAs(strategy); + } + + @Test + void resolveRetryStrategy_defaultNewRetries2026False_returnsNull() { + SdkClientConfiguration config = SdkClientConfiguration.builder() + .option(SdkClientOption.DEFAULT_NEW_RETRIES_2026, false) + .build(); + + assertThat(StsRetryStrategy.resolveRetryStrategy(config)).isNull(); + } + + @Test + void resolveRetryStrategy_newRetries2026False_returnsNull() { + System.setProperty(SdkSystemSetting.AWS_NEW_RETRIES_2026.property(), "false"); + try { + SdkClientConfiguration config = SdkClientConfiguration.builder().build(); + assertThat(StsRetryStrategy.resolveRetryStrategy(config)).isNull(); + } finally { + System.clearProperty(SdkSystemSetting.AWS_NEW_RETRIES_2026.property()); + } + } + + @ParameterizedTest + @MethodSource("retryModeResolutionCases") + void resolveRetryStrategy_returnsCorrectStrategyBasedOnMode(String mode, Class expected) { + System.setProperty(SdkSystemSetting.AWS_RETRY_MODE.property(), mode); + System.setProperty(SdkSystemSetting.AWS_NEW_RETRIES_2026.property(), "true"); + try { + SdkClientConfiguration config = SdkClientConfiguration.builder().build(); + RetryStrategy resolved = StsRetryStrategy.resolveRetryStrategy(config); + assertThat(resolved).isInstanceOf(expected); + assertRetriesOnIdpCommunicationException(resolved); + } finally { + System.clearProperty(SdkSystemSetting.AWS_RETRY_MODE.property()); + System.clearProperty(SdkSystemSetting.AWS_NEW_RETRIES_2026.property()); + } + } + + void assertRetriesOnIdpCommunicationException(RetryStrategy strategy) { + RetryToken token = strategy.acquireInitialToken(AcquireInitialTokenRequest.create("test")).token(); + + AwsErrorDetails errorDetails = AwsErrorDetails.builder() + .errorCode("IDPCommunicationError") + .build(); + + IdpCommunicationErrorException failure = IdpCommunicationErrorException.builder() + .awsErrorDetails(errorDetails) + .build(); + RefreshRetryTokenRequest refresh = RefreshRetryTokenRequest.builder() + .token(token) + .failure(failure) + .build(); + assertThat(strategy.refreshRetryToken(refresh).delay()).isGreaterThan(Duration.ZERO); + } + + private static Stream retryModeResolutionCases() { + return Stream.of( + Arguments.of("standard", StandardRetryStrategy.class), + Arguments.of("legacy", LegacyRetryStrategy.class), + Arguments.of("adaptive", AdaptiveRetryStrategy.class) + ); + } +}