From 56f4be3207a33df743edb9115a3c9b2f5990f83e Mon Sep 17 00:00:00 2001 From: John DiSanti Date: Fri, 9 Sep 2022 17:25:43 -0700 Subject: [PATCH] Add customization to retry STS `IDPCommunicationErrorException` (#1718) --- CHANGELOG.next.toml | 6 +++ .../smithy/rustsdk/AwsCodegenDecorator.kt | 6 ++- .../rustsdk/customize/sts/STSDecorator.kt | 50 +++++++++++++++++++ aws/sdk/integration-tests/sts/Cargo.toml | 1 + .../sts/tests/retry_idp_comms_err.rs | 32 ++++++++++++ 5 files changed, 93 insertions(+), 2 deletions(-) create mode 100644 aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/sts/STSDecorator.kt create mode 100644 aws/sdk/integration-tests/sts/tests/retry_idp_comms_err.rs diff --git a/CHANGELOG.next.toml b/CHANGELOG.next.toml index e5dcda5393..4e13f28cd6 100644 --- a/CHANGELOG.next.toml +++ b/CHANGELOG.next.toml @@ -186,3 +186,9 @@ message = "Smithy IDL v2 mixins are now supported" references = ["smithy-rs#1680"] meta = { "breaking" = false, "tada" = false, "bug" = false, "target" = "all"} author = "ogudavid" + +[[aws-sdk-rust]] +message = "The AWS STS SDK now automatically retries `IDPCommunicationError` when calling `AssumeRoleWithWebIdentity`" +references = ["smithy-rs#966", "smithy-rs#1718"] +meta = { "breaking" = false, "tada" = false, "bug" = true } +author = "jdisanti" diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCodegenDecorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCodegenDecorator.kt index c83e3822a4..d228a60949 100644 --- a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCodegenDecorator.kt +++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCodegenDecorator.kt @@ -15,6 +15,7 @@ import software.amazon.smithy.rustsdk.customize.ec2.Ec2Decorator import software.amazon.smithy.rustsdk.customize.glacier.GlacierDecorator import software.amazon.smithy.rustsdk.customize.route53.Route53Decorator import software.amazon.smithy.rustsdk.customize.s3.S3Decorator +import software.amazon.smithy.rustsdk.customize.sts.STSDecorator val DECORATORS = listOf( // General AWS Decorators @@ -35,12 +36,13 @@ val DECORATORS = listOf( AwsReadmeDecorator(), // Service specific decorators - DisabledAuthDecorator(), ApiGatewayDecorator(), - S3Decorator(), + DisabledAuthDecorator(), Ec2Decorator(), GlacierDecorator(), Route53Decorator(), + S3Decorator(), + STSDecorator(), // Only build docs-rs for linux to reduce load on docs.rs DocsRsMetadataDecorator(DocsRsMetadataSettings(targets = listOf("x86_64-unknown-linux-gnu"), allFeatures = true)), diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/sts/STSDecorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/sts/STSDecorator.kt new file mode 100644 index 0000000000..6dc1ccb36d --- /dev/null +++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/sts/STSDecorator.kt @@ -0,0 +1,50 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package software.amazon.smithy.rustsdk.customize.sts + +import software.amazon.smithy.model.Model +import software.amazon.smithy.model.shapes.ServiceShape +import software.amazon.smithy.model.shapes.Shape +import software.amazon.smithy.model.shapes.ShapeId +import software.amazon.smithy.model.shapes.StructureShape +import software.amazon.smithy.model.traits.ErrorTrait +import software.amazon.smithy.model.traits.RetryableTrait +import software.amazon.smithy.model.transform.ModelTransformer +import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext +import software.amazon.smithy.rust.codegen.client.smithy.CoreCodegenContext +import software.amazon.smithy.rust.codegen.client.smithy.customize.RustCodegenDecorator +import software.amazon.smithy.rust.codegen.client.smithy.letIf +import software.amazon.smithy.rust.codegen.client.util.hasTrait +import java.util.logging.Logger + +class STSDecorator : RustCodegenDecorator { + override val name: String = "STS" + override val order: Byte = 0 + private val logger: Logger = Logger.getLogger(javaClass.name) + + private fun applies(serviceId: ShapeId) = + serviceId == ShapeId.from("com.amazonaws.sts#AWSSecurityTokenServiceV20110615") + + private fun isIdpCommunicationError(shape: Shape): Boolean = + shape is StructureShape && shape.hasTrait() && + shape.id.namespace == "com.amazonaws.sts" && shape.id.name == "IDPCommunicationErrorException" + + override fun transformModel(service: ServiceShape, model: Model): Model { + return model.letIf(applies(service.id)) { + ModelTransformer.create().mapShapes(model) { shape -> + shape.letIf(isIdpCommunicationError(shape)) { + logger.info("Adding @retryable trait to $shape and setting its error type to 'server'") + (shape as StructureShape).toBuilder() + .removeTrait(ErrorTrait.ID) + .addTrait(ErrorTrait("server")) + .addTrait(RetryableTrait.builder().build()).build() + } + } + } + } + + override fun supportsCodegenContext(clazz: Class): Boolean = + clazz.isAssignableFrom(ClientCodegenContext::class.java) +} diff --git a/aws/sdk/integration-tests/sts/Cargo.toml b/aws/sdk/integration-tests/sts/Cargo.toml index 2c729f74e1..ead7a6a887 100644 --- a/aws/sdk/integration-tests/sts/Cargo.toml +++ b/aws/sdk/integration-tests/sts/Cargo.toml @@ -11,5 +11,6 @@ edition = "2021" aws-sdk-sts = { path = "../../build/aws-sdk/sdk/sts" } aws-smithy-client = { path = "../../build/aws-sdk/sdk/aws-smithy-client", features = ["test-util", "rustls"] } aws-smithy-http = { path = "../../build/aws-sdk/sdk/aws-smithy-http" } +aws-smithy-types = { path = "../../build/aws-sdk/sdk/aws-smithy-types" } tokio = { version = "1.8.4", features = ["full", "test-util"] } tracing-subscriber = { version = "0.3.15", features = ["env-filter"] } diff --git a/aws/sdk/integration-tests/sts/tests/retry_idp_comms_err.rs b/aws/sdk/integration-tests/sts/tests/retry_idp_comms_err.rs new file mode 100644 index 0000000000..6fe9895cd3 --- /dev/null +++ b/aws/sdk/integration-tests/sts/tests/retry_idp_comms_err.rs @@ -0,0 +1,32 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use aws_sdk_sts as sts; +use aws_smithy_types::error::Error as ErrorMeta; +use aws_smithy_types::retry::{ErrorKind, ProvideErrorKind}; +use sts::error::{ + AssumeRoleWithWebIdentityError, AssumeRoleWithWebIdentityErrorKind, + IdpCommunicationErrorException, +}; + +#[tokio::test] +async fn idp_comms_err_retryable() { + let error = AssumeRoleWithWebIdentityError::new( + AssumeRoleWithWebIdentityErrorKind::IdpCommunicationErrorException( + IdpCommunicationErrorException::builder() + .message("test") + .build(), + ), + ErrorMeta::builder() + .code("IDPCommunicationError") + .message("test") + .build(), + ); + assert_eq!( + Some(ErrorKind::ServerError), + error.retryable_error_kind(), + "IdpCommunicationErrorException should be a retryable server error" + ); +}