Add customization to retry STS `IDPCommunicationErrorException` (#1718) (56f4be32) · Commits · Public Repositories / Smithy Rust

CHANGELOG.next.toml

+6 −0

Original line number	Diff line number	Diff line
		@@ -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"

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCodegenDecorator.kt

+4 −2

Original line number	Diff line number	Diff line
		@@ -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)),

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/sts/STSDecorator.kt

0 → 100644

+50 −0

Original line number	Diff line number	Diff line
		/*
		* 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<ClientCodegenContext> {
		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<ErrorTrait>() &&
		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<out CoreCodegenContext>): Boolean =
		clazz.isAssignableFrom(ClientCodegenContext::class.java)
		}

aws/sdk/integration-tests/sts/Cargo.toml

+1 −0

Original line number	Diff line number	Diff line
		@@ -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"] }

aws/sdk/integration-tests/sts/tests/retry_idp_comms_err.rs

0 → 100644

+32 −0

Original line number	Diff line number	Diff line
		/*
		* 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"
		);
		}