Add DefaultEndpointResolver to the orchestrator

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

@@ -8,6 +8,7 @@ package software.amazon.smithy.rustsdk
 import software.amazon.smithy.aws.traits.auth.SigV4Trait
 import software.amazon.smithy.aws.traits.auth.UnsignedPayloadTrait
 import software.amazon.smithy.model.knowledge.ServiceIndex
+import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.model.traits.OptionalAuthTrait
 import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
 import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
@@ -38,6 +39,7 @@ class SigV4AuthDecorator : ClientCodegenDecorator {
 
     override fun operationRuntimePluginCustomizations(
         codegenContext: ClientCodegenContext,
+        operation: OperationShape,
         baseCustomizations: List<OperationRuntimePluginCustomization>,
     ): List<OperationRuntimePluginCustomization> =
         baseCustomizations.letIf(codegenContext.settings.codegenConfig.enableNewSmithyRuntime) {

aws/sra-test/integration-tests/aws-sdk-s3/Cargo.toml

@@ -6,18 +6,18 @@ publish = false
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-aws-credential-types = { path = "../../../rust-runtime/aws-credential-types", features = ["test-util"] }
-aws-http = { path = "../../../rust-runtime/aws-http" }
-aws-runtime = { path = "../../../rust-runtime/aws-runtime" }
-aws-sdk-s3 = { path = "../../build/sdk/aws-sdk-s3", features = ["test-util"] }
-aws-sigv4 = { path = "../../../rust-runtime/aws-sigv4" }
-aws-types = { path = "../../../rust-runtime/aws-types" }
-aws-smithy-async = { path = "../../../../rust-runtime/aws-smithy-async", features = ["rt-tokio"] }
-aws-smithy-client = { path = "../../../../rust-runtime/aws-smithy-client", features = ["test-util"] }
-aws-smithy-types = { path = "../../../../rust-runtime/aws-smithy-types" }
-aws-smithy-http = { path = "../../../../rust-runtime/aws-smithy-http" }
-aws-smithy-runtime = { path = "../../../../rust-runtime/aws-smithy-runtime", features = ["test-util"] }
-aws-smithy-runtime-api = { path = "../../../../rust-runtime/aws-smithy-runtime-api" }
+aws-credential-types = { path = "../../../sdk/build/aws-sdk/sdk/aws-credential-types", features = ["test-util"] }
+aws-http = { path = "../../../sdk/build/aws-sdk/sdk/aws-http" }
+aws-runtime = { path = "../../../sdk/build/aws-sdk/sdk/aws-runtime" }
+aws-sdk-s3 = { path = "../../../sdk/build/aws-sdk/sdk/s3/", features = ["test-util"] }
+aws-sigv4 = { path = "../../../sdk/build/aws-sdk/sdk/aws-sigv4" }
+aws-types = { path = "../../../sdk/build/aws-sdk/sdk/aws-types" }
+aws-smithy-async = { path = "../../../sdk/build/aws-sdk/sdk/aws-smithy-async", features = ["rt-tokio"] }
+aws-smithy-client = { path = "../../../sdk/build/aws-sdk/sdk/aws-smithy-client", features = ["test-util"] }
+aws-smithy-types = { path = "../../../sdk/build/aws-sdk/sdk/aws-smithy-types" }
+aws-smithy-http = { path = "../../../sdk/build/aws-sdk/sdk/aws-smithy-http" }
+aws-smithy-runtime = { path = "../../../sdk/build/aws-sdk/sdk/aws-smithy-runtime", features = ["test-util"] }
+aws-smithy-runtime-api = { path = "../../../sdk/build/aws-sdk/sdk/aws-smithy-runtime-api" }
 tokio = { version = "1.23.1", features = ["macros", "test-util", "rt-multi-thread"] }
 tracing = "0.1.37"
 tracing-subscriber = { version = "0.3.15", features = ["env-filter", "json"] }

aws/sra-test/integration-tests/aws-sdk-s3/tests/sra_test.rs

@@ -17,8 +17,10 @@ use aws_sdk_s3::primitives::SdkBody;
 use aws_smithy_client::erase::DynConnector;
 use aws_smithy_client::test_connection::TestConnection;
 use aws_smithy_runtime::client::connections::adapter::DynConnectorAdapter;
-use aws_smithy_runtime_api::client::endpoints::StaticUriEndpointResolver;
-use aws_smithy_runtime_api::client::interceptors::{Interceptor, InterceptorContext, Interceptors};
+use aws_smithy_runtime_api::client::endpoints::DefaultEndpointResolver;
+use aws_smithy_runtime_api::client::interceptors::{
+    Interceptor, InterceptorContext, InterceptorError, Interceptors,
+};
 use aws_smithy_runtime_api::client::orchestrator::{
     BoxError, ConfigBagAccessors, Connection, HttpRequest, HttpResponse, TraceProbe,
 };
@@ -27,7 +29,6 @@ use aws_smithy_runtime_api::config_bag::ConfigBag;
 use aws_smithy_runtime_api::type_erasure::TypedBox;
 use aws_types::region::SigningRegion;
 use aws_types::SigningService;
-use http::Uri;
 use std::sync::Arc;
 use std::time::{Duration, UNIX_EPOCH};
 
@@ -108,9 +109,16 @@ async fn sra_manual_test() {
                 ),
             );
 
-            cfg.set_endpoint_resolver(StaticUriEndpointResolver::uri(Uri::from_static(
-                "https://test-bucket.s3.us-east-1.amazonaws.com/",
-            )));
+            cfg.set_endpoint_resolver(DefaultEndpointResolver::new(
+                aws_smithy_http::endpoint::SharedEndpointResolver::new(
+                    aws_sdk_s3::endpoint::DefaultResolver::new(),
+                ),
+            ));
+
+            let params_builder = aws_sdk_s3::endpoint::Params::builder()
+                .set_region(Some("us-east-1".to_owned()))
+                .set_endpoint(Some("https://s3.us-east-1.amazonaws.com/".to_owned()));
+            cfg.put(params_builder);
 
             cfg.set_retry_strategy(
                 aws_smithy_runtime_api::client::retries::NeverRetryStrategy::new(),
@@ -162,6 +170,77 @@ async fn sra_manual_test() {
         }
     }
 
+    // This is a temporary operation runtime plugin until <Operation>EndpointParamsInterceptor and
+    // <Operation>EndpointParamsFinalizerInterceptor have been fully implemented, in which case
+    // `.with_operation_plugin(ManualOperationRuntimePlugin)` can be removed.
+    struct ManualOperationRuntimePlugin;
+
+    impl RuntimePlugin for ManualOperationRuntimePlugin {
+        fn configure(&self, cfg: &mut ConfigBag) -> Result<(), BoxError> {
+            #[derive(Debug)]
+            struct ListObjectsV2EndpointParamsInterceptor;
+            impl Interceptor<HttpRequest, HttpResponse> for ListObjectsV2EndpointParamsInterceptor {
+                fn read_before_execution(
+                    &self,
+                    context: &InterceptorContext<HttpRequest, HttpResponse>,
+                    cfg: &mut ConfigBag,
+                ) -> Result<(), BoxError> {
+                    let input = context.input()?;
+                    let input = input
+                        .downcast_ref::<ListObjectsV2Input>()
+                        .ok_or_else(|| InterceptorError::invalid_input_access())?;
+                    let mut params_builder = cfg
+                        .get::<aws_sdk_s3::endpoint::ParamsBuilder>()
+                        .ok_or(InterceptorError::read_before_execution(
+                            "missing endpoint params builder",
+                        ))?
+                        .clone();
+                    params_builder = params_builder.set_bucket(input.bucket.clone());
+                    cfg.put(params_builder);
+
+                    Ok(())
+                }
+            }
+
+            #[derive(Debug)]
+            struct ListObjectsV2EndpointParamsFinalizerInterceptor;
+            impl Interceptor<HttpRequest, HttpResponse> for ListObjectsV2EndpointParamsFinalizerInterceptor {
+                fn read_before_execution(
+                    &self,
+                    _context: &InterceptorContext<HttpRequest, HttpResponse>,
+                    cfg: &mut ConfigBag,
+                ) -> Result<(), BoxError> {
+                    let params_builder = cfg
+                        .get::<aws_sdk_s3::endpoint::ParamsBuilder>()
+                        .ok_or(InterceptorError::read_before_execution(
+                            "missing endpoint params builder",
+                        ))?
+                        .clone();
+                    let params = params_builder
+                        .build()
+                        .map_err(InterceptorError::read_before_execution)?;
+                    cfg.put(
+                        aws_smithy_runtime_api::client::orchestrator::EndpointResolverParams::new(
+                            params,
+                        ),
+                    );
+
+                    Ok(())
+                }
+            }
+
+            cfg.get::<Interceptors<HttpRequest, HttpResponse>>()
+                .expect("interceptors set")
+                .register_operation_interceptor(
+                    Arc::new(ListObjectsV2EndpointParamsInterceptor) as _
+                )
+                .register_operation_interceptor(Arc::new(
+                    ListObjectsV2EndpointParamsFinalizerInterceptor,
+                ) as _);
+            Ok(())
+        }
+    }
+
     let conn = TestConnection::new(vec![(
                 http::Request::builder()
                     .header("authorization", "AWS4-HMAC-SHA256 Credential=ANOTREAL/20210618/us-east-1/s3/aws4_request, SignedHeaders=host;x-amz-content-sha256;x-amz-date;x-amz-security-token;x-amz-user-agent, Signature=ae78f74d26b6b0c3a403d9e8cc7ec3829d6264a2b33db672bf2b151bbb901786")
@@ -187,7 +266,8 @@ async fn sra_manual_test() {
 
     let runtime_plugins = aws_smithy_runtime_api::client::runtime_plugin::RuntimePlugins::new()
         .with_client_plugin(ManualServiceRuntimePlugin(conn.clone()))
-        .with_operation_plugin(aws_sdk_s3::operation::list_objects_v2::ListObjectsV2::new());
+        .with_operation_plugin(aws_sdk_s3::operation::list_objects_v2::ListObjectsV2::new())
+        .with_operation_plugin(ManualOperationRuntimePlugin);
 
     let input = ListObjectsV2Input::builder()
         .bucket("test-bucket")

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/RustClientCodegenPlugin.kt

@@ -15,6 +15,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegen
 import software.amazon.smithy.rust.codegen.client.smithy.customize.CombinedClientCodegenDecorator
 import software.amazon.smithy.rust.codegen.client.smithy.customize.NoOpEventStreamSigningDecorator
 import software.amazon.smithy.rust.codegen.client.smithy.customize.RequiredCustomizations
+import software.amazon.smithy.rust.codegen.client.smithy.endpoint.EndpointParamsDecorator
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.EndpointsDecorator
 import software.amazon.smithy.rust.codegen.client.smithy.generators.client.FluentClientDecorator
 import software.amazon.smithy.rust.codegen.client.testutil.ClientDecoratableBuildPlugin
@@ -58,6 +59,7 @@ class RustClientCodegenPlugin : ClientDecoratableBuildPlugin() {
                 RequiredCustomizations(),
                 FluentClientDecorator(),
                 EndpointsDecorator(),
+                EndpointParamsDecorator(),
                 NoOpEventStreamSigningDecorator(),
                 ApiKeyAuthDecorator(),
                 *decorator,

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/EndpointPrefixGenerator.kt

@@ -7,16 +7,16 @@ package software.amazon.smithy.rust.codegen.client.smithy.customizations
 
 import software.amazon.smithy.model.shapes.OperationShape
 import software.amazon.smithy.model.traits.EndpointTrait
+import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
 import software.amazon.smithy.rust.codegen.client.smithy.generators.EndpointTraitBindings
 import software.amazon.smithy.rust.codegen.core.rustlang.Writable
 import software.amazon.smithy.rust.codegen.core.rustlang.rust
 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.CodegenContext
 import software.amazon.smithy.rust.codegen.core.smithy.customize.OperationCustomization
 import software.amazon.smithy.rust.codegen.core.smithy.customize.OperationSection
 
-class EndpointPrefixGenerator(private val codegenContext: CodegenContext, private val shape: OperationShape) :
+class EndpointPrefixGenerator(private val codegenContext: ClientCodegenContext, private val shape: OperationShape) :
     OperationCustomization() {
     override fun section(section: OperationSection): Writable = when (section) {
         is OperationSection.MutateRequest -> writable {
@@ -29,11 +29,16 @@ class EndpointPrefixGenerator(private val codegenContext: CodegenContext, privat
                     epTrait,
                 )
                 withBlock("let endpoint_prefix = ", "?;") {
-                    endpointTraitBindings.render(this, "self")
+                    endpointTraitBindings.render(
+                        this,
+                        "self",
+                        codegenContext.settings.codegenConfig.enableNewSmithyRuntime,
+                    )
                 }
                 rust("request.properties_mut().insert(endpoint_prefix);")
             }
         }
+
         else -> emptySection
     }
 }

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customize/ClientCodegenDecorator.kt

@@ -74,6 +74,7 @@ interface ClientCodegenDecorator : CoreCodegenDecorator<ClientCodegenContext> {
      */
     fun operationRuntimePluginCustomizations(
         codegenContext: ClientCodegenContext,
+        operation: OperationShape,
         baseCustomizations: List<OperationRuntimePluginCustomization>,
     ): List<OperationRuntimePluginCustomization> = baseCustomizations
 }
@@ -135,10 +136,11 @@ open class CombinedClientCodegenDecorator(decorators: List<ClientCodegenDecorato
 
     override fun operationRuntimePluginCustomizations(
         codegenContext: ClientCodegenContext,
+        operation: OperationShape,
         baseCustomizations: List<OperationRuntimePluginCustomization>,
     ): List<OperationRuntimePluginCustomization> =
         combineCustomizations(baseCustomizations) { decorator, customizations ->
-            decorator.operationRuntimePluginCustomizations(codegenContext, customizations)
+            decorator.operationRuntimePluginCustomizations(codegenContext, operation, customizations)
         }
 
     companion object {

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointConfigCustomization.kt

@@ -69,6 +69,7 @@ internal class EndpointConfigCustomization(
                         /// use aws_smithy_http::endpoint;
                         /// use $moduleUseName::endpoint::{Params as EndpointParams, DefaultResolver};
                         /// /// Endpoint resolver which adds a prefix to the generated endpoint
+                        /// ##[derive(Debug)]
                         /// struct PrefixResolver {
                         ///     base_resolver: DefaultResolver,
                         ///     prefix: String
@@ -132,6 +133,7 @@ internal class EndpointConfigCustomization(
                             RuntimeType.forInlineFun("MissingResolver", ClientRustModule.Endpoint) {
                                 rustTemplate(
                                     """
+                                    ##[derive(Debug)]
                                     pub(crate) struct MissingResolver;
                                     impl<T> #{ResolveEndpoint}<T> for MissingResolver {
                                         fn resolve_endpoint(&self, _params: &T) -> #{Result} {

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointParamsDecorator.kt

@@ -0,0 +1,56 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.client.smithy.endpoint
+
+import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
+import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
+import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationRuntimePluginCustomization
+import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationRuntimePluginSection
+import software.amazon.smithy.rust.codegen.core.rustlang.Writable
+import software.amazon.smithy.rust.codegen.core.rustlang.rust
+import software.amazon.smithy.rust.codegen.core.rustlang.writable
+import software.amazon.smithy.rust.codegen.core.util.letIf
+
+/**
+ * Decorator that injects operation-level interceptors that configure an endpoint parameters builder
+ * with operation specific information, e.g. a bucket name.
+ *
+ * Whenever a setter needs to be called on the endpoint parameters builder with operation specific information,
+ * this decorator must be used.
+ */
+class EndpointParamsDecorator : ClientCodegenDecorator {
+    override val name: String get() = "EndpointParamsDecorator"
+    override val order: Byte get() = 0
+
+    override fun operationRuntimePluginCustomizations(
+        codegenContext: ClientCodegenContext,
+        operation: OperationShape,
+        baseCustomizations: List<OperationRuntimePluginCustomization>,
+    ): List<OperationRuntimePluginCustomization> =
+        baseCustomizations.letIf(codegenContext.settings.codegenConfig.enableNewSmithyRuntime) {
+            it + listOf(EndpointParametersRuntimePluginCustomization(codegenContext, operation))
+        }
+}
+
+private class EndpointParametersRuntimePluginCustomization(
+    private val codegenContext: ClientCodegenContext,
+    private val operation: OperationShape,
+) : OperationRuntimePluginCustomization() {
+    override fun section(section: OperationRuntimePluginSection): Writable = writable {
+        val symbolProvider = codegenContext.symbolProvider
+        val operationName = symbolProvider.toSymbol(operation).name
+        if (section is OperationRuntimePluginSection.AdditionalConfig) {
+            section.registerInterceptor(codegenContext.runtimeConfig, this) {
+                rust("${operationName}EndpointParamsInterceptor")
+            }
+            // The finalizer interceptor should be registered last
+            section.registerInterceptor(codegenContext.runtimeConfig, this) {
+                rust("${operationName}EndpointParamsFinalizerInterceptor")
+            }
+        }
+    }
+}

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointTypesGenerator.kt

@@ -41,6 +41,7 @@ class EndpointTypesGenerator(
     }
 
     fun paramsStruct(): RuntimeType = EndpointParamsGenerator(params).paramsStruct()
+    fun paramsBuilder(): RuntimeType = EndpointParamsGenerator(params).paramsBuilder()
     fun defaultResolver(): RuntimeType? =
         rules?.let { EndpointResolverGenerator(stdlib, runtimeConfig).defaultEndpointResolver(it) }
 

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/generators/EndpointParamsGenerator.kt

@@ -118,7 +118,7 @@ internal class EndpointParamsGenerator(private val parameters: Parameters) {
         generateEndpointsStruct(this)
     }
 
-    private fun endpointsBuilder(): RuntimeType = RuntimeType.forInlineFun("ParamsBuilder", ClientRustModule.Endpoint) {
+    internal fun paramsBuilder(): RuntimeType = RuntimeType.forInlineFun("ParamsBuilder", ClientRustModule.Endpoint) {
         generateEndpointParamsBuilder(this)
     }
 
@@ -182,7 +182,7 @@ internal class EndpointParamsGenerator(private val parameters: Parameters) {
                     #{Builder}::default()
                 }
                 """,
-                "Builder" to endpointsBuilder(),
+                "Builder" to paramsBuilder(),
             )
             parameters.toList().forEach { parameter ->
                 val name = parameter.memberName()