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._
This commit is contained in:
John DiSanti 2023-06-14 13:05:16 -07:00 committed by GitHub
parent 9d56db1f3e
commit a694dd9dcd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
8 changed files with 123 additions and 31 deletions

View File

@ -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"] }

View File

@ -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::<SigV4OperationSigningConfig>().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::<SigV4OperationSigningConfig>(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<FrozenLayer> {
let mut layer = Layer::new("Presigning");
layer.set_retry_strategy(NeverRetryStrategy::new());
layer.put(disable_interceptor::<InvocationIdInterceptor>("presigning"));
layer.put(disable_interceptor::<RequestInfoInterceptor>("presigning"));
layer.put(disable_interceptor::<UserAgentInterceptor>("presigning"));

View File

@ -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 {

View File

@ -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,
) {

View File

@ -108,6 +108,7 @@ class ClientCodegenVisitor(
codegenContext,
ClientModuleDocProvider(codegenContext, service.serviceNameOrDefault("the service")),
),
protocolImpl = protocolGeneratorFactory.protocol(codegenContext),
)
rustCrate = RustCrate(

View File

@ -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)) {

View File

@ -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 {}

View File

@ -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"\