Port connection poisoning tests to orchestrator and fix issues discov…

…ered (#3029)

Issues fixed:
- `ConnectionPoisoningInterceptor` should use the
`read_after_deserialization` hook so that it's possible to poison
modeled transient errors.
- aws-config should enable connection poisoning on its IMDS/ECS/HTTP
clients.

----

_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-config/src/http_credential_provider.rs

@@ -115,6 +115,7 @@ impl Builder {
             .http_connector(SharedHttpConnector::new(DynConnectorAdapter::new(
                 connector,
             )))
+            .with_connection_poisoning()
             .endpoint_url(endpoint)
             .no_auth()
             .runtime_plugin(StaticRuntimePlugin::new().with_config({

aws/rust-runtime/aws-config/src/imds/client.rs

@@ -458,6 +458,7 @@ impl Builder {
                 config.time_source(),
                 self.token_ttl.unwrap_or(DEFAULT_TOKEN_TTL),
             ))
+            .with_connection_poisoning()
             .serializer(|path| {
                 Ok(http::Request::builder()
                     .uri(path)

aws/rust-runtime/aws-config/src/imds/client/token.rs

@@ -132,6 +132,7 @@ impl TokenResolver {
                     .operation_name("get-token")
                     .runtime_plugin(common_plugin)
                     .no_auth()
+                    .with_connection_poisoning()
                     .serializer(move |_| {
                         Ok(http::Request::builder()
                             .method("PUT")

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

@@ -54,6 +54,14 @@ impl<'a, I, O, E> BeforeSerializationInterceptorContextRef<'a, I, O, E> {
     pub fn input(&self) -> &I {
         expect!(self, input)
     }
+
+    /// Downgrade this wrapper struct, returning the underlying InterceptorContext.
+    ///
+    /// There's no good reason to use this unless you're writing tests or you have to
+    /// interact with an API that doesn't support the context wrapper structs.
+    pub fn inner(&self) -> &'_ InterceptorContext<I, O, E> {
+        self.inner
+    }
 }
 
 //
@@ -80,6 +88,22 @@ impl<'a, I, O, E> BeforeSerializationInterceptorContextMut<'a, I, O, E> {
     pub fn input_mut(&mut self) -> &mut I {
         expect!(self, input_mut)
     }
+
+    /// Downgrade this wrapper struct, returning the underlying InterceptorContext.
+    ///
+    /// There's no good reason to use this unless you're writing tests or you have to
+    /// interact with an API that doesn't support the context wrapper structs.
+    pub fn inner(&self) -> &'_ InterceptorContext<I, O, E> {
+        self.inner
+    }
+
+    /// Downgrade this wrapper struct, returning the underlying InterceptorContext.
+    ///
+    /// There's no good reason to use this unless you're writing tests or you have to
+    /// interact with an API that doesn't support the context wrapper structs.
+    pub fn inner_mut(&mut self) -> &'_ mut InterceptorContext<I, O, E> {
+        self.inner
+    }
 }
 
 //
@@ -101,6 +125,14 @@ impl<'a, I, O, E> BeforeTransmitInterceptorContextRef<'a, I, O, E> {
     pub fn request(&self) -> &Request {
         expect!(self, request)
     }
+
+    /// Downgrade this wrapper struct, returning the underlying InterceptorContext.
+    ///
+    /// There's no good reason to use this unless you're writing tests or you have to
+    /// interact with an API that doesn't support the context wrapper structs.
+    pub fn inner(&self) -> &'_ InterceptorContext<I, O, E> {
+        self.inner
+    }
 }
 
 //
@@ -127,6 +159,22 @@ impl<'a, I, O, E> BeforeTransmitInterceptorContextMut<'a, I, O, E> {
     pub fn request_mut(&mut self) -> &mut Request {
         expect!(self, request_mut)
     }
+
+    /// Downgrade this wrapper struct, returning the underlying InterceptorContext.
+    ///
+    /// There's no good reason to use this unless you're writing tests or you have to
+    /// interact with an API that doesn't support the context wrapper structs.
+    pub fn inner(&self) -> &'_ InterceptorContext<I, O, E> {
+        self.inner
+    }
+
+    /// Downgrade this wrapper struct, returning the underlying InterceptorContext.
+    ///
+    /// There's no good reason to use this unless you're writing tests or you have to
+    /// interact with an API that doesn't support the context wrapper structs.
+    pub fn inner_mut(&mut self) -> &'_ mut InterceptorContext<I, O, E> {
+        self.inner
+    }
 }
 
 //
@@ -148,6 +196,14 @@ impl<'a, I, O, E> BeforeDeserializationInterceptorContextRef<'a, I, O, E> {
     pub fn response(&self) -> &Response {
         expect!(self, response)
     }
+
+    /// Downgrade this wrapper struct, returning the underlying InterceptorContext.
+    ///
+    /// There's no good reason to use this unless you're writing tests or you have to
+    /// interact with an API that doesn't support the context wrapper structs.
+    pub fn inner(&self) -> &'_ InterceptorContext<I, O, E> {
+        self.inner
+    }
 }
 
 //
@@ -174,10 +230,18 @@ impl<'a, I, O, E> BeforeDeserializationInterceptorContextMut<'a, I, O, E> {
         expect!(self, response_mut)
     }
 
-    #[doc(hidden)]
-    /// Downgrade this helper struct, returning the underlying InterceptorContext. There's no good
-    /// reason to use this unless you're writing tests or you have to interact with an API that
-    /// doesn't support the helper structs.
+    /// Downgrade this wrapper struct, returning the underlying InterceptorContext.
+    ///
+    /// There's no good reason to use this unless you're writing tests or you have to
+    /// interact with an API that doesn't support the context wrapper structs.
+    pub fn inner(&self) -> &'_ InterceptorContext<I, O, E> {
+        self.inner
+    }
+
+    /// Downgrade this wrapper struct, returning the underlying InterceptorContext.
+    ///
+    /// There's no good reason to use this unless you're writing tests or you have to
+    /// interact with an API that doesn't support the context wrapper structs.
     pub fn inner_mut(&mut self) -> &'_ mut InterceptorContext<I, O, E> {
         self.inner
     }
@@ -206,6 +270,14 @@ impl<'a, I, O, E> AfterDeserializationInterceptorContextRef<'a, I, O, E> {
     pub fn output_or_error(&self) -> Result<&O, &OrchestratorError<E>> {
         expect!(self, output_or_error)
     }
+
+    /// Downgrade this wrapper struct, returning the underlying InterceptorContext.
+    ///
+    /// There's no good reason to use this unless you're writing tests or you have to
+    /// interact with an API that doesn't support the context wrapper structs.
+    pub fn inner(&self) -> &'_ InterceptorContext<I, O, E> {
+        self.inner
+    }
 }
 
 //
@@ -243,6 +315,14 @@ impl<'a, I, O, E> FinalizerInterceptorContextRef<'a, I, O, E> {
     pub fn output_or_error(&self) -> Option<Result<&O, &OrchestratorError<E>>> {
         self.inner.output_or_error.as_ref().map(|o| o.as_ref())
     }
+
+    /// Downgrade this wrapper struct, returning the underlying InterceptorContext.
+    ///
+    /// There's no good reason to use this unless you're writing tests or you have to
+    /// interact with an API that doesn't support the context wrapper structs.
+    pub fn inner(&self) -> &'_ InterceptorContext<I, O, E> {
+        self.inner
+    }
 }
 
 //
@@ -300,4 +380,20 @@ impl<'a, I, O, E> FinalizerInterceptorContextMut<'a, I, O, E> {
     pub fn output_or_error_mut(&mut self) -> Option<&mut Result<O, OrchestratorError<E>>> {
         self.inner.output_or_error.as_mut()
     }
+
+    /// Downgrade this wrapper struct, returning the underlying InterceptorContext.
+    ///
+    /// There's no good reason to use this unless you're writing tests or you have to
+    /// interact with an API that doesn't support the context wrapper structs.
+    pub fn inner(&self) -> &'_ InterceptorContext<I, O, E> {
+        self.inner
+    }
+
+    /// Downgrade this wrapper struct, returning the underlying InterceptorContext.
+    ///
+    /// There's no good reason to use this unless you're writing tests or you have to
+    /// interact with an API that doesn't support the context wrapper structs.
+    pub fn inner_mut(&mut self) -> &'_ mut InterceptorContext<I, O, E> {
+        self.inner
+    }
 }

rust-runtime/aws-smithy-runtime/src/client/connectors/connection_poisoning.rs

@@ -6,7 +6,7 @@
 use aws_smithy_http::connection::{CaptureSmithyConnection, ConnectionMetadata};
 use aws_smithy_runtime_api::box_error::BoxError;
 use aws_smithy_runtime_api::client::interceptors::context::{
-    BeforeDeserializationInterceptorContextMut, BeforeTransmitInterceptorContextMut,
+    AfterDeserializationInterceptorContextRef, BeforeTransmitInterceptorContextMut,
 };
 use aws_smithy_runtime_api::client::interceptors::Interceptor;
 use aws_smithy_runtime_api::client::retries::{ClassifyRetry, RetryReason};
@@ -61,9 +61,9 @@ impl Interceptor for ConnectionPoisoningInterceptor {
         Ok(())
     }
 
-    fn modify_before_deserialization(
+    fn read_after_deserialization(
         &self,
-        context: &mut BeforeDeserializationInterceptorContextMut<'_>,
+        context: &AfterDeserializationInterceptorContextRef<'_>,
         runtime_components: &RuntimeComponents,
         cfg: &mut ConfigBag,
     ) -> Result<(), BoxError> {
@@ -77,7 +77,7 @@ impl Interceptor for ConnectionPoisoningInterceptor {
             .ok_or("retry classifiers are required for connection poisoning to work")?;
 
         let error_is_transient = retry_classifiers
-            .classify_retry(context.inner_mut())
+            .classify_retry(context.inner())
             .map(|reason| reason == RetryReason::Error(ErrorKind::TransientError))
             .unwrap_or_default();
         let connection_poisoning_is_enabled =

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

@@ -4,6 +4,7 @@
  */
 
 use crate::client::auth::no_auth::{NoAuthScheme, NO_AUTH_SCHEME_ID};
+use crate::client::connectors::connection_poisoning::ConnectionPoisoningInterceptor;
 use crate::client::identity::no_auth::NoAuthIdentityResolver;
 use crate::client::orchestrator::endpoints::StaticUriEndpointResolver;
 use crate::client::retries::strategy::{NeverRetryStrategy, StandardRetryStrategy};
@@ -254,6 +255,11 @@ impl<I, O, E> OperationBuilder<I, O, E> {
         self
     }
 
+    /// Registers the [`ConnectionPoisoningInterceptor`].
+    pub fn with_connection_poisoning(self) -> Self {
+        self.interceptor(ConnectionPoisoningInterceptor::new())
+    }
+
     pub fn runtime_plugin(mut self, runtime_plugin: impl IntoShared<SharedRuntimePlugin>) -> Self {
         self.runtime_plugins.push(runtime_plugin.into_shared());
         self