From 5140fea8d1424c9887cd85d945be1e9c884ffb9e Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Fri, 13 Apr 2018 19:10:42 -0700
Subject: [PATCH] allow to use castom error handler for json extractor

---
 src/error.rs       | 37 ++++++++++++++++++++++++++++++++++---
 src/httpmessage.rs |  5 +++--
 src/json.rs        | 36 +++++++++++++++++++++++++++---------
 3 files changed, 64 insertions(+), 14 deletions(-)

diff --git a/src/error.rs b/src/error.rs
index 7435b504b..796183fc0 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -1,4 +1,5 @@
 //! Error and Result module
+use std::cell::RefCell;
 use std::io::Error as IoError;
 use std::str::Utf8Error;
 use std::string::FromUtf8Error;
@@ -545,18 +546,31 @@ impl From<UrlParseError> for UrlGenerationError {
 /// ```
 pub struct InternalError<T> {
     cause: T,
-    status: StatusCode,
+    status: InternalErrorType,
     backtrace: Backtrace,
 }
 
 unsafe impl<T> Sync for InternalError<T> {}
 unsafe impl<T> Send for InternalError<T> {}
 
+enum InternalErrorType {
+    Status(StatusCode),
+    Response(RefCell<Option<HttpResponse>>),
+}
+
 impl<T> InternalError<T> {
     pub fn new(cause: T, status: StatusCode) -> Self {
         InternalError {
             cause,
-            status,
+            status: InternalErrorType::Status(status),
+            backtrace: Backtrace::new(),
+        }
+    }
+
+    pub fn from_response(cause: T, response: HttpResponse) -> Self {
+        InternalError {
+            cause,
+            status: InternalErrorType::Response(RefCell::new(Some(response))),
             backtrace: Backtrace::new(),
         }
     }
@@ -594,7 +608,16 @@ where
     T: Send + Sync + fmt::Debug + 'static,
 {
     fn error_response(&self) -> HttpResponse {
-        HttpResponse::new(self.status)
+        match self.status {
+            InternalErrorType::Status(st) => HttpResponse::new(st),
+            InternalErrorType::Response(ref resp) => {
+                if let Some(resp) = resp.borrow_mut().take() {
+                    resp
+                } else {
+                    HttpResponse::new(StatusCode::INTERNAL_SERVER_ERROR)
+                }
+            }
+        }
     }
 }
 
@@ -859,4 +882,12 @@ mod tests {
             _ => env::remove_var(NAME),
         }
     }
+
+    #[test]
+    fn test_internal_error() {
+        let err = InternalError::from_response(
+            ExpectError::Encoding, HttpResponse::Ok().into());
+        let resp: HttpResponse = err.error_response();
+        assert_eq!(resp.status(), StatusCode::OK);
+    }
 }
diff --git a/src/httpmessage.rs b/src/httpmessage.rs
index 0b40a8128..b590172b9 100644
--- a/src/httpmessage.rs
+++ b/src/httpmessage.rs
@@ -118,11 +118,12 @@ pub trait HttpMessage {
     /// # extern crate actix_web;
     /// # extern crate futures;
     /// # #[macro_use] extern crate serde_derive;
-    /// use actix_web::*;
     /// use bytes::Bytes;
     /// use futures::future::Future;
+    /// use actix_web::{HttpMessage, HttpRequest, HttpResponse,
+    ///                 FutureResponse, AsyncResponder};
     ///
-    /// fn index(mut req: HttpRequest) -> Box<Future<Item=HttpResponse, Error=Error>> {
+    /// fn index(mut req: HttpRequest) -> FutureResponse<HttpResponse> {
     ///     req.body()                     // <- get Body future
     ///        .limit(1024)                // <- change max size of the body to a 1kb
     ///        .from_err()
diff --git a/src/json.rs b/src/json.rs
index 4eb51fb89..73128b975 100644
--- a/src/json.rs
+++ b/src/json.rs
@@ -2,6 +2,7 @@ use bytes::{Bytes, BytesMut};
 use futures::{Future, Poll, Stream};
 use http::header::CONTENT_LENGTH;
 use std::fmt;
+use std::rc::Rc;
 use std::ops::{Deref, DerefMut};
 
 use mime;
@@ -131,15 +132,17 @@ where
     T: DeserializeOwned + 'static,
     S: 'static,
 {
-    type Config = JsonConfig;
+    type Config = JsonConfig<S>;
     type Result = Box<Future<Item = Self, Error = Error>>;
 
     #[inline]
     fn from_request(req: &HttpRequest<S>, cfg: &Self::Config) -> Self::Result {
+        let req = req.clone();
+        let err = Rc::clone(&cfg.ehandler);
         Box::new(
             JsonBody::new(req.clone())
                 .limit(cfg.limit)
-                .from_err()
+                .map_err(move |e| (*err)(e, req))
                 .map(Json),
         )
     }
@@ -150,7 +153,7 @@ where
 /// ```rust
 /// # extern crate actix_web;
 /// #[macro_use] extern crate serde_derive;
-/// use actix_web::{App, Json, Result, http};
+/// use actix_web::{App, Json, HttpResponse, Result, http, error};
 ///
 /// #[derive(Deserialize)]
 /// struct Info {
@@ -167,25 +170,40 @@ where
 ///        "/index.html", |r| {
 ///            r.method(http::Method::POST)
 ///               .with(index)
-///               .limit(4096);} // <- change json extractor configuration
-///     );
+///               .limit(4096)   // <- change json extractor configuration
+///               .error_handler(|err, req| {  // <- create custom error response
+///                   error::InternalError::from_response(
+///                      err, HttpResponse::Conflict().finish()).into()
+///               });
+///        });
 /// }
 /// ```
-pub struct JsonConfig {
+pub struct JsonConfig<S> {
     limit: usize,
+    ehandler: Rc<Fn(JsonPayloadError, HttpRequest<S>) -> Error>,
 }
 
-impl JsonConfig {
+impl<S> JsonConfig<S> {
     /// Change max size of payload. By default max size is 256Kb
     pub fn limit(&mut self, limit: usize) -> &mut Self {
         self.limit = limit;
         self
     }
+
+    /// Set custom error handler
+    pub fn error_handler<F>(&mut self, f: F) -> &mut Self
+    where
+        F: Fn(JsonPayloadError, HttpRequest<S>) -> Error + 'static
+    {
+        self.ehandler = Rc::new(f);
+        self
+    }
 }
 
-impl Default for JsonConfig {
+impl<S> Default for JsonConfig<S> {
     fn default() -> Self {
-        JsonConfig { limit: 262_144 }
+        JsonConfig { limit: 262_144,
+                     ehandler: Rc::new(|e, _| e.into()) }
     }
 }