add middleware::logger::custom_response_replace (#2631)

Co-authored-by: Rob Ede <robjtede@icloud.com>

actix-web/CHANGES.md

@@ -4,7 +4,9 @@
 ### Added
 - Add `ContentDisposition::attachment` constructor. [#2867]
 - Add `ErrorHandlers::default_handler()` (as well as `default_handler_{server, client}()`) to make registering handlers for groups of response statuses easier. [#2784]
+- Add `Logger::custom_response_replace()`. [#2631]
 
+[#2631]: https://github.com/actix/actix-web/pull/2631
 [#2784]: https://github.com/actix/actix-web/pull/2784
 [#2867]: https://github.com/actix/actix-web/pull/2867
 

actix-web/src/middleware/logger.rs

@@ -26,7 +26,7 @@ use crate::{
     body::{BodySize, MessageBody},
     http::header::HeaderName,
     service::{ServiceRequest, ServiceResponse},
-    Error, HttpResponse, Result,
+    Error, Result,
 };
 
 /// Middleware for logging request and response summaries to the terminal.
@@ -69,10 +69,11 @@ use crate::{
 /// `%D` | Time taken to serve the request, in milliseconds
 /// `%U` | Request URL
 /// `%{r}a` | "Real IP" remote address **\***
-/// `%{FOO}i` |  `request.headers["FOO"]`
+/// `%{FOO}i` | `request.headers["FOO"]`
 /// `%{FOO}o` | `response.headers["FOO"]`
 /// `%{FOO}e` | `env_var["FOO"]`
 /// `%{FOO}xi` | [Custom request replacement](Logger::custom_request_replace) labelled "FOO"
+/// `%{FOO}xo` | [Custom response replacement](Logger::custom_response_replace) labelled "FOO"
 ///
 /// # Security
 /// **\*** "Real IP" remote address is calculated using
@@ -179,6 +180,55 @@ impl Logger {
 
         self
     }
+
+    /// Register a function that receives a `ServiceResponse` and returns a string for use in the
+    /// log line.
+    ///
+    /// The label passed as the first argument should match a replacement substring in
+    /// the logger format like `%{label}xo`.
+    ///
+    /// It is convention to print "-" to indicate no output instead of an empty string.
+    ///
+    /// The replacement function does not have access to the response body.
+    ///
+    /// # Examples
+    /// ```
+    /// # use actix_web::{dev::ServiceResponse, middleware::Logger};
+    /// fn log_if_error(res: &ServiceResponse) -> String {
+    ///     if res.status().as_u16() >= 400 {
+    ///         "ERROR".to_string()
+    ///     } else {
+    ///         "-".to_string()
+    ///     }
+    /// }
+    ///
+    /// Logger::new("example %{ERROR_STATUS}xo")
+    ///     .custom_response_replace("ERROR_STATUS", |res| log_if_error(res) );
+    /// ```
+    pub fn custom_response_replace(
+        mut self,
+        label: &str,
+        f: impl Fn(&ServiceResponse) -> String + 'static,
+    ) -> Self {
+        let inner = Rc::get_mut(&mut self.0).unwrap();
+
+        let ft = inner.format.0.iter_mut().find(
+            |ft| matches!(ft, FormatText::CustomResponse(unit_label, _) if label == unit_label),
+        );
+
+        if let Some(FormatText::CustomResponse(_, res_fn)) = ft {
+            *res_fn = Some(CustomResponseFn {
+                inner_fn: Rc::new(f),
+            });
+        } else {
+            debug!(
+                "Attempted to register custom response logging function for non-existent label: {}",
+                label
+            );
+        }
+
+        self
+    }
 }
 
 impl Default for Logger {
@@ -210,10 +260,16 @@ where
 
     fn new_transform(&self, service: S) -> Self::Future {
         for unit in &self.0.format.0 {
-            // missing request replacement function diagnostic
             if let FormatText::CustomRequest(label, None) = unit {
                 warn!(
-                    "No custom request replacement function was registered for label \"{}\".",
+                    "No custom request replacement function was registered for label: {}",
+                    label
+                );
+            }
+
+            if let FormatText::CustomResponse(label, None) = unit {
+                warn!(
+                    "No custom response replacement function was registered for label: {}",
                     label
                 );
             }
@@ -308,11 +364,25 @@ where
             debug!("Error in response: {:?}", error);
         }
 
-        if let Some(ref mut format) = this.format {
+        let res = if let Some(ref mut format) = this.format {
+            // to avoid polluting all the Logger types with the body parameter we swap the body
+            // out temporarily since it's not usable in custom response functions anyway
+
+            let (req, res) = res.into_parts();
+            let (res, body) = res.into_parts();
+
+            let temp_res = ServiceResponse::new(req, res.map_into_boxed_body());
+
             for unit in &mut format.0 {
-                unit.render_response(res.response());
+                unit.render_response(&temp_res);
             }
-        }
+
+            // re-construct original service response
+            let (req, res) = temp_res.into_parts();
+            ServiceResponse::new(req, res.set_body(body))
+        } else {
+            res
+        };
 
         let time = *this.time;
         let format = this.format.take();
@@ -399,7 +469,7 @@ impl Format {
     /// Returns `None` if the format string syntax is incorrect.
     pub fn new(s: &str) -> Format {
         log::trace!("Access log format: {}", s);
-        let fmt = Regex::new(r"%(\{([A-Za-z0-9\-_]+)\}([aioe]|xi)|[%atPrUsbTD]?)").unwrap();
+        let fmt = Regex::new(r"%(\{([A-Za-z0-9\-_]+)\}([aioe]|x[io])|[%atPrUsbTD]?)").unwrap();
 
         let mut idx = 0;
         let mut results = Vec::new();
@@ -417,7 +487,7 @@ impl Format {
                         if key.as_str() == "r" {
                             FormatText::RealIpRemoteAddr
                         } else {
-                            unreachable!()
+                            unreachable!("regex and code mismatch")
                         }
                     }
                     "i" => {
@@ -428,6 +498,7 @@ impl Format {
                     }
                     "e" => FormatText::EnvironHeader(key.as_str().to_owned()),
                     "xi" => FormatText::CustomRequest(key.as_str().to_owned(), None),
+                    "xo" => FormatText::CustomResponse(key.as_str().to_owned(), None),
                     _ => unreachable!(),
                 })
             } else {
@@ -475,6 +546,7 @@ enum FormatText {
     ResponseHeader(HeaderName),
     EnvironHeader(String),
     CustomRequest(String, Option<CustomRequestFn>),
+    CustomResponse(String, Option<CustomResponseFn>),
 }
 
 #[derive(Clone)]
@@ -494,6 +566,23 @@ impl fmt::Debug for CustomRequestFn {
     }
 }
 
+#[derive(Clone)]
+struct CustomResponseFn {
+    inner_fn: Rc<dyn Fn(&ServiceResponse) -> String>,
+}
+
+impl CustomResponseFn {
+    fn call(&self, res: &ServiceResponse) -> String {
+        (self.inner_fn)(res)
+    }
+}
+
+impl fmt::Debug for CustomResponseFn {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str("custom_response_fn")
+    }
+}
+
 impl FormatText {
     fn render(
         &self,
@@ -526,11 +615,12 @@ impl FormatText {
         }
     }
 
-    fn render_response<B>(&mut self, res: &HttpResponse<B>) {
+    fn render_response(&mut self, res: &ServiceResponse) {
         match self {
             FormatText::ResponseStatus => {
                 *self = FormatText::Str(format!("{}", res.status().as_u16()))
             }
+
             FormatText::ResponseHeader(ref name) => {
                 let s = if let Some(val) = res.headers().get(name) {
                     if let Ok(s) = val.to_str() {
@@ -543,6 +633,16 @@ impl FormatText {
                 };
                 *self = FormatText::Str(s.to_string())
             }
+
+            FormatText::CustomResponse(_, res_fn) => {
+                let text = match res_fn {
+                    Some(res_fn) => FormatText::Str(res_fn.call(res)),
+                    None => FormatText::Str("-".to_owned()),
+                };
+
+                *self = text;
+            }
+
             _ => {}
         }
     }
@@ -627,8 +727,11 @@ mod tests {
     use actix_utils::future::ok;
 
     use super::*;
-    use crate::http::{header, StatusCode};
+    use crate::{
-    use crate::test::{self, TestRequest};
+        http::{header, StatusCode},
+        test::{self, TestRequest},
+        HttpResponse,
+    };
 
     #[actix_rt::test]
     async fn test_logger() {
@@ -691,9 +794,10 @@ mod tests {
             unit.render_request(now, &req);
         }
 
-        let resp = HttpResponse::build(StatusCode::OK).force_close().finish();
+        let req = TestRequest::default().to_http_request();
+        let res = ServiceResponse::new(req, HttpResponse::Ok().finish());
         for unit in &mut format.0 {
-            unit.render_response(&resp);
+            unit.render_response(&res);
         }
 
         let entry_time = OffsetDateTime::now_utc();
@@ -723,9 +827,10 @@ mod tests {
             unit.render_request(now, &req);
         }
 
-        let resp = HttpResponse::build(StatusCode::OK).force_close().finish();
+        let req = TestRequest::default().to_http_request();
+        let res = ServiceResponse::new(req, HttpResponse::Ok().force_close().finish());
         for unit in &mut format.0 {
-            unit.render_response(&resp);
+            unit.render_response(&res);
         }
 
         let render = |fmt: &mut fmt::Formatter<'_>| {
@@ -755,9 +860,10 @@ mod tests {
             unit.render_request(now, &req);
         }
 
-        let resp = HttpResponse::build(StatusCode::OK).force_close().finish();
+        let req = TestRequest::default().to_http_request();
+        let res = ServiceResponse::new(req, HttpResponse::Ok().force_close().finish());
         for unit in &mut format.0 {
-            unit.render_response(&resp);
+            unit.render_response(&res);
         }
 
         let entry_time = OffsetDateTime::now_utc();
@@ -784,9 +890,10 @@ mod tests {
             unit.render_request(now, &req);
         }
 
-        let resp = HttpResponse::build(StatusCode::OK).force_close().finish();
+        let req = TestRequest::default().to_http_request();
+        let res = ServiceResponse::new(req, HttpResponse::Ok().force_close().finish());
         for unit in &mut format.0 {
-            unit.render_response(&resp);
+            unit.render_response(&res);
         }
 
         let render = |fmt: &mut fmt::Formatter<'_>| {
@@ -815,9 +922,10 @@ mod tests {
             unit.render_request(now, &req);
         }
 
-        let resp = HttpResponse::build(StatusCode::OK).force_close().finish();
+        let req = TestRequest::default().to_http_request();
+        let res = ServiceResponse::new(req, HttpResponse::Ok().finish());
         for unit in &mut format.0 {
-            unit.render_response(&resp);
+            unit.render_response(&res);
         }
 
         let entry_time = OffsetDateTime::now_utc();
@@ -832,7 +940,7 @@ mod tests {
     }
 
     #[actix_rt::test]
-    async fn test_custom_closure_log() {
+    async fn test_custom_closure_req_log() {
         let mut logger = Logger::new("test %{CUSTOM}xi")
             .custom_request_replace("CUSTOM", |_req: &ServiceRequest| -> String {
                 String::from("custom_log")
@@ -857,6 +965,38 @@ mod tests {
         assert_eq!(log_output, "custom_log");
     }
 
+    #[actix_rt::test]
+    async fn test_custom_closure_response_log() {
+        let mut logger = Logger::new("test %{CUSTOM}xo").custom_response_replace(
+            "CUSTOM",
+            |res: &ServiceResponse| -> String {
+                if res.status().as_u16() == 200 {
+                    String::from("custom_log")
+                } else {
+                    String::from("-")
+                }
+            },
+        );
+        let mut unit = Rc::get_mut(&mut logger.0).unwrap().format.0[1].clone();
+
+        let label = match &unit {
+            FormatText::CustomResponse(label, _) => label,
+            ft => panic!("expected CustomResponse, found {:?}", ft),
+        };
+
+        assert_eq!(label, "CUSTOM");
+
+        let req = TestRequest::default().to_http_request();
+        let resp_ok = ServiceResponse::new(req, HttpResponse::Ok().finish());
+        let now = OffsetDateTime::now_utc();
+        unit.render_response(&resp_ok);
+
+        let render = |fmt: &mut fmt::Formatter<'_>| unit.render(fmt, 1024, now);
+
+        let log_output = FormatDisplay(&render).to_string();
+        assert_eq!(log_output, "custom_log");
+    }
+
     #[actix_rt::test]
     async fn test_closure_logger_in_middleware() {
         let captured = "custom log replacement";