Follow up on DefaultEndpointResolver in the orchestrator (#2592)

## Motivation and Context
This PR incorporates post-merge feedback left in #2577.

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._

---------

Co-authored-by: Yuki Saito <[email protected]>

Author: ysaito1001

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

@@ -17,7 +17,7 @@ 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::DefaultEndpointResolver;
+use aws_smithy_runtime::client::orchestrator::endpoints::DefaultEndpointResolver;
 use aws_smithy_runtime_api::client::interceptors::{
     Interceptor, InterceptorContext, InterceptorError, Interceptors,
 };

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

@@ -32,6 +32,7 @@ class EndpointParamsInterceptorGenerator(
         arrayOf(
             "BoxError" to runtimeApi.resolve("client::runtime_plugin::BoxError"),
             "ConfigBag" to runtimeApi.resolve("config_bag::ConfigBag"),
+            "ContextAttachedError" to interceptors.resolve("error::ContextAttachedError"),
             "EndpointResolverParams" to orchestrator.resolve("EndpointResolverParams"),
             "HttpResponse" to orchestrator.resolve("HttpResponse"),
             "HttpRequest" to orchestrator.resolve("HttpRequest"),
@@ -83,10 +84,10 @@ class EndpointParamsInterceptorGenerator(
             let input = context.input()?;
             let _input = input
                 .downcast_ref::<${operationInput.name}>()
-                .ok_or_else(|| #{InterceptorError}::invalid_input_access())?;
+                .ok_or_else(|| "failed to downcast to ${operationInput.name}")?;
             let params_builder = cfg
                 .get::<#{ParamsBuilder}>()
-                .ok_or(#{InterceptorError}::read_before_execution("missing endpoint params builder"))?
+                .ok_or_else(|| "missing endpoint params builder")?
                 .clone();
             ${"" /* TODO(EndpointResolver): Call setters on `params_builder` to update its fields by using values from `_input` */}
             cfg.put(params_builder);
@@ -111,7 +112,7 @@ class EndpointParamsInterceptorGenerator(
             )
             withBlockTemplate(
                 "let endpoint_prefix = ",
-                ".map_err(#{InterceptorError}::read_before_execution)?;",
+                """.map_err(|err| #{ContextAttachedError}::new("endpoint prefix could not be built", err))?;""",
                 *codegenScope,
             ) {
                 endpointTraitBindings.render(
@@ -130,11 +131,11 @@ class EndpointParamsInterceptorGenerator(
             let _ = context;
             let params_builder = cfg
                 .get::<#{ParamsBuilder}>()
-                .ok_or(#{InterceptorError}::read_before_execution("missing endpoint params builder"))?
+                .ok_or_else(|| "missing endpoint params builder")?
                 .clone();
             let params = params_builder
                 .build()
-                .map_err(#{InterceptorError}::read_before_execution)?;
+                .map_err(|err| #{ContextAttachedError}::new("endpoint params could not be built", err))?;
             cfg.put(
                 #{EndpointResolverParams}::new(params)
             );

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

@@ -86,7 +86,7 @@ class ServiceRuntimePluginGenerator(
             "ConfigBagAccessors" to runtimeApi.resolve("client::orchestrator::ConfigBagAccessors"),
             "Connection" to runtimeApi.resolve("client::orchestrator::Connection"),
             "ConnectorSettings" to RuntimeType.smithyClient(rc).resolve("http_connector::ConnectorSettings"),
-            "DefaultEndpointResolver" to runtimeApi.resolve("client::endpoints::DefaultEndpointResolver"),
+            "DefaultEndpointResolver" to runtime.resolve("client::orchestrator::endpoints::DefaultEndpointResolver"),
             "DynConnectorAdapter" to runtime.resolve("client::connections::adapter::DynConnectorAdapter"),
             "HttpAuthSchemes" to runtimeApi.resolve("client::orchestrator::HttpAuthSchemes"),
             "IdentityResolvers" to runtimeApi.resolve("client::orchestrator::IdentityResolvers"),

rust-runtime/aws-smithy-runtime-api/src/client.rs

@@ -21,8 +21,5 @@ pub mod retries;
 /// Runtime plugin type definitions.
 pub mod runtime_plugin;
 
-/// Smithy endpoint resolution runtime plugins
-pub mod endpoints;
-
 /// Smithy auth runtime plugins
 pub mod auth;

rust-runtime/aws-smithy-runtime-api/src/client/endpoints.rs

@@ -169,3 +169,31 @@ impl std::error::Error for InterceptorError {
         self.source.as_ref().map(|err| err.as_ref() as _)
     }
 }
+
+/// A convenience error that allows for adding additional `context` to `source`
+#[derive(Debug)]
+pub struct ContextAttachedError {
+    context: String,
+    source: Option<BoxError>,
+}
+
+impl ContextAttachedError {
+    pub fn new(context: impl Into<String>, source: impl Into<BoxError>) -> Self {
+        Self {
+            context: context.into(),
+            source: Some(source.into()),
+        }
+    }
+}
+
+impl fmt::Display for ContextAttachedError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.context)
+    }
+}
+
+impl std::error::Error for ContextAttachedError {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        self.source.as_ref().map(|err| err.as_ref() as _)
+    }
+}

rust-runtime/aws-smithy-runtime-api/src/client/interceptors/error.rs

@@ -6,5 +6,5 @@ allowed_external_types = [
     "http::header::name::HeaderName",
     "http::request::Request",
     "http::response::Response",
-    "uri::Uri",
+    "http::uri::Uri",
 ]

rust-runtime/aws-smithy-runtime/external-types.toml

@@ -18,7 +18,8 @@ use aws_smithy_runtime_api::config_bag::ConfigBag;
 use tracing::{debug_span, Instrument};
 
 mod auth;
-mod endpoints;
+/// Defines types that implement a trait for endpoint resolution
+pub mod endpoints;
 mod http;
 pub(self) mod phase;
 

rust-runtime/aws-smithy-runtime/src/client/orchestrator.rs

@@ -3,12 +3,114 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-use aws_smithy_http::endpoint::EndpointPrefix;
+use aws_smithy_http::endpoint::error::ResolveEndpointError;
+use aws_smithy_http::endpoint::{
+    apply_endpoint, EndpointPrefix, ResolveEndpoint, SharedEndpointResolver,
+};
 use aws_smithy_runtime_api::client::interceptors::InterceptorContext;
 use aws_smithy_runtime_api::client::orchestrator::{
-    BoxError, ConfigBagAccessors, HttpRequest, HttpResponse,
+    BoxError, ConfigBagAccessors, EndpointResolver, EndpointResolverParams, HttpRequest,
+    HttpResponse,
 };
 use aws_smithy_runtime_api::config_bag::ConfigBag;
+use http::header::HeaderName;
+use http::{HeaderValue, Uri};
+use std::fmt::Debug;
+use std::str::FromStr;
+
+#[derive(Debug, Clone)]
+pub struct StaticUriEndpointResolver {
+    endpoint: Uri,
+}
+
+impl StaticUriEndpointResolver {
+    pub fn http_localhost(port: u16) -> Self {
+        Self {
+            endpoint: Uri::from_str(&format!("http://localhost:{port}"))
+                .expect("all u16 values are valid ports"),
+        }
+    }
+
+    pub fn uri(endpoint: Uri) -> Self {
+        Self { endpoint }
+    }
+}
+
+impl EndpointResolver for StaticUriEndpointResolver {
+    fn resolve_and_apply_endpoint(
+        &self,
+        _params: &EndpointResolverParams,
+        _endpoint_prefix: Option<&EndpointPrefix>,
+        request: &mut HttpRequest,
+    ) -> Result<(), BoxError> {
+        apply_endpoint(request.uri_mut(), &self.endpoint, None)?;
+        Ok(())
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct DefaultEndpointResolver<Params> {
+    inner: SharedEndpointResolver<Params>,
+}
+
+impl<Params> DefaultEndpointResolver<Params> {
+    pub fn new(resolve_endpoint: SharedEndpointResolver<Params>) -> Self {
+        Self {
+            inner: resolve_endpoint,
+        }
+    }
+}
+
+impl<Params> EndpointResolver for DefaultEndpointResolver<Params>
+where
+    Params: Debug + Send + Sync + 'static,
+{
+    fn resolve_and_apply_endpoint(
+        &self,
+        params: &EndpointResolverParams,
+        endpoint_prefix: Option<&EndpointPrefix>,
+        request: &mut HttpRequest,
+    ) -> Result<(), BoxError> {
+        let endpoint = match params.get::<Params>() {
+            Some(params) => self.inner.resolve_endpoint(params)?,
+            None => {
+                return Err(Box::new(ResolveEndpointError::message(
+                    "params of expected type was not present",
+                )));
+            }
+        };
+
+        let uri: Uri = endpoint.url().parse().map_err(|err| {
+            ResolveEndpointError::from_source("endpoint did not have a valid uri", err)
+        })?;
+
+        apply_endpoint(request.uri_mut(), &uri, endpoint_prefix).map_err(|err| {
+            ResolveEndpointError::message(format!(
+                "failed to apply endpoint `{:?}` to request `{:?}`",
+                uri, request,
+            ))
+            .with_source(Some(err.into()))
+        })?;
+
+        for (header_name, header_values) in endpoint.headers() {
+            request.headers_mut().remove(header_name);
+            for value in header_values {
+                request.headers_mut().insert(
+                    HeaderName::from_str(header_name).map_err(|err| {
+                        ResolveEndpointError::message("invalid header name")
+                            .with_source(Some(err.into()))
+                    })?,
+                    HeaderValue::from_str(value).map_err(|err| {
+                        ResolveEndpointError::message("invalid header value")
+                            .with_source(Some(err.into()))
+                    })?,
+                );
+            }
+        }
+
+        Ok(())
+    }
+}
 
 pub(super) fn orchestrate_endpoint(
     ctx: &mut InterceptorContext<HttpRequest, HttpResponse>,