diff --git a/actix-web/src/middleware/err_handlers.rs b/actix-web/src/middleware/err_handlers.rs index f74220cd2..657b4786c 100644 --- a/actix-web/src/middleware/err_handlers.rs +++ b/actix-web/src/middleware/err_handlers.rs @@ -30,6 +30,8 @@ pub enum ErrorHandlerResponse { type ErrorHandler = dyn Fn(ServiceResponse) -> Result>; +type DefaultHandler = Option>>; + /// Middleware for registering custom status code based error handlers. /// /// Register handlers with the `ErrorHandlers::handler()` method to register a custom error handler @@ -54,6 +56,8 @@ type ErrorHandler = dyn Fn(ServiceResponse) -> Result { + default_client: DefaultHandler, + default_server: DefaultHandler, handlers: Handlers, } @@ -62,6 +66,8 @@ type Handlers = Rc>>>; impl Default for ErrorHandlers { fn default() -> Self { ErrorHandlers { + default_client: Default::default(), + default_server: Default::default(), handlers: Default::default(), } } @@ -83,6 +89,46 @@ impl ErrorHandlers { .insert(status, Box::new(handler)); self } + + pub fn default_handler(self, handler: F) -> Self + where + F: Fn(ServiceResponse) -> Result> + 'static, + { + let handler = Rc::new(handler); + Self { + default_server: Some(handler.clone()), + default_client: Some(handler), + ..self + } + } + + pub fn default_handler_server(self, handler: F) -> Self + where + F: Fn(ServiceResponse) -> Result> + 'static, + { + Self { default_server: Some(Rc::new(handler)), ..self } + } + + pub fn default_handler_client(self, handler: F) -> Self + where + F: Fn(ServiceResponse) -> Result> + 'static, + { + Self { default_client: Some(Rc::new(handler)), ..self } + } + + fn get_handler<'a>( + status: &StatusCode, + default_client: Option<&'a dyn Fn(ServiceResponse) -> Result>>, + default_server: Option<&'a dyn Fn(ServiceResponse) -> Result>>, + handlers: &'a Handlers, + ) -> Option<&'a dyn Fn(ServiceResponse) -> Result>> + { + handlers + .get(status) + .map(|h| h.as_ref()) + .or_else(|| status.is_client_error().then(|| default_client).flatten()) + .or_else(|| status.is_server_error().then(|| default_server).flatten()) + } } impl Transform for ErrorHandlers @@ -99,13 +145,22 @@ where fn new_transform(&self, service: S) -> Self::Future { let handlers = self.handlers.clone(); - Box::pin(async move { Ok(ErrorHandlersMiddleware { service, handlers }) }) + let default_client = self.default_client.clone(); + let default_server = self.default_server.clone(); + Box::pin(async move { Ok(ErrorHandlersMiddleware { + service, + default_client, + default_server, + handlers + }) }) } } #[doc(hidden)] pub struct ErrorHandlersMiddleware { service: S, + default_client: DefaultHandler, + default_server: DefaultHandler, handlers: Handlers, } @@ -123,8 +178,10 @@ where fn call(&self, req: ServiceRequest) -> Self::Future { let handlers = self.handlers.clone(); + let default_client = self.default_client.clone(); + let default_server = self.default_server.clone(); let fut = self.service.call(req); - ErrorHandlersFuture::ServiceFuture { fut, handlers } + ErrorHandlersFuture::ServiceFuture { fut, default_client, default_server, handlers } } } @@ -137,6 +194,8 @@ pin_project! { ServiceFuture { #[pin] fut: Fut, + default_client: DefaultHandler, + default_server: DefaultHandler, handlers: Handlers, }, ErrorHandlerFuture { @@ -153,10 +212,17 @@ where fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { match self.as_mut().project() { - ErrorHandlersProj::ServiceFuture { fut, handlers } => { + ErrorHandlersProj::ServiceFuture { fut, default_client, default_server, handlers } => { let res = ready!(fut.poll(cx))?; + let status = res.status(); - match handlers.get(&res.status()) { + let handler = ErrorHandlers::get_handler( + &status, + default_client.as_mut().map(|f| Rc::as_ref(f)), + default_server.as_mut().map(|f| Rc::as_ref(f)), + handlers + ); + match handler { Some(handler) => match handler(res)? { ErrorHandlerResponse::Response(res) => Poll::Ready(Ok(res)), ErrorHandlerResponse::Future(fut) => { @@ -165,8 +231,7 @@ where self.poll(cx) } - }, - + } None => Poll::Ready(Ok(res.map_into_left_body())), } } @@ -298,4 +363,103 @@ mod tests { "error in error handler" ); } + + #[actix_rt::test] + async fn default_error_handler() { + #[allow(clippy::unnecessary_wraps)] + fn error_handler(mut res: ServiceResponse) -> Result> { + res.response_mut() + .headers_mut() + .insert(CONTENT_TYPE, HeaderValue::from_static("0001")); + Ok(ErrorHandlerResponse::Response(res.map_into_left_body())) + } + + let make_mw = |status| async move { + ErrorHandlers::new() + .default_handler(error_handler) + .new_transform(test::status_service(status).into_service()) + .await + .unwrap() + }; + let mw_server = make_mw(StatusCode::INTERNAL_SERVER_ERROR).await; + let mw_client = make_mw(StatusCode::BAD_REQUEST).await; + + let resp = test::call_service(&mw_client, TestRequest::default().to_srv_request()).await; + assert_eq!(resp.headers().get(CONTENT_TYPE).unwrap(), "0001"); + + let resp = test::call_service(&mw_server, TestRequest::default().to_srv_request()).await; + assert_eq!(resp.headers().get(CONTENT_TYPE).unwrap(), "0001"); + } + + #[actix_rt::test] + async fn default_handlers_separate_client_server() { + #[allow(clippy::unnecessary_wraps)] + fn error_handler_client(mut res: ServiceResponse) -> Result> { + res.response_mut() + .headers_mut() + .insert(CONTENT_TYPE, HeaderValue::from_static("0001")); + Ok(ErrorHandlerResponse::Response(res.map_into_left_body())) + } + + #[allow(clippy::unnecessary_wraps)] + fn error_handler_server(mut res: ServiceResponse) -> Result> { + res.response_mut() + .headers_mut() + .insert(CONTENT_TYPE, HeaderValue::from_static("0002")); + Ok(ErrorHandlerResponse::Response(res.map_into_left_body())) + } + + let make_mw = |status| async move { + ErrorHandlers::new() + .default_handler_server(error_handler_server) + .default_handler_client(error_handler_client) + .new_transform(test::status_service(status).into_service()) + .await + .unwrap() + }; + let mw_server = make_mw(StatusCode::INTERNAL_SERVER_ERROR).await; + let mw_client = make_mw(StatusCode::BAD_REQUEST).await; + + let resp = test::call_service(&mw_client, TestRequest::default().to_srv_request()).await; + assert_eq!(resp.headers().get(CONTENT_TYPE).unwrap(), "0001"); + + let resp = test::call_service(&mw_server, TestRequest::default().to_srv_request()).await; + assert_eq!(resp.headers().get(CONTENT_TYPE).unwrap(), "0002"); + } + + #[actix_rt::test] + async fn default_handlers_specialization() { + #[allow(clippy::unnecessary_wraps)] + fn error_handler_client(mut res: ServiceResponse) -> Result> { + res.response_mut() + .headers_mut() + .insert(CONTENT_TYPE, HeaderValue::from_static("0001")); + Ok(ErrorHandlerResponse::Response(res.map_into_left_body())) + } + + #[allow(clippy::unnecessary_wraps)] + fn error_handler_specific(mut res: ServiceResponse) -> Result> { + res.response_mut() + .headers_mut() + .insert(CONTENT_TYPE, HeaderValue::from_static("0003")); + Ok(ErrorHandlerResponse::Response(res.map_into_left_body())) + } + + let make_mw = |status| async move { + ErrorHandlers::new() + .default_handler_client(error_handler_client) + .handler(StatusCode::UNPROCESSABLE_ENTITY, error_handler_specific) + .new_transform(test::status_service(status).into_service()) + .await + .unwrap() + }; + let mw_client = make_mw(StatusCode::BAD_REQUEST).await; + let mw_specific = make_mw(StatusCode::UNPROCESSABLE_ENTITY).await; + + let resp = test::call_service(&mw_client, TestRequest::default().to_srv_request()).await; + assert_eq!(resp.headers().get(CONTENT_TYPE).unwrap(), "0001"); + + let resp = test::call_service(&mw_specific, TestRequest::default().to_srv_request()).await; + assert_eq!(resp.headers().get(CONTENT_TYPE).unwrap(), "0003"); + } }