register fns for custom request-derived logging units (#1749)

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

CHANGES.md

@@ -4,6 +4,7 @@
 ### Added
 * Implement `exclude_regex` for Logger middleware. [#1723]
 * Add request-local data extractor `web::ReqData`. [#1748]
+* Add ability to register closure for request middleware logging. [#1749]
 * Add `app_data` to `ServiceConfig`. [#1757]
 
 ### Changed
@@ -15,6 +16,7 @@
 [#1743]: https://github.com/actix/actix-web/pull/1743
 [#1748]: https://github.com/actix/actix-web/pull/1748
 [#1750]: https://github.com/actix/actix-web/pull/1750
+[#1749]: https://github.com/actix/actix-web/pull/1749
 
 
 ## 3.1.0 - 2020-09-29

src/middleware/logger.rs

@@ -34,21 +34,19 @@ use crate::HttpResponse;
 /// Default `Logger` could be created with `default` method, it uses the
 /// default format:
 ///
-/// ```ignore
+/// ```plain
-///  %a "%r" %s %b "%{Referer}i" "%{User-Agent}i" %T
+/// %a "%r" %s %b "%{Referer}i" "%{User-Agent}i" %T
 /// ```
+///
 /// ```rust
-/// use actix_web::middleware::Logger;
+/// use actix_web::{middleware::Logger, App};
-/// use actix_web::App;
 ///
-/// fn main() {
+/// std::env::set_var("RUST_LOG", "actix_web=info");
-///     std::env::set_var("RUST_LOG", "actix_web=info");
+/// env_logger::init();
-///     env_logger::init();
 ///
-///     let app = App::new()
+/// let app = App::new()
-///         .wrap(Logger::default())
+///     .wrap(Logger::default())
-///         .wrap(Logger::new("%a %{User-Agent}i"));
+///     .wrap(Logger::new("%a %{User-Agent}i"));
-/// }
 /// ```
 ///
 /// ## Format
@@ -80,6 +78,8 @@ use crate::HttpResponse;
 ///
 /// `%{FOO}e`  os.environ['FOO']
 ///
+/// `%{FOO}xi`  [custom request replacement](Logger::custom_request_replace) labelled "FOO"
+///
 /// # Security
 ///  **\*** It is calculated using
 ///  [`ConnectionInfo::realip_remote_addr()`](../dev/struct.ConnectionInfo.html#method.realip_remote_addr)
@@ -123,12 +123,52 @@ impl Logger {
         inner.exclude_regex = regex_set;
         self
     }
+
+    /// Register a function that receives a ServiceRequest 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}xi`.
+    ///
+    /// It is convention to print "-" to indicate no output instead of an empty string.
+    ///
+    /// # Example
+    /// ```rust
+    /// # use actix_web::{http::HeaderValue, middleware::Logger};
+    /// # fn parse_jwt_id (_req: Option<&HeaderValue>) -> String { "jwt_uid".to_owned() }
+    /// Logger::new("example %{JWT_ID}xi")
+    ///     .custom_request_replace("JWT_ID", |req| parse_jwt_id(req.headers().get("Authorization")));
+    /// ```
+    pub fn custom_request_replace(
+        mut self,
+        label: &str,
+        f: impl Fn(&ServiceRequest) -> String + 'static,
+    ) -> Self {
+        let inner = Rc::get_mut(&mut self.0).unwrap();
+
+        let ft = inner.format.0.iter_mut().find(|ft| {
+            matches!(ft, FormatText::CustomRequest(unit_label, _) if label == unit_label)
+        });
+
+        if let Some(FormatText::CustomRequest(_, request_fn)) = ft {
+            // replace into None or previously registered fn using same label
+            request_fn.replace(CustomRequestFn {
+                inner_fn: Rc::new(f),
+            });
+        } else {
+            // non-printed request replacement function diagnostic
+            debug!(
+                "Attempted to register custom request logging function for nonexistent label: {}",
+                label
+            );
+        }
+
+        self
+    }
 }
 
 impl Default for Logger {
     /// Create `Logger` middleware with format:
     ///
-    /// ```ignore
+    /// ```plain
     /// %a "%r" %s %b "%{Referer}i" "%{User-Agent}i" %T
     /// ```
     fn default() -> Logger {
@@ -153,6 +193,17 @@ where
     type Future = Ready<Result<Self::Transform, Self::InitError>>;
 
     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 {
+                debug!(
+                    "No custom request replacement function was registered for label {} in\
+                    logger format.",
+                    label
+                );
+            }
+        }
+
         ok(LoggerMiddleware {
             service,
             inner: self.0.clone(),
@@ -311,7 +362,6 @@ impl<B: MessageBody> MessageBody for StreamLog<B> {
 /// A formatting style for the `Logger`, consisting of multiple
 /// `FormatText`s concatenated into one line.
 #[derive(Clone)]
-#[doc(hidden)]
 struct Format(Vec<FormatText>);
 
 impl Default for Format {
@@ -327,7 +377,8 @@ 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])|[atPrUsbTD]?)").unwrap();
+        let fmt =
+            Regex::new(r"%(\{([A-Za-z0-9\-_]+)\}([aioe]|xi)|[atPrUsbTD]?)").unwrap();
 
         let mut idx = 0;
         let mut results = Vec::new();
@@ -355,6 +406,7 @@ impl Format {
                         HeaderName::try_from(key.as_str()).unwrap(),
                     ),
                     "e" => FormatText::EnvironHeader(key.as_str().to_owned()),
+                    "xi" => FormatText::CustomRequest(key.as_str().to_owned(), None),
                     _ => unreachable!(),
                 })
             } else {
@@ -384,7 +436,9 @@ impl Format {
 /// A string of text to be logged. This is either one of the data
 /// fields supported by the `Logger`, or a custom `String`.
 #[doc(hidden)]
+#[non_exhaustive]
 #[derive(Debug, Clone)]
+// TODO: remove pub on next breaking change
 pub enum FormatText {
     Str(String),
     Percent,
@@ -400,6 +454,26 @@ pub enum FormatText {
     RequestHeader(HeaderName),
     ResponseHeader(HeaderName),
     EnvironHeader(String),
+    CustomRequest(String, Option<CustomRequestFn>),
+}
+
+// TODO: remove pub on next breaking change
+#[doc(hidden)]
+#[derive(Clone)]
+pub struct CustomRequestFn {
+    inner_fn: Rc<dyn Fn(&ServiceRequest) -> String>,
+}
+
+impl CustomRequestFn {
+    fn call(&self, req: &ServiceRequest) -> String {
+        (self.inner_fn)(req)
+    }
+}
+
+impl fmt::Debug for CustomRequestFn {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str("custom_request_fn")
+    }
 }
 
 impl FormatText {
@@ -456,7 +530,7 @@ impl FormatText {
     }
 
     fn render_request(&mut self, now: OffsetDateTime, req: &ServiceRequest) {
-        match *self {
+        match &*self {
             FormatText::RequestLine => {
                 *self = if req.query_string().is_empty() {
                     FormatText::Str(format!(
@@ -508,11 +582,20 @@ impl FormatText {
                 };
                 *self = s;
             }
+            FormatText::CustomRequest(_, request_fn) => {
+                let s = match request_fn {
+                    Some(f) => FormatText::Str(f.call(req)),
+                    None => FormatText::Str("-".to_owned()),
+                };
+
+                *self = s;
+            }
             _ => (),
         }
     }
 }
 
+/// Converter to get a String from something that writes to a Formatter.
 pub(crate) struct FormatDisplay<'a>(
     &'a dyn Fn(&mut Formatter<'_>) -> Result<(), fmt::Error>,
 );
@@ -530,7 +613,7 @@ mod tests {
 
     use super::*;
     use crate::http::{header, StatusCode};
-    use crate::test::TestRequest;
+    use crate::test::{self, TestRequest};
 
     #[actix_rt::test]
     async fn test_logger() {
@@ -699,4 +782,45 @@ mod tests {
         println!("{}", s);
         assert!(s.contains("192.0.2.60"));
     }
+
+    #[actix_rt::test]
+    async fn test_custom_closure_log() {
+        let mut logger = Logger::new("test %{CUSTOM}xi")
+            .custom_request_replace("CUSTOM", |_req: &ServiceRequest| -> String {
+                String::from("custom_log")
+            });
+        let mut unit = Rc::get_mut(&mut logger.0).unwrap().format.0[1].clone();
+
+        let label = match &unit {
+            FormatText::CustomRequest(label, _) => label,
+            ft => panic!("expected CustomRequest, found {:?}", ft),
+        };
+
+        assert_eq!(label, "CUSTOM");
+
+        let req = TestRequest::default().to_srv_request();
+        let now = OffsetDateTime::now_utc();
+
+        unit.render_request(now, &req);
+
+        let render = |fmt: &mut 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";
+
+        let logger = Logger::new("%{CUSTOM}xi")
+            .custom_request_replace("CUSTOM", move |_req: &ServiceRequest| -> String {
+                captured.to_owned()
+            });
+
+        let mut srv = logger.new_transform(test::ok_service()).await.unwrap();
+
+        let req = TestRequest::default().to_srv_request();
+        srv.call(req).await.unwrap();
+    }
 }