From 46cfe23d34b284a7240b40278aed2ae180d08dda Mon Sep 17 00:00:00 2001 From: benny-n Date: Sat, 5 Feb 2022 11:55:53 +0200 Subject: [PATCH] feat: add custom response replace to middleware logger Initial implementation, only exposes the response's status --- actix-web/src/middleware/logger.rs | 119 ++++++++++++++++++++++++++++- 1 file changed, 116 insertions(+), 3 deletions(-) diff --git a/actix-web/src/middleware/logger.rs b/actix-web/src/middleware/logger.rs index 53a3550de..067f2aba9 100644 --- a/actix-web/src/middleware/logger.rs +++ b/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), + CustomResponse(String, Option), } #[derive(Clone)] @@ -494,6 +551,23 @@ impl fmt::Debug for CustomRequestFn { } } +#[derive(Clone)] +struct CustomResponseFn { + inner_fn: Rc 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";