From 9b5ecd5bc4ea309d7eb86144a728f89469ea5b71 Mon Sep 17 00:00:00 2001
From: Quentin Gliech <quenting@element.io>
Date: Fri, 25 Mar 2022 10:09:01 +0100
Subject: [PATCH] Make the ServerLayer work properly with axum

---
 Cargo.lock                        |  2 ++
 crates/cli/src/commands/server.rs |  4 +++-
 crates/handlers/Cargo.toml        |  3 ++-
 crates/http/Cargo.toml            |  1 +
 crates/http/src/layers/server.rs  | 31 ++++++++++---------------------
 5 files changed, 18 insertions(+), 23 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 121a294f..05aa924a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2074,6 +2074,7 @@ dependencies = [
  "mas-config",
  "mas-data-model",
  "mas-email",
+ "mas-http",
  "mas-iana",
  "mas-jose",
  "mas-static-files",
@@ -2112,6 +2113,7 @@ dependencies = [
  "opentelemetry",
  "opentelemetry-http",
  "opentelemetry-semantic-conventions",
+ "pin-project-lite",
  "rustls 0.20.4",
  "serde",
  "serde_json",
diff --git a/crates/cli/src/commands/server.rs b/crates/cli/src/commands/server.rs
index df9ddf21..9a0b2ee5 100644
--- a/crates/cli/src/commands/server.rs
+++ b/crates/cli/src/commands/server.rs
@@ -25,6 +25,7 @@ use hyper::Server;
 use mas_axum_utils::UrlBuilder;
 use mas_config::RootConfig;
 use mas_email::{MailTransport, Mailer};
+use mas_http::ServerLayer;
 use mas_storage::MIGRATOR;
 use mas_tasks::TaskQueue;
 use mas_templates::Templates;
@@ -215,7 +216,8 @@ impl Options {
             &encrypter,
             &mailer,
             &url_builder,
-        );
+        )
+        .layer(ServerLayer::default());
 
         info!("Listening on http://{}", listener.local_addr().unwrap());
 
diff --git a/crates/handlers/Cargo.toml b/crates/handlers/Cargo.toml
index 9204cb27..0d4c6dde 100644
--- a/crates/handlers/Cargo.toml
+++ b/crates/handlers/Cargo.toml
@@ -56,16 +56,17 @@ rand = "0.8.5"
 headers = "0.3.7"
 
 oauth2-types = { path = "../oauth2-types" }
+mas-axum-utils = {  path = "../axum-utils" }
 mas-config = { path = "../config" }
 mas-data-model = { path = "../data-model" }
 mas-email = { path = "../email" }
+mas-http = {  path = "../http" }
 mas-iana = { path = "../iana" }
 mas-jose = { path = "../jose" }
 mas-static-files = { path = "../static-files" }
 mas-storage = { path = "../storage" }
 mas-templates = { path = "../templates" }
 mas-warp-utils = { path = "../warp-utils" }
-mas-axum-utils = {  path = "../axum-utils" }
 
 [dev-dependencies]
 indoc = "1.0.4"
diff --git a/crates/http/Cargo.toml b/crates/http/Cargo.toml
index 16bf7fa1..edb4542c 100644
--- a/crates/http/Cargo.toml
+++ b/crates/http/Cargo.toml
@@ -15,6 +15,7 @@ hyper-rustls = { version = "0.23.0", features = ["http1", "http2"] }
 opentelemetry = "0.17.0"
 opentelemetry-http = "0.6.0"
 opentelemetry-semantic-conventions = "0.9.0"
+pin-project-lite = "0.2.8"
 rustls = "0.20.4"
 serde = "1.0.136"
 serde_json = "1.0.79"
diff --git a/crates/http/src/layers/server.rs b/crates/http/src/layers/server.rs
index 24d93323..4a136ddf 100644
--- a/crates/http/src/layers/server.rs
+++ b/crates/http/src/layers/server.rs
@@ -12,45 +12,34 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::{marker::PhantomData, time::Duration};
+use std::marker::PhantomData;
 
 use http::{Request, Response};
-use http_body::combinators::BoxBody;
-use tower::{
-    timeout::TimeoutLayer, util::BoxCloneService, Layer, Service, ServiceBuilder, ServiceExt,
-};
-use tower_http::compression::{CompressionBody, CompressionLayer};
+use tower::{util::BoxCloneService, Layer, Service, ServiceBuilder, ServiceExt};
+use tower_http::{compression::CompressionBody, ServiceBuilderExt};
 
 use super::otel::TraceLayer;
-use crate::BoxError;
 
 #[derive(Debug, Default)]
 pub struct ServerLayer<ReqBody> {
     _t: PhantomData<ReqBody>,
 }
 
-impl<ReqBody, ResBody, S, E> Layer<S> for ServerLayer<ReqBody>
+impl<ReqBody, ResBody, S> Layer<S> for ServerLayer<ReqBody>
 where
-    S: Service<Request<ReqBody>, Response = Response<ResBody>, Error = E> + Clone + Send + 'static,
-    ReqBody: http_body::Body + 'static,
-    ResBody: http_body::Body + Sync + Send + 'static,
-    ResBody::Error: std::fmt::Display + 'static,
+    S: Service<Request<ReqBody>, Response = Response<ResBody>> + Clone + Send + 'static,
     S::Future: Send + 'static,
-    E: std::error::Error + Into<BoxError>,
+    S::Error: std::fmt::Display,
+    ReqBody: http_body::Body + 'static,
+    ResBody: http_body::Body + Send + 'static,
 {
     #[allow(clippy::type_complexity)]
-    type Service = BoxCloneService<
-        Request<ReqBody>,
-        Response<CompressionBody<BoxBody<ResBody::Data, ResBody::Error>>>,
-        BoxError,
-    >;
+    type Service = BoxCloneService<Request<ReqBody>, Response<CompressionBody<ResBody>>, S::Error>;
 
     fn layer(&self, inner: S) -> Self::Service {
         ServiceBuilder::new()
-            .layer(CompressionLayer::new())
+            .compression()
             .layer(TraceLayer::http_server())
-            .map_response(|r: Response<_>| r.map(BoxBody::new))
-            .layer(TimeoutLayer::new(Duration::from_secs(10)))
             .service(inner)
             .boxed_clone()
     }