feat: add custom response replace to middleware logger

Initial implementation, only exposes the response's status

actix-web/src/middleware/logger.rs

@@ -24,7 +24,7 @@ use time::{format_description::well_known::Rfc3339, OffsetDateTime};
 
 use crate::{
     body::{BodySize, MessageBody},
-    http::header::HeaderName,
+    http::{header::HeaderName, StatusCode},
     service::{ServiceRequest, ServiceResponse},
     Error, HttpResponse, Result,
 };
@@ -73,6 +73,7 @@ use crate::{
 /// `%{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,53 @@ impl Logger {
 
         self
     }
+
+    /// Register a function that receives a StatusCode 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.
+    ///
+    /// # Examples
+    /// ```
+    /// # use actix_web::http::{header::HeaderValue, StatusCode};
+    /// # use actix_web::middleware::Logger;
+    /// fn log_if_error (status: &StatusCode) -> String {
+    ///     if status.as_u16() >= 400 {
+    ///         "ERROR".to_string()
+    ///     } else {
+    ///         "-".to_string()
+    ///     }
+    /// }
+    /// Logger::new("example %{ERROR_STATUS}xo")
+    ///     .custom_response_replace("ERROR_STATUS", |status| log_if_error(status) );
+    /// ```
+    pub fn custom_response_replace(
+        mut self,
+        label: &str,
+        f: impl Fn(&StatusCode) -> 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(_, response_fn)) = ft {
+            // replace into None or previously registered fn using same label
+            response_fn.replace(CustomResponseFn {
+                inner_fn: Rc::new(f),
+            });
+        } else {
+            // non-printed response replacement function diagnostic
+            debug!(
+                "Attempted to register custom response logging function for nonexistent label: {}",
+                label
+            );
+        }
+
+        self
+    }
 }
 
 impl Default for Logger {
@@ -217,6 +265,13 @@ where
                     label
                 );
             }
+            // missing response replacement function diagnostic
+            if let FormatText::CustomResponse(label, None) = unit {
+                warn!(
+                    "No custom response replacement function was registered for label \"{}\".",
+                    label
+                );
+            }
         }
 
         ready(Ok(LoggerMiddleware {
@@ -399,7 +454,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();
@@ -428,6 +483,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 +531,7 @@ enum FormatText {
     ResponseHeader(HeaderName),
     EnvironHeader(String),
     CustomRequest(String, Option<CustomRequestFn>),
+    CustomResponse(String, Option<CustomResponseFn>),
 }
 
 #[derive(Clone)]
@@ -494,6 +551,23 @@ impl fmt::Debug for CustomRequestFn {
     }
 }
 
+#[derive(Clone)]
+struct CustomResponseFn {
+    inner_fn: Rc<dyn Fn(&StatusCode) -> String>,
+}
+
+impl CustomResponseFn {
+    fn call(&self, res: &StatusCode) -> 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,
@@ -543,6 +617,14 @@ impl FormatText {
                 };
                 *self = FormatText::Str(s.to_string())
             }
+            FormatText::CustomResponse(_, response_fn) => {
+                let s = match response_fn {
+                    Some(f) => FormatText::Str(f.call(&res.status())),
+                    None => FormatText::Str("-".to_owned()),
+                };
+
+                *self = s;
+            }
             _ => {}
         }
     }
@@ -832,7 +914,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 +939,37 @@ mod tests {
         assert_eq!(log_output, "custom_log");
     }
 
+    #[actix_rt::test]
+    async fn test_custom_closure_res_log() {
+        let mut logger = Logger::new("test %{CUSTOM}xo").custom_response_replace(
+            "CUSTOM",
+            |sc: &StatusCode| -> String {
+                if sc.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 resp_ok = HttpResponse::build(StatusCode::OK).force_close().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";