From a694dd9dcd028d3c4ddd4a90091c1a994da9d008 Mon Sep 17 00:00:00 2001 From: John DiSanti Date: Wed, 14 Jun 2023 13:05:16 -0700 Subject: [PATCH] Fix Polly presigning in the orchestrator (#2769) This PR fixes presigning for Amazon Polly in the orchestrator implementation. ---- _By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice._ --- aws/rust-runtime/aws-inlineable/Cargo.toml | 2 +- .../src/presigning_interceptors.rs | 21 +++-- .../smithy/rustsdk/AwsPresigningDecorator.kt | 83 +++++++++++++++++-- .../client/smithy/ClientCodegenContext.kt | 2 + .../client/smithy/ClientCodegenVisitor.kt | 1 + .../protocol/RequestSerializerGenerator.kt | 42 ++++++---- .../src/client/retries/strategy/never.rs | 1 + .../check-aws-sdk-orchestrator-impl | 2 +- 8 files changed, 123 insertions(+), 31 deletions(-) diff --git a/aws/rust-runtime/aws-inlineable/Cargo.toml b/aws/rust-runtime/aws-inlineable/Cargo.toml index d941f059b..e061a335e 100644 --- a/aws/rust-runtime/aws-inlineable/Cargo.toml +++ b/aws/rust-runtime/aws-inlineable/Cargo.toml @@ -23,6 +23,7 @@ aws-smithy-client = { path = "../../../rust-runtime/aws-smithy-client" } aws-smithy-http = { path = "../../../rust-runtime/aws-smithy-http" } aws-smithy-http-tower = { path = "../../../rust-runtime/aws-smithy-http-tower" } aws-smithy-runtime-api = { path = "../../../rust-runtime/aws-smithy-runtime-api" } +aws-smithy-runtime = { path = "../../../rust-runtime/aws-smithy-runtime" } aws-smithy-types = { path = "../../../rust-runtime/aws-smithy-types" } aws-smithy-async = { path = "../../../rust-runtime/aws-smithy-async" } aws-types = { path = "../aws-types" } @@ -42,7 +43,6 @@ tracing = "0.1" aws-credential-types = { path = "../aws-credential-types", features = ["test-util"] } aws-smithy-client = { path = "../../../rust-runtime/aws-smithy-client", features = ["test-util"] } aws-smithy-http = { path = "../../../rust-runtime/aws-smithy-http", features = ["rt-tokio"] } -aws-smithy-runtime = { path = "../../../rust-runtime/aws-smithy-runtime" } tempfile = "3.6.0" tracing-subscriber = { version = "0.3.17", features = ["env-filter"] } aws-smithy-async = { path = "../../../rust-runtime/aws-smithy-async", features = ["test-util"] } diff --git a/aws/rust-runtime/aws-inlineable/src/presigning_interceptors.rs b/aws/rust-runtime/aws-inlineable/src/presigning_interceptors.rs index 55642794a..8d67e3ae6 100644 --- a/aws/rust-runtime/aws-inlineable/src/presigning_interceptors.rs +++ b/aws/rust-runtime/aws-inlineable/src/presigning_interceptors.rs @@ -11,7 +11,9 @@ use aws_runtime::auth::sigv4::{HttpSignatureType, SigV4OperationSigningConfig}; use aws_runtime::invocation_id::InvocationIdInterceptor; use aws_runtime::request_info::RequestInfoInterceptor; use aws_runtime::user_agent::UserAgentInterceptor; +use aws_sigv4::http_request::SignableBody; use aws_smithy_async::time::{SharedTimeSource, StaticTimeSource}; +use aws_smithy_runtime::client::retries::strategy::NeverRetryStrategy; use aws_smithy_runtime_api::client::interceptors::{ disable_interceptor, BeforeSerializationInterceptorContextMut, BeforeTransmitInterceptorContextMut, BoxError, Interceptor, InterceptorRegistrar, @@ -26,11 +28,15 @@ use aws_smithy_types::config_bag::{ConfigBag, FrozenLayer, Layer}; #[derive(Debug)] pub(crate) struct SigV4PresigningInterceptor { config: PresigningConfig, + payload_override: SignableBody<'static>, } impl SigV4PresigningInterceptor { - pub(crate) fn new(config: PresigningConfig) -> Self { - Self { config } + pub(crate) fn new(config: PresigningConfig, payload_override: SignableBody<'static>) -> Self { + Self { + config, + payload_override, + } } } @@ -60,8 +66,7 @@ impl Interceptor for SigV4PresigningInterceptor { if let Some(mut config) = cfg.get::().cloned() { config.signing_options.expires_in = Some(self.config.expires()); config.signing_options.signature_type = HttpSignatureType::HttpRequestQueryParams; - config.signing_options.payload_override = - Some(aws_sigv4::http_request::SignableBody::UnsignedPayload); + config.signing_options.payload_override = Some(self.payload_override.clone()); cfg.interceptor_state() .put::(config); Ok(()) @@ -81,9 +86,12 @@ pub(crate) struct SigV4PresigningRuntimePlugin { } impl SigV4PresigningRuntimePlugin { - pub(crate) fn new(config: PresigningConfig) -> Self { + pub(crate) fn new(config: PresigningConfig, payload_override: SignableBody<'static>) -> Self { Self { - interceptor: SharedInterceptor::new(SigV4PresigningInterceptor::new(config)), + interceptor: SharedInterceptor::new(SigV4PresigningInterceptor::new( + config, + payload_override, + )), } } } @@ -91,6 +99,7 @@ impl SigV4PresigningRuntimePlugin { impl RuntimePlugin for SigV4PresigningRuntimePlugin { fn config(&self) -> Option { let mut layer = Layer::new("Presigning"); + layer.set_retry_strategy(NeverRetryStrategy::new()); layer.put(disable_interceptor::("presigning")); layer.put(disable_interceptor::("presigning")); layer.put(disable_interceptor::("presigning")); diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsPresigningDecorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsPresigningDecorator.kt index bbf5f80fb..6c85ce361 100644 --- a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsPresigningDecorator.kt +++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsPresigningDecorator.kt @@ -23,17 +23,20 @@ import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationSec import software.amazon.smithy.rust.codegen.client.smithy.generators.client.FluentClientCustomization import software.amazon.smithy.rust.codegen.client.smithy.generators.client.FluentClientSection import software.amazon.smithy.rust.codegen.client.smithy.generators.protocol.MakeOperationGenerator +import software.amazon.smithy.rust.codegen.client.smithy.generators.protocol.RequestSerializerGenerator import software.amazon.smithy.rust.codegen.client.smithy.protocols.ClientHttpBoundProtocolPayloadGenerator import software.amazon.smithy.rust.codegen.core.rustlang.CargoDependency import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter import software.amazon.smithy.rust.codegen.core.rustlang.Writable import software.amazon.smithy.rust.codegen.core.rustlang.docs +import software.amazon.smithy.rust.codegen.core.rustlang.rust import software.amazon.smithy.rust.codegen.core.rustlang.rustBlock import software.amazon.smithy.rust.codegen.core.rustlang.rustBlockTemplate import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate import software.amazon.smithy.rust.codegen.core.rustlang.withBlock import software.amazon.smithy.rust.codegen.core.rustlang.writable import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType +import software.amazon.smithy.rust.codegen.core.smithy.contextName import software.amazon.smithy.rust.codegen.core.util.cloneOperation import software.amazon.smithy.rust.codegen.core.util.expectTrait import software.amazon.smithy.rust.codegen.core.util.hasTrait @@ -305,14 +308,24 @@ class AwsPresignedFluentBuilderMethod( } private fun RustWriter.renderPresignedMethodBody(section: FluentClientSection.FluentBuilderImpl) { + val presignableOp = PRESIGNABLE_OPERATIONS.getValue(section.operationShape.id) + val operationShape = if (presignableOp.hasModelTransforms()) { + codegenContext.model.expectShape(syntheticShapeId(section.operationShape.id), OperationShape::class.java) + } else { + section.operationShape + } + rustTemplate( """ + #{alternate_presigning_serializer} + let runtime_plugins = #{Operation}::register_runtime_plugins( - #{RuntimePlugins}::new() - .with_client_plugin(#{SigV4PresigningRuntimePlugin}::new(presigning_config)), + #{RuntimePlugins}::new(), self.handle.clone(), - self.config_override, - ); + self.config_override + ) + .with_client_plugin(#{SigV4PresigningRuntimePlugin}::new(presigning_config, #{payload_override})) + #{alternate_presigning_serializer_registration}; let input = self.inner.build().map_err(#{SdkError}::construction_failure)?; let mut context = #{Operation}::orchestrate_with_stop_point(&runtime_plugins, input, #{StopPoint}::BeforeTransmit) @@ -332,14 +345,74 @@ class AwsPresignedFluentBuilderMethod( "OperationError" to section.operationErrorType, "RuntimePlugins" to RuntimeType.smithyRuntimeApi(runtimeConfig) .resolve("client::runtime_plugin::RuntimePlugins"), - "SharedInterceptor" to RuntimeType.smithyRuntimeApi(runtimeConfig).resolve("client::interceptors").resolve("SharedInterceptor"), + "SharedInterceptor" to RuntimeType.smithyRuntimeApi(runtimeConfig).resolve("client::interceptors") + .resolve("SharedInterceptor"), "SigV4PresigningRuntimePlugin" to AwsRuntimeType.presigningInterceptor(runtimeConfig) .resolve("SigV4PresigningRuntimePlugin"), "StopPoint" to RuntimeType.smithyRuntime(runtimeConfig).resolve("client::orchestrator::StopPoint"), "TypedBox" to RuntimeType.smithyTypes(runtimeConfig).resolve("type_erasure::TypedBox"), "USER_AGENT" to CargoDependency.Http.toType().resolve("header::USER_AGENT"), + "alternate_presigning_serializer" to writable { + if (presignableOp.hasModelTransforms()) { + val smithyTypes = RuntimeType.smithyTypes(codegenContext.runtimeConfig) + rustTemplate( + """ + ##[derive(Debug)] + struct AlternatePresigningSerializerRuntimePlugin; + impl #{RuntimePlugin} for AlternatePresigningSerializerRuntimePlugin { + fn config(&self) -> Option<#{FrozenLayer}> { + use #{ConfigBagAccessors}; + let mut cfg = #{Layer}::new("presigning_serializer"); + cfg.set_request_serializer(#{AlternateSerializer}); + Some(cfg.freeze()) + } + } + """, + "AlternateSerializer" to alternateSerializer(operationShape), + "ConfigBagAccessors" to RuntimeType.smithyRuntimeApi(codegenContext.runtimeConfig) + .resolve("client::orchestrator::ConfigBagAccessors"), + "FrozenLayer" to smithyTypes.resolve("config_bag::FrozenLayer"), + "Layer" to smithyTypes.resolve("config_bag::Layer"), + "RuntimePlugin" to RuntimeType.runtimePlugin(codegenContext.runtimeConfig), + ) + } + }, + "alternate_presigning_serializer_registration" to writable { + if (presignableOp.hasModelTransforms()) { + rust(".with_operation_plugin(AlternatePresigningSerializerRuntimePlugin)") + } + }, + "payload_override" to writable { + rustTemplate( + "#{aws_sigv4}::http_request::SignableBody::" + + when (presignableOp.payloadSigningType) { + PayloadSigningType.EMPTY -> "Bytes(b\"\")" + PayloadSigningType.UNSIGNED_PAYLOAD -> "UnsignedPayload" + }, + "aws_sigv4" to AwsRuntimeType.awsSigv4(runtimeConfig), + ) + }, ) } + + private fun alternateSerializer(transformedOperationShape: OperationShape): RuntimeType = + transformedOperationShape.contextName(codegenContext.serviceShape).replaceFirstChar { + it.uppercase() + }.let { baseName -> + "${baseName}PresigningRequestSerializer".let { name -> + RuntimeType.forInlineFun(name, codegenContext.symbolProvider.moduleForShape(transformedOperationShape)) { + RequestSerializerGenerator( + codegenContext, + codegenContext.protocolImpl!!, + null, + nameOverride = name, + ).render( + this, + transformedOperationShape, + ) + } + } + } } interface PresignModelTransform { diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenContext.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenContext.kt index 92a8a9f8c..18db61982 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenContext.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenContext.kt @@ -13,6 +13,7 @@ import software.amazon.smithy.rust.codegen.core.smithy.CodegenContext import software.amazon.smithy.rust.codegen.core.smithy.CodegenTarget import software.amazon.smithy.rust.codegen.core.smithy.ModuleDocProvider import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider +import software.amazon.smithy.rust.codegen.core.smithy.protocols.Protocol /** * [ClientCodegenContext] contains code-generation context that is _specific_ to the [RustClientCodegenPlugin] plugin @@ -30,6 +31,7 @@ data class ClientCodegenContext( // Expose the `rootDecorator`, enabling customizations to compose by referencing information from the root codegen // decorator val rootDecorator: ClientCodegenDecorator, + val protocolImpl: Protocol? = null, ) : CodegenContext( model, symbolProvider, moduleDocProvider, serviceShape, protocol, settings, CodegenTarget.CLIENT, ) { diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenVisitor.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenVisitor.kt index 99bb528fe..7583bc342 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenVisitor.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenVisitor.kt @@ -108,6 +108,7 @@ class ClientCodegenVisitor( codegenContext, ClientModuleDocProvider(codegenContext, service.serviceNameOrDefault("the service")), ), + protocolImpl = protocolGeneratorFactory.protocol(codegenContext), ) rustCrate = RustCrate( diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/protocol/RequestSerializerGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/protocol/RequestSerializerGenerator.kt index f3b53c360..99adee16b 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/protocol/RequestSerializerGenerator.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/protocol/RequestSerializerGenerator.kt @@ -29,7 +29,8 @@ import software.amazon.smithy.rust.codegen.core.util.inputShape class RequestSerializerGenerator( private val codegenContext: ClientCodegenContext, private val protocol: Protocol, - private val bodyGenerator: ProtocolPayloadGenerator, + private val bodyGenerator: ProtocolPayloadGenerator?, + private val nameOverride: String? = null, ) { private val httpBindingResolver = protocol.httpBindingResolver private val symbolProvider = codegenContext.symbolProvider @@ -63,11 +64,12 @@ class RequestSerializerGenerator( val inputShape = operationShape.inputShape(codegenContext.model) val operationName = symbolProvider.toSymbol(operationShape).name val inputSymbol = symbolProvider.toSymbol(inputShape) + val serializerName = nameOverride ?: "${operationName}RequestSerializer" writer.rustTemplate( """ ##[derive(Debug)] - struct ${operationName}RequestSerializer; - impl #{RequestSerializer} for ${operationName}RequestSerializer { + struct $serializerName; + impl #{RequestSerializer} for $serializerName { ##[allow(unused_mut, clippy::let_and_return, clippy::needless_borrow, clippy::useless_conversion)] fn serialize_input(&self, input: #{Input}, _cfg: &mut #{ConfigBag}) -> Result<#{HttpRequest}, #{BoxError}> { let input = #{TypedBox}::<#{ConcreteInput}>::assume_from(input).expect("correct type").unwrap(); @@ -85,22 +87,26 @@ class RequestSerializerGenerator( "ConcreteInput" to inputSymbol, "create_http_request" to createHttpRequest(operationShape), "generate_body" to writable { - val body = writable { - bodyGenerator.generatePayload( - this, - "input", - operationShape, - ClientAdditionalPayloadContext(propertyBagAvailable = false), - ) - } - val streamingMember = inputShape.findStreamingMember(codegenContext.model) - val isBlobStreaming = - streamingMember != null && codegenContext.model.expectShape(streamingMember.target) is BlobShape - if (isBlobStreaming) { - // Consume the `ByteStream` into its inner `SdkBody`. - rust("#T.into_inner()", body) + if (bodyGenerator != null) { + val body = writable { + bodyGenerator.generatePayload( + this, + "input", + operationShape, + ClientAdditionalPayloadContext(propertyBagAvailable = false), + ) + } + val streamingMember = inputShape.findStreamingMember(codegenContext.model) + val isBlobStreaming = + streamingMember != null && codegenContext.model.expectShape(streamingMember.target) is BlobShape + if (isBlobStreaming) { + // Consume the `ByteStream` into its inner `SdkBody`. + rust("#T.into_inner()", body) + } else { + rustTemplate("#{SdkBody}::from(#{body})", *codegenScope, "body" to body) + } } else { - rustTemplate("#{SdkBody}::from(#{body})", *codegenScope, "body" to body) + rustTemplate("#{SdkBody}::empty()", *codegenScope) } }, "add_content_length" to if (needsContentLength(operationShape)) { diff --git a/rust-runtime/aws-smithy-runtime/src/client/retries/strategy/never.rs b/rust-runtime/aws-smithy-runtime/src/client/retries/strategy/never.rs index 599b6ea5a..3f0a2325e 100644 --- a/rust-runtime/aws-smithy-runtime/src/client/retries/strategy/never.rs +++ b/rust-runtime/aws-smithy-runtime/src/client/retries/strategy/never.rs @@ -8,6 +8,7 @@ use aws_smithy_runtime_api::client::orchestrator::BoxError; use aws_smithy_runtime_api::client::retries::{RetryStrategy, ShouldAttempt}; use aws_smithy_types::config_bag::ConfigBag; +#[non_exhaustive] #[derive(Debug, Clone, Default)] pub struct NeverRetryStrategy {} diff --git a/tools/ci-scripts/check-aws-sdk-orchestrator-impl b/tools/ci-scripts/check-aws-sdk-orchestrator-impl index a02eb14ad..e9889a298 100755 --- a/tools/ci-scripts/check-aws-sdk-orchestrator-impl +++ b/tools/ci-scripts/check-aws-sdk-orchestrator-impl @@ -16,7 +16,6 @@ cd smithy-rs services_that_compile=(\ "aws-config"\ "dynamodb"\ - "polly"\ "route53"\ "s3"\ "sts"\ @@ -30,6 +29,7 @@ services_that_pass_tests=(\ "iam"\ "kms"\ "lambda"\ + "polly"\ "qldbsession"\ "s3control"\ "sso"\