From 3b93bff602239ed6a947183292ed35e126df9900 Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Mon, 2 Apr 2018 21:37:00 -0700
Subject: [PATCH] add ErrorHandlers middleware

---
 guide/src/qs_10.md            |  34 ++++++++++
 src/middleware/errhandlers.rs | 115 ++++++++++++++++++++++++++++++++++
 src/middleware/mod.rs         |   2 +
 3 files changed, 151 insertions(+)
 create mode 100644 src/middleware/errhandlers.rs

diff --git a/guide/src/qs_10.md b/guide/src/qs_10.md
index c69d26f6..9a4addba 100644
--- a/guide/src/qs_10.md
+++ b/guide/src/qs_10.md
@@ -210,3 +210,37 @@ fn main() {
 #     let _ = sys.run();
 }
 ```
+
+## Error handlers
+
+`ErrorHandlers` middleware allows to provide custom handlers for responses.
+
+You can use `ErrorHandlers::handler()` method  to register a custom error handler
+for specific status code. You can modify existing response or create completly new
+one. Error handler can return response immediately or return future that resolves
+to a response.
+
+```rust
+# extern crate actix_web;
+use actix_web::{
+    App, HttpRequest, HttpResponse, Result,
+    http, middleware::Response, middleware::ErrorHandlers};
+
+fn render_500<S>(_: &mut HttpRequest<S>, resp: HttpResponse) -> Result<Response> {
+   let mut builder = resp.into_builder();
+   builder.header(http::header::CONTENT_TYPE, "application/json");
+   Ok(Response::Done(builder.into()))
+}
+
+fn main() {
+    let app = App::new()
+        .middleware(
+            ErrorHandlers::new()
+                .handler(http::StatusCode::INTERNAL_SERVER_ERROR, render_500))
+        .resource("/test", |r| {
+             r.method(http::Method::GET).f(|_| HttpResponse::Ok());
+             r.method(http::Method::HEAD).f(|_| HttpResponse::MethodNotAllowed());
+        })
+        .finish();
+}
+```
diff --git a/src/middleware/errhandlers.rs b/src/middleware/errhandlers.rs
new file mode 100644
index 00000000..2906e39e
--- /dev/null
+++ b/src/middleware/errhandlers.rs
@@ -0,0 +1,115 @@
+use std::collections::HashMap;
+
+use error::Result;
+use http::StatusCode;
+use httprequest::HttpRequest;
+use httpresponse::HttpResponse;
+use middleware::{Middleware, Response};
+
+
+type ErrorHandler<S> = Fn(&mut HttpRequest<S>, HttpResponse) -> Result<Response>;
+
+/// `Middleware` for allowing custom handlers for responses.
+///
+/// You can use `ErrorHandlers::handler()` method  to register a custom error handler
+/// for specific status code. You can modify existing response or create completly new
+/// one.
+///
+/// ## Example
+///
+/// ```rust
+/// # extern crate actix_web;
+/// use actix_web::{
+///     App, HttpRequest, HttpResponse, Result,
+///     http, middleware::Response, middleware::ErrorHandlers};
+///
+/// fn render_500<S>(_: &mut HttpRequest<S>, resp: HttpResponse) -> Result<Response> {
+///    let mut builder = resp.into_builder();
+///    builder.header(http::header::CONTENT_TYPE, "application/json");
+///    Ok(Response::Done(builder.into()))
+/// }
+///
+/// fn main() {
+///     let app = App::new()
+///         .middleware(
+///             ErrorHandlers::new()
+///                 .handler(http::StatusCode::INTERNAL_SERVER_ERROR, render_500))
+///         .resource("/test", |r| {
+///              r.method(http::Method::GET).f(|_| HttpResponse::Ok());
+///              r.method(http::Method::HEAD).f(|_| HttpResponse::MethodNotAllowed());
+///         })
+///         .finish();
+/// }
+/// ```
+pub struct ErrorHandlers<S> {
+    handlers: HashMap<StatusCode, Box<ErrorHandler<S>>>,
+}
+
+impl<S> Default for ErrorHandlers<S> {
+    fn default() -> Self {
+        ErrorHandlers {
+            handlers: HashMap::new(),
+        }
+    }
+}
+
+impl<S> ErrorHandlers<S> {
+
+    /// Construct new `ErrorHandlers` instance
+    pub fn new() -> Self {
+        ErrorHandlers::default()
+    }
+
+    /// Register error handler for specified status code
+    pub fn handler<F>(mut self, status: StatusCode, handler: F) -> Self
+        where F: Fn(&mut HttpRequest<S>, HttpResponse) -> Result<Response> + 'static
+    {
+        self.handlers.insert(status, Box::new(handler));
+        self
+    }
+}
+
+impl<S: 'static> Middleware<S> for ErrorHandlers<S> {
+
+    fn response(&self, req: &mut HttpRequest<S>, resp: HttpResponse) -> Result<Response> {
+        if let Some(handler) = self.handlers.get(&resp.status()) {
+            handler(req, resp)
+        } else {
+            Ok(Response::Done(resp))
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use http::StatusCode;
+    use http::header::CONTENT_TYPE;
+
+    fn render_500<S>(_: &mut HttpRequest<S>, resp: HttpResponse) -> Result<Response> {
+        let mut builder = resp.into_builder();
+        builder.header(CONTENT_TYPE, "0001");
+        Ok(Response::Done(builder.into()))
+    }
+    
+    #[test]
+    fn test_handler() {
+        let mw = ErrorHandlers::new()
+            .handler(StatusCode::INTERNAL_SERVER_ERROR, render_500);
+
+        let mut req = HttpRequest::default();
+        let resp = HttpResponse::InternalServerError().finish();
+        let resp = match mw.response(&mut req, resp) {
+            Ok(Response::Done(resp)) => resp,
+            _ => panic!(),
+        };
+        assert_eq!(resp.headers().get(CONTENT_TYPE).unwrap(), "0001");
+
+        let resp = HttpResponse::Ok().finish();
+        let resp = match mw.response(&mut req, resp) {
+            Ok(Response::Done(resp)) => resp,
+            _ => panic!(),
+        };
+        assert!(!resp.headers().contains_key(CONTENT_TYPE));
+    }
+}
diff --git a/src/middleware/mod.rs b/src/middleware/mod.rs
index 6f85dfeb..97f3fa2b 100644
--- a/src/middleware/mod.rs
+++ b/src/middleware/mod.rs
@@ -10,9 +10,11 @@ mod logger;
 #[cfg(feature = "session")]
 mod session;
 mod defaultheaders;
+mod errhandlers;
 pub mod cors;
 pub mod csrf;
 pub use self::logger::Logger;
+pub use self::errhandlers::ErrorHandlers;
 pub use self::defaultheaders::{DefaultHeaders, DefaultHeadersBuilder};
 
 #[cfg(feature = "session")]