Better Tower layers

crates/axum-utils/src/client_authorization.rs

@@ -184,7 +184,8 @@ fn jwks_key_store(jwks: &JwksOrJwksUri) -> Either<StaticJwksStore, DynamicJwksSt
 
             // TODO: get the client from somewhere else?
             let exporter = mas_http::client("fetch-jwks")
-                .json::<JsonWebKeySet>()
+                .response_body_to_bytes()
+                .json_response::<JsonWebKeySet>()
                 .map_request(move |_: ()| {
                     http::Request::builder()
                         .method("GET")

crates/cli/src/commands/debug.rs

@@ -89,7 +89,9 @@ impl Options {
                 json: true,
                 url,
             } => {
-                let mut client = mas_http::client("cli-debug-http").json();
+                let mut client = mas_http::client("cli-debug-http")
+                    .response_body_to_bytes()
+                    .json_response();
                 let request = hyper::Request::builder()
                     .uri(url)
                     .body(hyper::Body::empty())?;

crates/http/src/ext.rs

@@ -14,9 +14,14 @@
 
 use http::header::HeaderName;
 use once_cell::sync::OnceCell;
+use tower::{layer::util::Stack, ServiceBuilder};
 use tower_http::cors::CorsLayer;
 
-use crate::layers::json::Json;
+use crate::layers::{
+    body_to_bytes::{BodyToBytes, BodyToBytesLayer},
+    json_request::{JsonRequest, JsonRequestLayer},
+    json_response::{JsonResponse, JsonResponseLayer},
+};
 
 static PROPAGATOR_HEADERS: OnceCell<Vec<HeaderName>> = OnceCell::new();
 
@@ -60,11 +65,37 @@ impl CorsLayerExt for CorsLayer {
 }
 
 pub trait ServiceExt: Sized {
-    fn json<T>(self) -> Json<Self, T>;
+    fn response_body_to_bytes(self) -> BodyToBytes<Self> {
+        BodyToBytes::new(self)
     }
 
-impl<S> ServiceExt for S {
-    fn json<T>(self) -> Json<Self, T> {
-        Json::new(self)
+    fn json_response<T>(self) -> JsonResponse<Self, T> {
+        JsonResponse::new(self)
+    }
+
+    fn json_request<T>(self) -> JsonRequest<Self, T> {
+        JsonRequest::new(self)
+    }
+}
+
+impl<S> ServiceExt for S {}
+
+pub trait ServiceBuilderExt<L>: Sized {
+    fn response_to_bytes(self) -> ServiceBuilder<Stack<BodyToBytesLayer, L>>;
+    fn json_response<T>(self) -> ServiceBuilder<Stack<JsonResponseLayer<T>, L>>;
+    fn json_request<T>(self) -> ServiceBuilder<Stack<JsonRequestLayer<T>, L>>;
+}
+
+impl<L> ServiceBuilderExt<L> for ServiceBuilder<L> {
+    fn response_to_bytes(self) -> ServiceBuilder<Stack<BodyToBytesLayer, L>> {
+        self.layer(BodyToBytesLayer::default())
+    }
+
+    fn json_response<T>(self) -> ServiceBuilder<Stack<JsonResponseLayer<T>, L>> {
+        self.layer(JsonResponseLayer::default())
+    }
+
+    fn json_request<T>(self) -> ServiceBuilder<Stack<JsonRequestLayer<T>, L>> {
+        self.layer(JsonRequestLayer::default())
     }
 }

crates/http/src/layers/body_to_bytes.rs

@@ -0,0 +1,96 @@
+// Copyright 2022 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use bytes::Bytes;
+use futures_util::future::BoxFuture;
+use http::{Request, Response};
+use http_body::Body;
+use thiserror::Error;
+use tower::{Layer, Service};
+
+#[derive(Debug, Error)]
+pub enum Error<ServiceError, BodyError> {
+    #[error(transparent)]
+    Service { inner: ServiceError },
+
+    #[error(transparent)]
+    Body { inner: BodyError },
+}
+
+impl<S, B> Error<S, B> {
+    fn service(inner: S) -> Self {
+        Self::Service { inner }
+    }
+
+    fn body(inner: B) -> Self {
+        Self::Body { inner }
+    }
+}
+
+#[derive(Clone)]
+pub struct BodyToBytes<S> {
+    inner: S,
+}
+
+impl<S> BodyToBytes<S> {
+    pub const fn new(inner: S) -> Self {
+        Self { inner }
+    }
+}
+
+impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for BodyToBytes<S>
+where
+    S: Service<Request<ReqBody>, Response = Response<ResBody>>,
+    S::Future: Send + 'static,
+    ResBody: Body + Send,
+    ResBody::Data: Send,
+{
+    type Error = Error<S::Error, ResBody::Error>;
+    type Response = Response<Bytes>;
+    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+    fn poll_ready(
+        &mut self,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx).map_err(Error::service)
+    }
+
+    fn call(&mut self, request: Request<ReqBody>) -> Self::Future {
+        let inner = self.inner.call(request);
+
+        let fut = async {
+            let response = inner.await.map_err(Error::service)?;
+            let (parts, body) = response.into_parts();
+
+            let body = hyper::body::to_bytes(body).await.map_err(Error::body)?;
+
+            let response = Response::from_parts(parts, body);
+            Ok(response)
+        };
+
+        Box::pin(fut)
+    }
+}
+
+#[derive(Default, Clone, Copy)]
+pub struct BodyToBytesLayer;
+
+impl<S> Layer<S> for BodyToBytesLayer {
+    type Service = BodyToBytes<S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        BodyToBytes::new(inner)
+    }
+}

crates/http/src/layers/json_request.rs

@@ -0,0 +1,123 @@
+// Copyright 2022 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::{future::Ready, marker::PhantomData, task::Poll};
+
+use bytes::Bytes;
+use futures_util::{
+    future::{Either, MapErr},
+    FutureExt, TryFutureExt,
+};
+use http::{header::CONTENT_TYPE, HeaderValue, Request};
+use http_body::Full;
+use serde::Serialize;
+use thiserror::Error;
+use tower::{Layer, Service};
+
+#[derive(Debug, Error)]
+pub enum Error<Service> {
+    #[error(transparent)]
+    Service { inner: Service },
+
+    #[error("could not serialize JSON payload")]
+    Json {
+        #[source]
+        inner: serde_json::Error,
+    },
+}
+
+impl<S> Error<S> {
+    fn service(source: S) -> Self {
+        Self::Service { inner: source }
+    }
+
+    fn json(source: serde_json::Error) -> Self {
+        Self::Json { inner: source }
+    }
+}
+
+#[derive(Clone)]
+pub struct JsonRequest<S, T> {
+    inner: S,
+    _t: PhantomData<T>,
+}
+
+impl<S, T> JsonRequest<S, T> {
+    pub const fn new(inner: S) -> Self {
+        Self {
+            inner,
+            _t: PhantomData,
+        }
+    }
+}
+
+impl<S, T> Service<Request<T>> for JsonRequest<S, T>
+where
+    S: Service<Request<Full<Bytes>>>,
+    S::Future: Send + 'static,
+    S::Error: 'static,
+    T: Serialize,
+{
+    type Error = Error<S::Error>;
+    type Response = S::Response;
+    type Future = Either<
+        Ready<Result<Self::Response, Self::Error>>,
+        MapErr<S::Future, fn(S::Error) -> Self::Error>,
+    >;
+
+    fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx).map_err(Error::service)
+    }
+
+    fn call(&mut self, request: Request<T>) -> Self::Future {
+        let (mut parts, body) = request.into_parts();
+
+        parts
+            .headers
+            .insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
+
+        let body = match serde_json::to_vec(&body) {
+            Ok(body) => Full::new(Bytes::from(body)),
+            Err(err) => return std::future::ready(Err(Error::json(err))).left_future(),
+        };
+
+        let request = Request::from_parts(parts, body);
+
+        self.inner
+            .call(request)
+            .map_err(Error::service as fn(S::Error) -> Self::Error)
+            .right_future()
+    }
+}
+
+#[derive(Clone, Copy)]
+pub struct JsonRequestLayer<T> {
+    _t: PhantomData<T>,
+}
+
+impl<T> Default for JsonRequestLayer<T> {
+    fn default() -> Self {
+        Self {
+            _t: PhantomData::default(),
+        }
+    }
+}
+
+impl<S, T> Layer<S> for JsonRequestLayer<T> {
+    type Service = JsonRequest<S, T>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        JsonRequest::new(inner)
+    }
+}

crates/http/src/layers/json_response.rs

@@ -14,24 +14,18 @@
 
 use std::{marker::PhantomData, task::Poll};
 
-use futures_util::future::BoxFuture;
+use bytes::Buf;
+use futures_util::FutureExt;
 use http::{header::ACCEPT, HeaderValue, Request, Response};
-use http_body::Body;
 use serde::de::DeserializeOwned;
 use thiserror::Error;
 use tower::{Layer, Service};
 
 #[derive(Debug, Error)]
-pub enum Error<Service, Body> {
+pub enum Error<Service> {
     #[error(transparent)]
     Service { inner: Service },
 
-    #[error("failed to fully read the request body")]
-    Body {
-        #[source]
-        inner: Body,
-    },
-
     #[error("could not parse JSON payload")]
     Json {
         #[source]
@@ -39,27 +33,23 @@ pub enum Error<Service, Body> {
     },
 }
 
-impl<S, B> Error<S, B> {
+impl<S> Error<S> {
     fn service(source: S) -> Self {
         Self::Service { inner: source }
     }
 
-    fn body(source: B) -> Self {
-        Self::Body { inner: source }
-    }
-
     fn json(source: serde_json::Error) -> Self {
         Self::Json { inner: source }
     }
 }
 
 #[derive(Clone)]
-pub struct Json<S, T> {
+pub struct JsonResponse<S, T> {
     inner: S,
     _t: PhantomData<T>,
 }
 
-impl<S, T> Json<S, T> {
+impl<S, T> JsonResponse<S, T> {
     pub const fn new(inner: S) -> Self {
         Self {
             inner,
@@ -68,59 +58,64 @@ impl<S, T> Json<S, T> {
     }
 }
 
-impl<S, T, B, C> Service<Request<B>> for Json<S, T>
+impl<S, T, B, C> Service<Request<B>> for JsonResponse<S, T>
 where
     S: Service<Request<B>, Response = Response<C>>,
     S::Future: Send + 'static,
-    C: Body + Send + 'static,
-    C::Data: Send + 'static,
+    C: Buf,
     T: DeserializeOwned,
 {
-    type Error = Error<S::Error, C::Error>;
+    type Error = Error<S::Error>;
     type Response = Response<T>;
-    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
+    type Future = futures_util::future::Map<
+        S::Future,
+        fn(Result<Response<C>, S::Error>) -> Result<Self::Response, Self::Error>,
+    >;
 
     fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<(), Self::Error>> {
         self.inner.poll_ready(cx).map_err(Error::service)
     }
 
     fn call(&mut self, mut request: Request<B>) -> Self::Future {
+        fn mapper<C, T, E>(res: Result<Response<C>, E>) -> Result<Response<T>, Error<E>>
+        where
+            C: Buf,
+            T: DeserializeOwned,
+        {
+            let response = res.map_err(Error::service)?;
+            let (parts, body) = response.into_parts();
+
+            let body = serde_json::from_reader(body.reader()).map_err(Error::json)?;
+
+            let res = Response::from_parts(parts, body);
+            Ok(res)
+        }
+
         request
             .headers_mut()
             .insert(ACCEPT, HeaderValue::from_static("application/json"));
 
-        let fut = self.inner.call(request);
-
-        let fut = async {
-            let response = fut.await.map_err(Error::service)?;
-            let (parts, body) = response.into_parts();
-
-            futures_util::pin_mut!(body);
-            let bytes = hyper::body::to_bytes(&mut body)
-                .await
-                .map_err(Error::body)?;
-
-            let body = serde_json::from_slice(&bytes).map_err(Error::json)?;
-
-            let res = Response::from_parts(parts, body);
-            Ok(res)
-        };
-
-        Box::pin(fut)
+        self.inner.call(request).map(mapper::<C, T, S::Error>)
     }
 }
 
-#[derive(Default, Clone, Copy)]
-pub struct JsonResponseLayer<T, ReqBody>(PhantomData<(T, ReqBody)>);
+#[derive(Clone, Copy)]
+pub struct JsonResponseLayer<T> {
+    _t: PhantomData<T>,
+}
 
-impl<ReqBody, ResBody, S, T> Layer<S> for JsonResponseLayer<T, ReqBody>
-where
-    S: Service<Request<ReqBody>, Response = Response<ResBody>>,
-    T: serde::de::DeserializeOwned,
-{
-    type Service = Json<S, T>;
+impl<T> Default for JsonResponseLayer<T> {
+    fn default() -> Self {
+        Self {
+            _t: PhantomData::default(),
+        }
+    }
+}
+
+impl<S, T> Layer<S> for JsonResponseLayer<T> {
+    type Service = JsonResponse<S, T>;
 
     fn layer(&self, inner: S) -> Self::Service {
-        Json::new(inner)
+        JsonResponse::new(inner)
     }
 }

crates/http/src/layers/mod.rs

@@ -12,7 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+pub(crate) mod body_to_bytes;
 pub(crate) mod client;
-pub(crate) mod json;
+pub(crate) mod json_request;
+pub(crate) mod json_response;
 pub mod otel;
 pub(crate) mod server;

crates/http/src/lib.rs

@@ -50,7 +50,10 @@ mod layers;
 pub use self::{
     ext::{set_propagator, CorsLayerExt, ServiceExt as HttpServiceExt},
     future_service::FutureService,
-    layers::{client::ClientLayer, json::JsonResponseLayer, otel, server::ServerLayer},
+    layers::{
+        body_to_bytes::BodyToBytesLayer, client::ClientLayer, json_request::JsonRequestLayer,
+        json_response::JsonResponseLayer, otel, server::ServerLayer,
+    },
 };
 
 pub(crate) type BoxError = Box<dyn std::error::Error + Send + Sync>;