From ad7e67f940f0ebd25e0ad2bf1642b65b92df8e18 Mon Sep 17 00:00:00 2001 From: Benny Nazimov <66024037+benny-n@users.noreply.github.com> Date: Mon, 26 Sep 2022 21:44:51 +0300 Subject: [PATCH] add `middleware::logger::custom_response_replace` (#2631) Co-authored-by: Rob Ede --- actix-web/CHANGES.md | 2 + actix-web/src/middleware/logger.rs | 186 +++++++++++++++++++++++++---- 2 files changed, 165 insertions(+), 23 deletions(-) diff --git a/actix-web/CHANGES.md b/actix-web/CHANGES.md index 00ab6c9c8..a018bc248 100644 --- a/actix-web/CHANGES.md +++ b/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 diff --git a/actix-web/src/middleware/logger.rs b/actix-web/src/middleware/logger.rs index 53a3550de..5fec5a013 100644 --- a/actix-web/src/middleware/logger.rs +++ b/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), + CustomResponse(String, Option), } #[derive(Clone)] @@ -494,6 +566,23 @@ impl fmt::Debug for CustomRequestFn { } } +#[derive(Clone)] +struct CustomResponseFn { + inner_fn: Rc 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(&mut self, res: &HttpResponse) { + 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::test::{self, TestRequest}; + use crate::{ + 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";