From 8f8094052055c7adaa0cea3f95124d2fdc753b12 Mon Sep 17 00:00:00 2001
From: Harry Barber <hlbarber@amazon.co.uk>
Date: Fri, 19 Aug 2022 12:58:50 +0000
Subject: [PATCH] Allow handler extensions

---
 generated/examples/usage.rs                   |  41 ++++--
 generated/src/operations/empty_operation.rs   |   2 +-
 .../src/operations/get_pokemon_species.rs     |   2 +-
 generated/src/services/pokemon_service.rs     |  14 +-
 runtime/src/operation/flattened.rs            | 135 ++++++++++++++++++
 runtime/src/operation/handler.rs              | 121 ++++++++--------
 runtime/src/operation/http_conversions.rs     | 107 ++++++++++++++
 runtime/src/operation/mod.rs                  |  42 ++++--
 runtime/src/operation/shape.rs                |  31 ++--
 runtime/src/operation/upgrade.rs              |  87 ++++++-----
 10 files changed, 425 insertions(+), 157 deletions(-)
 create mode 100644 runtime/src/operation/flattened.rs
 create mode 100644 runtime/src/operation/http_conversions.rs

diff --git a/generated/examples/usage.rs b/generated/examples/usage.rs
index 2b31304..a490b67 100644
--- a/generated/examples/usage.rs
+++ b/generated/examples/usage.rs
@@ -1,34 +1,35 @@
 use std::{convert::Infallible, future::Ready, task::Poll};
 
 use generated::{operations::*, services::*, structures::*};
-use runtime::operation::{AdjoinState, OperationError, OperationShapeExt};
+use runtime::operation::{Extension, OperationError, OperationShapeExt};
 use tower::{util::MapResponseLayer, Service};
 
-/// Fallible handler with state
+// Fallible handler with extensions.
 async fn get_pokemon_species_stateful(
     _input: GetPokemonSpeciesInput,
-    _state: usize,
+    _ext_a: Extension<usize>,
+    _ext_b: Extension<String>,
 ) -> Result<GetPokemonSpeciesOutput, ResourceNotFoundException> {
     todo!()
 }
 
-/// Fallible handler without state
+// Fallible handler without extensions.
 async fn get_pokemon_species(
     _input: GetPokemonSpeciesInput,
 ) -> Result<GetPokemonSpeciesOutput, ResourceNotFoundException> {
     todo!()
 }
 
-/// Infallible handler without state
+// Infallible handler without extensions.
 async fn empty_operation(_input: EmptyOperationInput) -> EmptyOperationOutput {
     todo!()
 }
 
-/// Bespoke implementation of `EmptyOperation`.
+// Bespoke implementation of `EmptyOperation`.
 #[derive(Clone)]
-struct EmptyOperationService;
+struct EmptyOperationServiceA;
 
-impl Service<EmptyOperationInput> for EmptyOperationService {
+impl Service<EmptyOperationInput> for EmptyOperationServiceA {
     type Response = EmptyOperationOutput;
     type Error = OperationError<String, Infallible>;
     type Future = Ready<Result<Self::Response, Self::Error>>;
@@ -42,13 +43,31 @@ impl Service<EmptyOperationInput> for EmptyOperationService {
     }
 }
 
+// Bespoke implementation of `EmptyOperation` with an extension.
+#[derive(Clone)]
+struct EmptyOperationServiceB;
+
+impl Service<(EmptyOperationInput, Extension<String>)> for EmptyOperationServiceB {
+    type Response = EmptyOperationOutput;
+    type Error = OperationError<String, Infallible>;
+    type Future = Ready<Result<Self::Response, Self::Error>>;
+
+    fn poll_ready(&mut self, _cx: &mut std::task::Context<'_>) -> Poll<Result<(), Self::Error>> {
+        todo!()
+    }
+
+    fn call(&mut self, _req: (EmptyOperationInput, Extension<String>)) -> Self::Future {
+        todo!()
+    }
+}
+
 fn main() {
     // Various ways of constructing operations
     let _get_pokemon_species = GetPokemonSpecies::from_handler(get_pokemon_species);
     let _empty_operation = EmptyOperation::from_handler(empty_operation);
-    let empty_operation = EmptyOperation::from_service(EmptyOperationService);
-    let get_pokemon_species =
-        GetPokemonSpecies::from_handler(get_pokemon_species_stateful.with_state(29));
+    let _empty_operation = EmptyOperation::from_service(EmptyOperationServiceA);
+    let empty_operation = EmptyOperation::from_service(EmptyOperationServiceB);
+    let get_pokemon_species = GetPokemonSpecies::from_handler(get_pokemon_species_stateful);
 
     // We can apply a layer to them
     let get_pokemon_species = get_pokemon_species.layer(MapResponseLayer::new(|resp| resp));
diff --git a/generated/src/operations/empty_operation.rs b/generated/src/operations/empty_operation.rs
index 374cb46..02c8b8d 100644
--- a/generated/src/operations/empty_operation.rs
+++ b/generated/src/operations/empty_operation.rs
@@ -32,7 +32,7 @@ impl<B> FromRequest<AWSRestJsonV1, EmptyOperation, B> for EmptyOperationInput {
 
     type Future = Ready<Result<Self, Self::Error>>;
 
-    fn from_request(_request: http::Request<B>) -> Self::Future {
+    fn from_request(_request: &mut http::Request<B>) -> Self::Future {
         todo!()
     }
 }
diff --git a/generated/src/operations/get_pokemon_species.rs b/generated/src/operations/get_pokemon_species.rs
index 1fb0b18..6f171e8 100644
--- a/generated/src/operations/get_pokemon_species.rs
+++ b/generated/src/operations/get_pokemon_species.rs
@@ -34,7 +34,7 @@ impl<B> FromRequest<AWSRestJsonV1, GetPokemonSpecies, B> for GetPokemonSpeciesIn
 
     type Future = Ready<Result<Self, Self::Error>>;
 
-    fn from_request(_request: http::Request<B>) -> Self::Future {
+    fn from_request(_request: &mut http::Request<B>) -> Self::Future {
         todo!()
     }
 }
diff --git a/generated/src/services/pokemon_service.rs b/generated/src/services/pokemon_service.rs
index 8106983..4ac555f 100644
--- a/generated/src/services/pokemon_service.rs
+++ b/generated/src/services/pokemon_service.rs
@@ -112,12 +112,12 @@ impl<Op1, Op2> PokemonServiceBuilder<Op1, Op2> {
 }
 
 impl<S1, S2, L1, L2> PokemonServiceBuilder<Operation<S1, L1>, Operation<S2, L2>> {
-    pub fn build<B>(self) -> PokemonService<RouteService<B>>
+    pub fn build<B, Exts1, Exts2>(self) -> PokemonService<RouteService<B>>
     where
         // GetPokemonSpecies composition
-        UpgradeLayer<AWSRestJsonV1, GetPokemonSpecies, B>: Layer<S1>,
-        L1: Layer<UpgradedService<AWSRestJsonV1, GetPokemonSpecies, B, S1>>,
-        S1: Service<<GetPokemonSpecies as OperationShape>::Input>,
+        UpgradeLayer<AWSRestJsonV1, GetPokemonSpecies, Exts1, B>: Layer<S1>,
+        L1: Layer<UpgradedService<AWSRestJsonV1, GetPokemonSpecies, Exts1, B, S1>>,
+        S1: Service<(<GetPokemonSpecies as OperationShape>::Input, Exts1)>,
         L1::Service: Service<http::Request<B>, Response = http::Response<BoxBody<Bytes, Error>>>,
         L1::Service: Clone + Send + 'static,
         <L1::Service as Service<http::Request<B>>>::Future: Send + 'static,
@@ -125,9 +125,9 @@ impl<S1, S2, L1, L2> PokemonServiceBuilder<Operation<S1, L1>, Operation<S2, L2>>
             Into<Box<dyn std::error::Error + Send + Sync>> + 'static,
 
         // EmptyOperation composition
-        UpgradeLayer<AWSRestJsonV1, EmptyOperation, B>: Layer<S2>,
-        L2: Layer<UpgradedService<AWSRestJsonV1, EmptyOperation, B, S2>>,
-        S2: Service<<EmptyOperation as OperationShape>::Input>,
+        UpgradeLayer<AWSRestJsonV1, EmptyOperation, Exts2, B>: Layer<S2>,
+        L2: Layer<UpgradedService<AWSRestJsonV1, EmptyOperation, Exts2, B, S2>>,
+        S2: Service<(<EmptyOperation as OperationShape>::Input, Exts2)>,
         L2::Service: Service<http::Request<B>, Response = http::Response<BoxBody<Bytes, Error>>>,
         L2::Service: Clone + Send + 'static,
         <L2::Service as Service<http::Request<B>>>::Future: Send + 'static,
diff --git a/runtime/src/operation/flattened.rs b/runtime/src/operation/flattened.rs
new file mode 100644
index 0000000..addea0b
--- /dev/null
+++ b/runtime/src/operation/flattened.rs
@@ -0,0 +1,135 @@
+use std::{
+    marker::PhantomData,
+    task::{Context, Poll},
+};
+
+use tower::Service;
+
+use super::{OperationError, OperationShape};
+
+/// A utility trait used to provide an even interface for all operation services.
+///
+/// This serves to take [`Service`]s of the form `Service<(Input, Arg0, Arg1, ...)>` to the canonical representation of
+/// `Service<(Input, (Arg0, Arg1, ...))>` inline with [`IntoService`](super::IntoService).
+pub trait Flattened<Op, Exts, PollError>:
+    Service<Self::Flattened, Response = Op::Output, Error = OperationError<PollError, Op::Error>>
+where
+    Op: OperationShape,
+{
+    type Flattened;
+
+    // Unflatten the request type.
+    fn unflatten(input: Op::Input, exts: Exts) -> Self::Flattened;
+}
+
+// `Service<Op::Input>`
+impl<Op, S, PollError> Flattened<Op, (), PollError> for S
+where
+    Op: OperationShape,
+    S: Service<Op::Input, Response = Op::Output, Error = OperationError<PollError, Op::Error>>,
+{
+    type Flattened = Op::Input;
+
+    fn unflatten(input: Op::Input, _exts: ()) -> Self::Flattened {
+        input
+    }
+}
+
+// `Service<(Op::Input, Arg0)>`
+impl<Op, Arg0, S, PollError> Flattened<Op, (Arg0,), PollError> for S
+where
+    Op: OperationShape,
+    S: Service<
+        (Op::Input, Arg0),
+        Response = Op::Output,
+        Error = OperationError<PollError, Op::Error>,
+    >,
+{
+    type Flattened = (Op::Input, Arg0);
+
+    fn unflatten(input: Op::Input, exts: (Arg0,)) -> Self::Flattened {
+        (input, exts.0)
+    }
+}
+
+// `Service<(Op::Input, Arg0, Arg1)>`
+impl<Op, Arg0, Arg1, S, PollError> Flattened<Op, (Arg0, Arg1), PollError> for S
+where
+    Op: OperationShape,
+    S: Service<
+        (Op::Input, Arg0, Arg1),
+        Response = Op::Output,
+        Error = OperationError<PollError, Op::Error>,
+    >,
+{
+    type Flattened = (Op::Input, Arg0, Arg1);
+
+    fn unflatten(input: Op::Input, exts: (Arg0, Arg1)) -> Self::Flattened {
+        (input, exts.0, exts.1)
+    }
+}
+
+/// An extension trait of [`Flattened`].
+pub trait FlattenedExt<Op, Exts, PollError>: Flattened<Op, Exts, PollError>
+where
+    Op: OperationShape,
+{
+    /// Convert the [`Flattened`] into a canonicalized [`Service`].
+    fn into_unflatten(self) -> IntoUnflattened<Op, Self, PollError>
+    where
+        Self: Sized,
+    {
+        IntoUnflattened {
+            inner: self,
+            _operation: PhantomData,
+            _poll_error: PhantomData,
+        }
+    }
+}
+
+impl<F, Op, Exts, PollError> FlattenedExt<Op, Exts, PollError> for F
+where
+    Op: OperationShape,
+    F: Flattened<Op, Exts, PollError>,
+{
+}
+
+/// A [`Service`] canonicalizing the request type of a [`Flattened`].
+#[derive(Debug)]
+pub struct IntoUnflattened<Op, S, PollError> {
+    inner: S,
+    _operation: PhantomData<Op>,
+    _poll_error: PhantomData<PollError>,
+}
+
+impl<Op, S, PollError> Clone for IntoUnflattened<Op, S, PollError>
+where
+    S: Clone,
+{
+    fn clone(&self) -> Self {
+        Self {
+            inner: self.inner.clone(),
+            _operation: PhantomData,
+            _poll_error: PhantomData,
+        }
+    }
+}
+
+impl<Op, S, Exts, PollError> Service<(Op::Input, Exts)> for IntoUnflattened<Op, S, PollError>
+where
+    Op: OperationShape,
+    S: Flattened<Op, Exts, PollError>,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = <S as Service<S::Flattened>>::Future;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, (input, exts): (Op::Input, Exts)) -> Self::Future {
+        let req = S::unflatten(input, exts);
+        self.inner.call(req)
+    }
+}
diff --git a/runtime/src/operation/handler.rs b/runtime/src/operation/handler.rs
index 29652d1..dc2058e 100644
--- a/runtime/src/operation/handler.rs
+++ b/runtime/src/operation/handler.rs
@@ -11,20 +11,16 @@ use futures::{
 };
 use tower::Service;
 
-/// The operation [`Service`] has two classes of failure modes - the failure models specified by
-/// the Smithy model and failures to [`Service::poll_ready`].
-pub enum OperationError<PollError, SmithyError> {
-    /// A [`Service::poll_ready`] failure occured.
-    Poll(PollError),
-    /// An error modelled by the Smithy model occured.
-    Smithy(SmithyError),
-}
+use super::{OperationError, OperationShape};
 
 /// A utility trait used to provide an even interface for all handlers.
-pub trait Handler<Input, Output, Error> {
-    type Future: Future<Output = Result<Output, Error>>;
+pub trait Handler<Op, Exts>
+where
+    Op: OperationShape,
+{
+    type Future: Future<Output = Result<Op::Output, Op::Error>>;
 
-    fn call(&mut self, req: Input) -> Self::Future;
+    fn call(&mut self, input: Op::Input, exts: Exts) -> Self::Future;
 }
 
 /// A utility trait used to provide an even interface over return types `Result<Ok, Error>`/`Ok`.
@@ -47,109 +43,108 @@ impl<Ok> ToResult<Ok, Infallible> for Ok {
 }
 
 // fn(Input) -> Output
-impl<Input, Output, Error, F, Fut> Handler<Input, Output, Error> for F
+impl<Op, F, Fut> Handler<Op, ()> for F
 where
-    F: FnMut(Input) -> Fut,
+    Op: OperationShape,
+    F: Fn(Op::Input) -> Fut,
     Fut: Future,
-    Fut::Output: ToResult<Output, Error>,
+    Fut::Output: ToResult<Op::Output, Op::Error>,
 {
-    type Future = Map<Fut, fn(Fut::Output) -> Result<Output, Error>>;
+    type Future = Map<Fut, fn(Fut::Output) -> Result<Op::Output, Op::Error>>;
 
-    fn call(&mut self, req: Input) -> Self::Future {
-        (self)(req).map(ToResult::into_result)
+    fn call(&mut self, input: Op::Input, _exts: ()) -> Self::Future {
+        (self)(input).map(ToResult::into_result)
     }
 }
 
-/// Adjoins state to a `fn(Input, State) -> Output` to create a [`Handler`].
-#[derive(Clone)]
-pub struct StatefulHandler<F, T> {
-    f: F,
-    state: T,
-}
-
-// fn(Input, State) -> Output
-impl<Input, Output, Error, F, Fut, T> Handler<Input, Output, Error> for StatefulHandler<F, T>
+// fn(Input, Arg0) -> Output
+impl<Op, F, Fut, Ext0> Handler<Op, (Ext0,)> for F
 where
-    T: Clone,
-    F: Fn(Input, T) -> Fut,
+    Op: OperationShape,
+    F: Fn(Op::Input, Ext0) -> Fut,
     Fut: Future,
-    Fut::Output: ToResult<Output, Error>,
+    Fut::Output: ToResult<Op::Output, Op::Error>,
 {
-    type Future = Map<Fut, fn(Fut::Output) -> Result<Output, Error>>;
+    type Future = Map<Fut, fn(Fut::Output) -> Result<Op::Output, Op::Error>>;
 
-    fn call(&mut self, req: Input) -> Self::Future {
-        let Self { f, state } = self;
-        f(req, state.clone()).map(ToResult::into_result)
+    fn call(&mut self, input: Op::Input, exts: (Ext0,)) -> Self::Future {
+        (self)(input, exts.0).map(ToResult::into_result)
     }
 }
 
-/// Provides the ability to [`AdjoinState::with_state`] on closures of the form
-/// `(Input, State) -> Output` converting them to a [`StatefulHandler`] and therefore causing them
-/// to implement [`Handler`].
-pub trait AdjoinState {
-    fn with_state<T>(self, state: T) -> StatefulHandler<Self, T>
-    where
-        Self: Sized,
-    {
-        StatefulHandler { f: self, state }
+// fn(Input, Arg0, Arg1) -> Output
+impl<Op, F, Fut, Ext0, Ext1> Handler<Op, (Ext0, Ext1)> for F
+where
+    Op: OperationShape,
+    F: Fn(Op::Input, Ext0, Ext1) -> Fut,
+    Fut: Future,
+    Fut::Output: ToResult<Op::Output, Op::Error>,
+{
+    type Future = Map<Fut, fn(Fut::Output) -> Result<Op::Output, Op::Error>>;
+
+    fn call(&mut self, input: Op::Input, exts: (Ext0, Ext1)) -> Self::Future {
+        (self)(input, exts.0, exts.1).map(ToResult::into_result)
     }
 }
 
-impl<F> AdjoinState for F {}
-
 /// An extension trait for [`Handler`].
-pub trait HandlerExt<Input, Output, Error>: Handler<Input, Output, Error> {
+pub trait HandlerExt<Op, Exts>: Handler<Op, Exts>
+where
+    Op: OperationShape,
+{
     /// Convert the [`Handler`] into a [`Service`].
-    fn into_service(self) -> IntoService<Output, Error, Self>
+    fn into_service(self) -> IntoService<Op, Self>
     where
         Self: Sized,
     {
         IntoService {
             handler: self,
-            _error: PhantomData,
-            _output: PhantomData,
+            _operation: PhantomData,
         }
     }
 }
 
-impl<Input, Output, Error, H> HandlerExt<Input, Output, Error> for H where
-    H: Handler<Input, Output, Error>
+impl<Op, Exts, H> HandlerExt<Op, Exts> for H
+where
+    Op: OperationShape,
+    H: Handler<Op, Exts>,
 {
 }
 
 /// A [`Service`] provided for every [`Handler`].
-pub struct IntoService<Output, Error, H> {
+pub struct IntoService<Op, H> {
     handler: H,
-    _output: PhantomData<Output>,
-    _error: PhantomData<Error>,
+    _operation: PhantomData<Op>,
 }
 
-impl<Output, Error, H> Clone for IntoService<Output, Error, H>
+impl<Op, H> Clone for IntoService<Op, H>
 where
     H: Clone,
 {
     fn clone(&self) -> Self {
         Self {
             handler: self.handler.clone(),
-            _output: PhantomData,
-            _error: PhantomData,
+            _operation: PhantomData,
         }
     }
 }
 
-impl<Input, Output, Error, H> Service<Input> for IntoService<Output, Error, H>
+impl<Op, Exts, H> Service<(Op::Input, Exts)> for IntoService<Op, H>
 where
-    H: Handler<Input, Output, Error>,
+    Op: OperationShape,
+    H: Handler<Op, Exts>,
 {
-    type Response = Output;
-    type Error = OperationError<Infallible, Error>;
-    type Future = MapErr<H::Future, fn(Error) -> Self::Error>;
+    type Response = Op::Output;
+    type Error = OperationError<Infallible, Op::Error>;
+    type Future = MapErr<H::Future, fn(Op::Error) -> Self::Error>;
 
     fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         Poll::Ready(Ok(()))
     }
 
-    fn call(&mut self, req: Input) -> Self::Future {
-        self.handler.call(req).map_err(OperationError::Smithy)
+    fn call(&mut self, (input, exts): (Op::Input, Exts)) -> Self::Future {
+        self.handler
+            .call(input, exts)
+            .map_err(OperationError::Smithy)
     }
 }
diff --git a/runtime/src/operation/http_conversions.rs b/runtime/src/operation/http_conversions.rs
new file mode 100644
index 0000000..6321c3f
--- /dev/null
+++ b/runtime/src/operation/http_conversions.rs
@@ -0,0 +1,107 @@
+use std::{convert::Infallible, future::Ready};
+
+use futures::{
+    future::{MapErr, MapOk, TryJoin},
+    Future, TryFutureExt,
+};
+use http_body::combinators::BoxBody;
+use hyper::body::Bytes;
+
+/// A protocol and operation aware conversion from [`http`] types to Smithy types.
+pub trait FromRequest<Protocol, Operation, B>: Sized {
+    /// Conversion failure.
+    type Error: IntoResponse<Protocol, Operation>;
+    type Future: Future<Output = Result<Self, Self::Error>>;
+
+    fn from_request(request: &mut http::Request<B>) -> Self::Future;
+}
+
+impl<P, Op, B> FromRequest<P, Op, B> for () {
+    type Error = Infallible;
+    type Future = Ready<Result<(), Infallible>>;
+
+    fn from_request(_request: &mut http::Request<B>) -> Self::Future {
+        std::future::ready(Ok(()))
+    }
+}
+
+impl<P, Op, B, Arg0> FromRequest<P, Op, B> for (Arg0,)
+where
+    Arg0: FromRequest<P, Op, B>,
+{
+    type Error = Arg0::Error;
+    type Future = MapOk<Arg0::Future, fn(Arg0) -> Self>;
+
+    fn from_request(request: &mut http::Request<B>) -> Self::Future {
+        Arg0::from_request(request).map_ok(|arg0| (arg0,))
+    }
+}
+
+/// An wrapper for state stored in the [`http::Extensions`] map.
+pub struct Extension<T>(pub T);
+
+impl<P, Op, B, T> FromRequest<P, Op, B> for Extension<T>
+where
+    T: Sync + Send + 'static,
+{
+    type Error = Infallible;
+    type Future = Ready<Result<Self, Self::Error>>;
+
+    fn from_request(request: &mut http::Request<B>) -> Self::Future {
+        std::future::ready(Ok(Extension(
+            request
+                .extensions_mut()
+                .remove()
+                .expect("request does not contain extension"),
+        )))
+    }
+}
+
+/// Represents one of two errors.
+/// 
+/// Implements [`IntoResponse`] if both inner also implement it.
+pub enum Either<A, B> {
+    Left(A),
+    Right(B),
+}
+
+impl<A, B, P, Op> IntoResponse<P, Op> for Either<A, B>
+where
+    A: IntoResponse<P, Op>,
+    B: IntoResponse<P, Op>,
+{
+    fn into_response(self) -> http::Response<BoxBody<Bytes, hyper::Error>> {
+        match self {
+            Either::Left(left) => left.into_response(),
+            Either::Right(right) => right.into_response(),
+        }
+    }
+}
+
+impl<P, Op, B, Arg0, Arg1> FromRequest<P, Op, B> for (Arg0, Arg1)
+where
+    Arg0: FromRequest<P, Op, B>,
+    Arg1: FromRequest<P, Op, B>,
+{
+    type Error = Either<Arg0::Error, Arg1::Error>;
+
+    type Future = TryJoin<
+        MapErr<Arg0::Future, fn(Arg0::Error) -> Self::Error>,
+        MapErr<Arg1::Future, fn(Arg1::Error) -> Self::Error>,
+    >;
+
+    fn from_request(_request: &mut http::Request<B>) -> Self::Future {
+        todo!()
+    }
+}
+
+/// A protocol and operation aware conversion from Smithy types to [`http`] types.
+pub trait IntoResponse<Protocol = (), Operation = ()>: Sized {
+    fn into_response(self) -> http::Response<BoxBody<Bytes, hyper::Error>>;
+}
+
+impl<P, Op> IntoResponse<P, Op> for Infallible {
+    fn into_response(self) -> http::Response<BoxBody<Bytes, hyper::Error>> {
+        match self {}
+    }
+}
diff --git a/runtime/src/operation/mod.rs b/runtime/src/operation/mod.rs
index 1c63e9f..b239f47 100644
--- a/runtime/src/operation/mod.rs
+++ b/runtime/src/operation/mod.rs
@@ -1,4 +1,6 @@
+mod flattened;
 mod handler;
+mod http_conversions;
 mod shape;
 mod upgrade;
 
@@ -7,7 +9,9 @@ use tower::{
     Layer,
 };
 
+pub use flattened::*;
 pub use handler::*;
+pub use http_conversions::*;
 pub use shape::*;
 pub use upgrade::*;
 
@@ -17,15 +21,18 @@ pub struct Operation<S, L = Identity> {
     layer: L,
 }
 
+type StackedUpgradeService<P, Op, E, B, L, S> =
+    <Stack<UpgradeLayer<P, Op, E, B>, L> as Layer<S>>::Service;
+
 impl<S, L> Operation<S, L> {
     /// Takes the [`Operation`], which contains the inner [`Service`](tower::Service), the HTTP [`Layer`] `L` and
     /// composes them together using [`UpgradeLayer`] for a specific protocol and [`OperationShape`].
     ///
     /// The composition is made explicit in the method constraints and return type.
-    pub fn upgrade<P, Op, B>(self) -> <Stack<UpgradeLayer<P, Op, B>, L> as Layer<S>>::Service
+    pub fn upgrade<P, Op, E, B>(self) -> StackedUpgradeService<P, Op, E, B, L, S>
     where
-        UpgradeLayer<P, Op, B>: Layer<S>,
-        L: Layer<<UpgradeLayer<P, Op, B> as Layer<S>>::Service>,
+        UpgradeLayer<P, Op, E, B>: Layer<S>,
+        L: Layer<<UpgradeLayer<P, Op, E, B> as Layer<S>>::Service>,
     {
         let Self { inner, layer } = self;
         let layer = Stack::new(UpgradeLayer::new(), layer);
@@ -33,21 +40,26 @@ impl<S, L> Operation<S, L> {
     }
 }
 
-impl<S> Operation<S> {
+impl<Op, S, PollError> Operation<IntoUnflattened<Op, S, PollError>> {
     /// Creates an [`Operation`] from a [`Service`](tower::Service).
-    pub fn from_service(inner: S) -> Self {
+    pub fn from_service<Exts>(inner: S) -> Self
+    where
+        Op: OperationShape,
+        S: Flattened<Op, Exts, PollError>,
+    {
         Self {
-            inner,
+            inner: inner.into_unflatten(),
             layer: Identity::new(),
         }
     }
 }
 
-impl<Output, Error, H> Operation<IntoService<Output, Error, H>> {
+impl<Op, H> Operation<IntoService<Op, H>> {
     /// Creates an [`Operation`] from a [`Handler`].
-    pub fn from_handler<Input>(handler: H) -> Self
+    pub fn from_handler<Exts>(handler: H) -> Self
     where
-        H: Handler<Input, Output, Error>,
+        Op: OperationShape,
+        H: Handler<Op, Exts>,
     {
         Self {
             inner: handler.into_service(),
@@ -57,8 +69,7 @@ impl<Output, Error, H> Operation<IntoService<Output, Error, H>> {
 }
 
 impl<S, L> Operation<S, L> {
-    /// Applies a [`Layer`] to the operation after is has been upgraded to a [`Service`] accepting
-    /// and returning [`http`] types.
+    /// Applies a [`Layer`] to the operation _after_ it has been upgraded via [`Operation::upgrade`].
     pub fn layer<NewL>(self, layer: NewL) -> Operation<S, Stack<L, NewL>> {
         Operation {
             inner: self.inner,
@@ -69,3 +80,12 @@ impl<S, L> Operation<S, L> {
 
 /// A marker struct indicating an [`Operation`] has not been set in a builder.
 pub struct OperationNotSet;
+
+/// The operation [`Service`] has two classes of failure modes - the failure models specified by
+/// the Smithy model and failures to [`Service::poll_ready`].
+pub enum OperationError<PollError, SmithyError> {
+    /// A [`Service::poll_ready`] failure occurred.
+    Poll(PollError),
+    /// An error modelled by the Smithy model occurred.
+    Smithy(SmithyError),
+}
diff --git a/runtime/src/operation/shape.rs b/runtime/src/operation/shape.rs
index d1a1dc6..32b0ed1 100644
--- a/runtime/src/operation/shape.rs
+++ b/runtime/src/operation/shape.rs
@@ -1,40 +1,33 @@
-use tower::Service;
-
-use super::{Handler, HandlerExt, IntoService, Operation, OperationError};
+use super::{Flattened, Handler, IntoService, IntoUnflattened, Operation};
 
 /// Mirrors the Smithy Operation shape.
 pub trait OperationShape {
     const NAME: &'static str;
 
+    /// The operation input.
     type Input;
+    /// The operation output.
     type Output;
+    /// The operation error.
     type Error;
 }
 
 /// An extension trait over [`OperationShape`].
 pub trait OperationShapeExt: OperationShape {
     /// Creates a new [`Operation`] for well-formed [`Handler`]s.
-    fn from_handler<H>(
-        handler: H,
-    ) -> Operation<IntoService<<Self as OperationShape>::Output, <Self as OperationShape>::Error, H>>
+    fn from_handler<H, Exts>(handler: H) -> Operation<IntoService<Self, H>>
     where
-        H: Handler<
-            <Self as OperationShape>::Input,
-            <Self as OperationShape>::Output,
-            <Self as OperationShape>::Error,
-        >,
+        H: Handler<Self, Exts>,
+        Self: Sized,
     {
-        Operation::from_service(handler.into_service())
+        Operation::from_handler(handler)
     }
 
-    /// Creates a new [`Operation`] for well-formed [`Service`]s.
-    fn from_service<S, PollError>(svc: S) -> Operation<S>
+    /// Creates a new [`Operation`] for well-formed [`Service`](tower::Service)s.
+    fn from_service<S, Exts, PollError>(svc: S) -> Operation<IntoUnflattened<Self, S, PollError>>
     where
-        S: Service<
-            <Self as OperationShape>::Input,
-            Response = <Self as OperationShape>::Output,
-            Error = OperationError<PollError, <Self as OperationShape>::Error>,
-        >,
+        S: Flattened<Self, Exts, PollError>,
+        Self: Sized,
     {
         Operation::from_service(svc)
     }
diff --git a/runtime/src/operation/upgrade.rs b/runtime/src/operation/upgrade.rs
index 0661e34..a1eae54 100644
--- a/runtime/src/operation/upgrade.rs
+++ b/runtime/src/operation/upgrade.rs
@@ -1,5 +1,6 @@
+#![allow(clippy::type_complexity)]
+
 use std::{
-    convert::Infallible,
     future::Future,
     marker::PhantomData,
     pin::Pin,
@@ -12,82 +13,68 @@ use hyper::body::Bytes;
 use pin_project_lite::pin_project;
 use tower::{Layer, Service};
 
-use super::{OperationError, OperationShape};
-
-/// A protocol and operation aware conversion from [`http`] types to Smithy types.
-pub trait FromRequest<Protocol, Operation, B>: Sized {
-    /// Conversion failure.
-    type Error: IntoResponse<Protocol, Operation>;
-    type Future: Future<Output = Result<Self, Self::Error>>;
-
-    fn from_request(request: http::Request<B>) -> Self::Future;
-}
+use crate::operation::FromRequest;
 
-/// A protocol and operation aware conversion from Smithy types to [`http`] types.
-pub trait IntoResponse<Protocol = (), Operation = ()>: Sized {
-    fn into_response(self) -> http::Response<BoxBody<Bytes, hyper::Error>>;
-}
-
-impl<P, Op> IntoResponse<P, Op> for Infallible {
-    fn into_response(self) -> http::Response<BoxBody<Bytes, hyper::Error>> {
-        match self {}
-    }
-}
+use super::{IntoResponse, OperationError, OperationShape};
 
 /// A [`Layer`] responsible for taking an operation [`Service`], accepting and returning Smithy
 /// types and converting it into a [`Service`] taking and returning [`http`] types.
 ///
 /// See [`Upgrade`].
 #[derive(Debug, Clone)]
-pub struct UpgradeLayer<Protocol, Operation, B> {
+pub struct UpgradeLayer<Protocol, Operation, Exts, B> {
     _protocol: PhantomData<Protocol>,
     _operation: PhantomData<Operation>,
+    _exts: PhantomData<Exts>,
     _body: PhantomData<B>,
 }
 
-impl<Protocol, Operation, B> Default for UpgradeLayer<Protocol, Operation, B> {
+impl<P, Op, E, B> Default for UpgradeLayer<P, Op, E, B> {
     fn default() -> Self {
         Self {
             _protocol: PhantomData,
             _operation: PhantomData,
+            _exts: PhantomData,
             _body: PhantomData,
         }
     }
 }
 
-impl<Protocol, Operation, B> UpgradeLayer<Protocol, Operation, B> {
+impl<Protocol, Operation, Exts, B> UpgradeLayer<Protocol, Operation, Exts, B> {
     /// Creates a new [`UpgradeLayer`].
     pub fn new() -> Self {
         Self::default()
     }
 }
 
-impl<S, P, Op, B> Layer<S> for UpgradeLayer<P, Op, B> {
-    type Service = Upgrade<P, Op, B, S>;
+impl<S, P, Op, E, B> Layer<S> for UpgradeLayer<P, Op, E, B> {
+    type Service = Upgrade<P, Op, E, B, S>;
 
     fn layer(&self, inner: S) -> Self::Service {
         Upgrade {
             _protocol: PhantomData,
             _operation: PhantomData,
             _body: PhantomData,
+            _exts: PhantomData,
             inner,
         }
     }
 }
 
 /// A alias allowing for quick access to [`UpgradeLayer`]s target [`Service`].
-pub type UpgradedService<P, Op, B, S> = <UpgradeLayer<P, Op, B> as Layer<S>>::Service;
+pub type UpgradedService<P, Op, E, B, S> = <UpgradeLayer<P, Op, E, B> as Layer<S>>::Service;
 
 /// A [`Service`] responsible for wrapping an operation [`Service`] accepting and returning Smithy
 /// types, and converting it into a [`Service`] accepting and returning [`http`] types.
-pub struct Upgrade<Protocol, Operation, B, S> {
+pub struct Upgrade<Protocol, Operation, Exts, B, S> {
     _protocol: PhantomData<Protocol>,
     _operation: PhantomData<Operation>,
+    _exts: PhantomData<Exts>,
     _body: PhantomData<B>,
     inner: S,
 }
 
-impl<P, Op, B, S> Clone for Upgrade<P, Op, B, S>
+impl<P, Op, E, B, S> Clone for Upgrade<P, Op, E, B, S>
 where
     S: Clone,
 {
@@ -96,6 +83,7 @@ where
             _protocol: PhantomData,
             _operation: PhantomData,
             _body: PhantomData,
+            _exts: PhantomData,
             inner: self.inner.clone(),
         }
     }
@@ -103,15 +91,15 @@ where
 
 pin_project! {
     /// The [`Service::Future`] of [`Upgrade`].
-    pub struct UpgradeFuture<Protocol, Operation, B, S>
+    pub struct UpgradeFuture<Protocol, Operation, Exts, B, S>
     where
         Operation: OperationShape,
-        Operation::Input: FromRequest<Protocol, Operation, B>,
-        S: Service<Operation::Input>,
+        (Operation::Input, Exts): FromRequest<Protocol, Operation, B>,
+        S: Service<(Operation::Input, Exts)>,
     {
         service: S,
         #[pin]
-        inner: Inner<<Operation::Input as FromRequest<Protocol, Operation, B>>::Future, S::Future>
+        inner: Inner<<(Operation::Input, Exts) as FromRequest<Protocol, Operation, B>>::Future, S::Future>
     }
 }
 
@@ -130,20 +118,22 @@ pin_project! {
     }
 }
 
-impl<P, Op, B, S, PollError, OpError> Future for UpgradeFuture<P, Op, B, S>
+impl<P, Op, E, B, S, PollError, OpError> Future for UpgradeFuture<P, Op, E, B, S>
 where
-    // Op is used to specify the operation shape
+    // `Op` is used to specify the operation shape
     Op: OperationShape,
-    // Smithy input must be convert from a HTTP request
+    // Smithy input must convert from a HTTP request
     Op::Input: FromRequest<P, Op, B>,
-    // Smithy output must be convert into a HTTP response
+    // Smithy output must convert into a HTTP response
     Op::Output: IntoResponse<P, Op>,
     // Smithy error must convert into a HTTP response
     OpError: IntoResponse<P, Op>,
 
+    // Must be able to convert extensions
+    E: FromRequest<P, Op, B>,
+
     // The signature of the inner service is correct
-    S: Service<Op::Input, Response = Op::Output, Error = OperationError<PollError, OpError>>
-        + Clone,
+    S: Service<(Op::Input, E), Response = Op::Output, Error = OperationError<PollError, OpError>>,
 {
     type Output = Result<http::Response<BoxBody<Bytes, hyper::Error>>, PollError>;
 
@@ -178,18 +168,27 @@ where
     }
 }
 
-impl<P, Op, B, S, PollError, OpError> Service<http::Request<B>> for Upgrade<P, Op, B, S>
+impl<P, Op, E, B, S, PollError, OpError> Service<http::Request<B>> for Upgrade<P, Op, E, B, S>
 where
+    // `Op` is used to specify the operation shape
     Op: OperationShape,
+    // Smithy input must convert from a HTTP request
     Op::Input: FromRequest<P, Op, B>,
+    // Smithy output must convert into a HTTP response
     Op::Output: IntoResponse<P, Op>,
+    // Smithy error must convert into a HTTP response
     OpError: IntoResponse<P, Op>,
-    S: Service<Op::Input, Response = Op::Output, Error = OperationError<PollError, OpError>>
+
+    // Must be able to convert extensions
+    E: FromRequest<P, Op, B>,
+
+    // The signature of the inner service is correct
+    S: Service<(Op::Input, E), Response = Op::Output, Error = OperationError<PollError, OpError>>
         + Clone,
 {
     type Response = http::Response<BoxBody<Bytes, hyper::Error>>;
     type Error = PollError;
-    type Future = UpgradeFuture<P, Op, B, S>;
+    type Future = UpgradeFuture<P, Op, E, B, S>;
 
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         self.inner.poll_ready(cx).map_err(|err| match err {
@@ -198,11 +197,11 @@ where
         })
     }
 
-    fn call(&mut self, req: http::Request<B>) -> Self::Future {
+    fn call(&mut self, mut req: http::Request<B>) -> Self::Future {
         UpgradeFuture {
             service: self.inner.clone(),
             inner: Inner::FromRequest {
-                inner: <Op::Input as FromRequest<P, Op, B>>::from_request(req),
+                inner: <(Op::Input, E) as FromRequest<P, Op, B>>::from_request(&mut req),
             },
         }
     }