From fba7cff032797b05f80f4388805eedd32e3fa8a1 Mon Sep 17 00:00:00 2001
From: Harry Barber <hlbarber@amazon.co.uk>
Date: Tue, 3 Jan 2023 14:09:14 +0000
Subject: [PATCH] Introduce OperationExtensionPlugin

---
 .../aws-smithy-http-server/src/extension.rs   | 127 +++++++++++++++++-
 1 file changed, 121 insertions(+), 6 deletions(-)

diff --git a/rust-runtime/aws-smithy-http-server/src/extension.rs b/rust-runtime/aws-smithy-http-server/src/extension.rs
index 265f1cee4f..5c302a757a 100644
--- a/rust-runtime/aws-smithy-http-server/src/extension.rs
+++ b/rust-runtime/aws-smithy-http-server/src/extension.rs
@@ -19,20 +19,24 @@
 //!
 //! [extensions]: https://docs.rs/http/latest/http/struct.Extensions.html
 
-use std::ops::Deref;
+use std::{fmt, future::Future, ops::Deref, pin::Pin, task::Context, task::Poll};
 
+use futures_util::ready;
+use futures_util::TryFuture;
 use thiserror::Error;
+use tower::{layer::util::Stack, Layer, Service};
 
+use crate::operation::{Operation, OperationShape};
+use crate::plugin::{plugin_from_operation_name_fn, OperationNameFn, Plugin, PluginPipeline, PluginStack};
 #[allow(deprecated)]
 use crate::request::RequestParts;
 
-pub use crate::request::extension::Extension;
-pub use crate::request::extension::MissingExtension;
+pub use crate::request::extension::{Extension, MissingExtension};
 
 /// Extension type used to store information about Smithy operations in HTTP responses.
-/// This extension type is set when it has been correctly determined that the request should be
-/// routed to a particular operation. The operation handler might not even get invoked because the
-/// request fails to deserialize into the modeled operation input.
+/// This extension type is inserted, via the [`OperationExtensionPlugin`], whenever it has been correctly determined
+/// that the request should be routed to a particular operation. The operation handler might not even get invoked
+/// because the request fails to deserialize into the modeled operation input.
 ///
 /// The format given must be the absolute shape ID with `#` replaced with a `.`.
 #[derive(Debug, Clone)]
@@ -81,6 +85,117 @@ impl OperationExtension {
     }
 }
 
+pin_project_lite::pin_project! {
+    /// The [`Service::Future`] of [`OperationExtensionService`] - inserts an [`OperationExtension`] into the
+    /// [`http::Response]`.
+    pub struct OperationExtensionFuture<Fut> {
+        #[pin]
+        inner: Fut,
+        operation_extension: Option<OperationExtension>
+    }
+}
+
+impl<Fut, RespB> Future for OperationExtensionFuture<Fut>
+where
+    Fut: TryFuture<Ok = http::Response<RespB>>,
+{
+    type Output = Result<http::Response<RespB>, Fut::Error>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+        let ext = this
+            .operation_extension
+            .take()
+            .expect("futures cannot be polled after completion");
+        let resp = ready!(this.inner.try_poll(cx));
+        Poll::Ready(resp.map(|mut resp| {
+            resp.extensions_mut().insert(ext);
+            resp
+        }))
+    }
+}
+
+/// Inserts a [`OperationExtension`] into the extensions of the [`http::Response`].
+#[derive(Debug, Clone)]
+pub struct OperationExtensionService<S> {
+    inner: S,
+    operation_extension: OperationExtension,
+}
+
+impl<S, B, RespBody> Service<http::Request<B>> for OperationExtensionService<S>
+where
+    S: Service<http::Request<B>, Response = http::Response<RespBody>>,
+{
+    type Response = http::Response<RespBody>;
+    type Error = S::Error;
+    type Future = OperationExtensionFuture<S::Future>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, req: http::Request<B>) -> Self::Future {
+        OperationExtensionFuture {
+            inner: self.inner.call(req),
+            operation_extension: Some(self.operation_extension.clone()),
+        }
+    }
+}
+
+/// A [`Layer`] applying the [`OperationExtensionService`] to an inner [`Service`].
+#[derive(Debug, Clone)]
+pub struct OperationExtensionLayer(OperationExtension);
+
+impl<S> Layer<S> for OperationExtensionLayer {
+    type Service = OperationExtensionService<S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        OperationExtensionService {
+            inner,
+            operation_extension: self.0.clone(),
+        }
+    }
+}
+
+/// A [`Plugin`] which applies [`OperationExtensionLayer`] to every operation.
+pub struct OperationExtensionPlugin(OperationNameFn<fn(&'static str) -> OperationExtensionLayer>);
+
+impl fmt::Debug for OperationExtensionPlugin {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_tuple("OperationExtensionPlugin").field(&"...").finish()
+    }
+}
+
+impl<P, Op, S, L> Plugin<P, Op, S, L> for OperationExtensionPlugin
+where
+    Op: OperationShape,
+{
+    type Service = S;
+    type Layer = Stack<L, OperationExtensionLayer>;
+
+    fn map(&self, input: Operation<S, L>) -> Operation<Self::Service, Self::Layer> {
+        <OperationNameFn<fn(&'static str) -> OperationExtensionLayer> as Plugin<P, Op, S, L>>::map(&self.0, input)
+    }
+}
+
+/// An extension trait on [`PluginPipeline`] allowing the application of [`OperationExtensionPlugin`].
+///
+/// See [`module`](crate::extension) documentation for more info.
+pub trait OperationExtensionExt<P> {
+    /// Apply the [`OperationExtensionPlugin`], which inserts the [`OperationExtension`] into every [`http::Response`].
+    fn insert_operation_extension(self) -> PluginPipeline<PluginStack<OperationExtensionPlugin, P>>;
+}
+
+impl<P> OperationExtensionExt<P> for PluginPipeline<P> {
+    fn insert_operation_extension(self) -> PluginPipeline<PluginStack<OperationExtensionPlugin, P>> {
+        let plugin = OperationExtensionPlugin(plugin_from_operation_name_fn(|name| {
+            let operation_extension = OperationExtension::new(name).expect("Operation name is malformed, this should never happen. Please file an issue against https://github.com/awslabs/smithy-rs");
+            OperationExtensionLayer(operation_extension)
+        }));
+        self.push(plugin)
+    }
+}
+
 /// Extension type used to store the type of user-modeled error returned by an operation handler.
 /// These are modeled errors, defined in the Smithy model.
 #[derive(Debug, Clone)]