fix: respond with 500 when middleware fails

LeoBorai · LeoBorai · commit 1dac121b9c32 · 2021-07-12T19:41:45.000-04:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -44,6 +44,7 @@ tokio = { version = "1", features = ["fs", "rt-multi-thread", "signal", "macros"
 tokio-rustls = "0.22"
 toml = "0.5"
 serde = { version = "1", features = ["derive"] }
+serde_json = "1.0.64"
 structopt = { version = "0.3", default-features = false }
 
 [dev-dependencies]
diff --git a/src/server/middleware/gzip.rs b/src/server/middleware/gzip.rs
@@ -1,16 +1,26 @@
-use http::{Request, Response};
+use http::{Request, Response, StatusCode};
 use hyper::Body;
 use std::sync::Arc;
 use tokio::sync::Mutex;
 
 use crate::addon::compression::gzip::compress_http_response;
+use crate::utils::error::make_http_error_response;
 
 use super::MiddlewareAfter;
 
 pub fn make_gzip_compression_middleware() -> MiddlewareAfter {
     Box::new(
         move |request: Arc<Request<Body>>, response: Arc<Mutex<Response<Body>>>| {
-            Box::pin(async move { compress_http_response(request, response).await })
+            Box::pin(async move {
+                compress_http_response(request, response)
+                    .await
+                    .map_err(|err| {
+                        make_http_error_response(
+                            StatusCode::INTERNAL_SERVER_ERROR,
+                            &err.to_string(),
+                        )
+                    })
+            })
         },
     )
 }
diff --git a/src/server/middleware/mod.rs b/src/server/middleware/mod.rs
@@ -1,7 +1,7 @@
 pub mod cors;
 pub mod gzip;
 
-use anyhow::{Error, Result};
+use anyhow::Error;
 use futures::Future;
 use http::{Request, Response};
 use hyper::Body;
@@ -15,15 +15,34 @@ use crate::config::Config;
 use self::cors::make_cors_middleware;
 use self::gzip::make_gzip_compression_middleware;
 
-pub type MiddlewareBefore = Box<dyn Fn(&mut Request<Body>) + Send + Sync>;
+/// Middleware chain `Result` which specifies the `Err` variant as a
+/// HTTP response.
+pub type Result = std::result::Result<(), Response<Body>>;
+
+/// Middleware chain to execute before the main handler digests the
+/// HTTP request. No HTTP response is available at this point.
+pub type MiddlewareBefore = Box<
+    dyn Fn(&mut Request<Body>) -> Pin<Box<dyn Future<Output = Result> + Send + Sync>> + Send + Sync,
+>;
+
+/// Middleware chain to execute after the main handler digests the
+/// HTTP request. The HTTP response is created by the handler and
+/// consumed by every middleware after chain.
 pub type MiddlewareAfter = Box<
     dyn Fn(
             Arc<Request<Body>>,
             Arc<Mutex<Response<Body>>>,
-        ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + Sync>>
+        ) -> Pin<Box<dyn Future<Output = Result> + Send + Sync>>
         + Send
         + Sync,
 >;
+
+/// The main handler for the HTTP request, a HTTP response is created
+/// as a result of this handler.
+///
+/// This handler will be executed against the HTTP request after every
+/// "Middleware Before" chain is executed but before any "Middleware After"
+/// chain is executed
 pub type Handler = Box<
     dyn Fn(Arc<Request<Body>>) -> Pin<Box<dyn Future<Output = Response<Body>> + Send + Sync>>
         + Send
@@ -57,16 +76,18 @@ impl Middleware {
     /// with the HTTP Response from the handler
     pub async fn handle(&self, mut request: Request<Body>, handler: Handler) -> Response<Body> {
         for fx in self.before.iter() {
-            fx(&mut request);
+            if let Err(err) = fx(&mut request).await {
+                return err;
+            }
         }
 
         let request = Arc::new(request);
         let response = handler(Arc::clone(&request)).await;
         let response = Arc::new(Mutex::new(response));
 
         for fx in self.after.iter() {
-            if let Err(error) = fx(Arc::clone(&request), Arc::clone(&response)).await {
-                eprintln!("{:?}", error);
+            if let Err(err) = fx(Arc::clone(&request), Arc::clone(&response)).await {
+                return err;
             }
         }
 
@@ -88,7 +109,7 @@ impl Default for Middleware {
 impl TryFrom<Arc<Config>> for Middleware {
     type Error = Error;
 
-    fn try_from(config: Arc<Config>) -> Result<Self, Self::Error> {
+    fn try_from(config: Arc<Config>) -> std::result::Result<Self, Self::Error> {
         let mut middleware = Middleware::default();
 
         if let Some(cors_config) = config.cors() {
diff --git a/src/utils/error.rs b/src/utils/error.rs
@@ -0,0 +1,22 @@
+use http::{Response, StatusCode};
+use hyper::Body;
+use serde::Serialize;
+
+#[derive(Debug, Serialize)]
+struct ErrorResponseBody {
+    status_code: u16,
+    message: String,
+}
+
+pub fn make_http_error_response(status: StatusCode, message: &str) -> Response<Body> {
+    Response::builder()
+        .status(status)
+        .body(Body::from(
+            serde_json::ser::to_string(&ErrorResponseBody {
+                status_code: status.as_u16(),
+                message: message.to_string(),
+            })
+            .unwrap(),
+        ))
+        .unwrap()
+}
diff --git a/src/utils/mod.rs b/src/utils/mod.rs
@@ -1 +1,2 @@
+pub mod error;
 pub mod signal;
