Add feature for limit request to log in middleware logger using status code

This commit is contained in:
Tglman 2023-07-10 23:59:01 +01:00
parent ede0201aa4
commit ca51ede874
1 changed files with 37 additions and 13 deletions

View File

@ -7,6 +7,7 @@ use std::{
fmt::{self, Display as _}, fmt::{self, Display as _},
future::Future, future::Future,
marker::PhantomData, marker::PhantomData,
ops::{Bound, RangeBounds},
pin::Pin, pin::Pin,
rc::Rc, rc::Rc,
task::{Context, Poll}, task::{Context, Poll},
@ -23,7 +24,7 @@ use time::{format_description::well_known::Rfc3339, OffsetDateTime};
use crate::{ use crate::{
body::{BodySize, MessageBody}, body::{BodySize, MessageBody},
http::header::HeaderName, http::{header::HeaderName, StatusCode},
service::{ServiceRequest, ServiceResponse}, service::{ServiceRequest, ServiceResponse},
Error, Result, Error, Result,
}; };
@ -89,6 +90,7 @@ struct Inner {
exclude: HashSet<String>, exclude: HashSet<String>,
exclude_regex: RegexSet, exclude_regex: RegexSet,
log_target: Cow<'static, str>, log_target: Cow<'static, str>,
status_range: (Bound<StatusCode>, Bound<StatusCode>),
} }
impl Logger { impl Logger {
@ -99,6 +101,10 @@ impl Logger {
exclude: HashSet::new(), exclude: HashSet::new(),
exclude_regex: RegexSet::empty(), exclude_regex: RegexSet::empty(),
log_target: Cow::Borrowed(module_path!()), log_target: Cow::Borrowed(module_path!()),
status_range: (
Bound::Included(StatusCode::from_u16(100).unwrap()),
Bound::Included(StatusCode::from_u16(999).unwrap()),
),
})) }))
} }
@ -121,6 +127,13 @@ impl Logger {
self self
} }
/// Set a range of status to include in the logging
pub fn status_range<R: RangeBounds<StatusCode>>(mut self, status: R) -> Self {
let inner = Rc::get_mut(&mut self.0).unwrap();
inner.status_range = (status.start_bound().cloned(), status.end_bound().cloned());
self
}
/// Sets the logging target to `target`. /// Sets the logging target to `target`.
/// ///
/// By default, the log target is `module_path!()` of the log call location. In our case, that /// By default, the log target is `module_path!()` of the log call location. In our case, that
@ -242,6 +255,10 @@ impl Default for Logger {
exclude: HashSet::new(), exclude: HashSet::new(),
exclude_regex: RegexSet::empty(), exclude_regex: RegexSet::empty(),
log_target: Cow::Borrowed(module_path!()), log_target: Cow::Borrowed(module_path!()),
status_range: (
Bound::Included(StatusCode::from_u16(100).unwrap()),
Bound::Included(StatusCode::from_u16(999).unwrap()),
),
})) }))
} }
} }
@ -306,6 +323,7 @@ where
LoggerResponse { LoggerResponse {
fut: self.service.call(req), fut: self.service.call(req),
format: None, format: None,
status_range: self.inner.status_range,
time: OffsetDateTime::now_utc(), time: OffsetDateTime::now_utc(),
log_target: Cow::Borrowed(""), log_target: Cow::Borrowed(""),
_phantom: PhantomData, _phantom: PhantomData,
@ -321,6 +339,7 @@ where
LoggerResponse { LoggerResponse {
fut: self.service.call(req), fut: self.service.call(req),
format: Some(format), format: Some(format),
status_range: self.inner.status_range,
time: now, time: now,
log_target: self.inner.log_target.clone(), log_target: self.inner.log_target.clone(),
_phantom: PhantomData, _phantom: PhantomData,
@ -339,6 +358,7 @@ pin_project! {
fut: S::Future, fut: S::Future,
time: OffsetDateTime, time: OffsetDateTime,
format: Option<Format>, format: Option<Format>,
status_range:(Bound<StatusCode>,Bound<StatusCode>),
log_target: Cow<'static, str>, log_target: Cow<'static, str>,
_phantom: PhantomData<B>, _phantom: PhantomData<B>,
} }
@ -363,22 +383,26 @@ where
debug!("Error in response: {:?}", error); debug!("Error in response: {:?}", error);
} }
let res = if let Some(ref mut format) = this.format { let res = if this.status_range.contains(&res.status()) {
// to avoid polluting all the Logger types with the body parameter we swap the body if let Some(ref mut format) = this.format {
// out temporarily since it's not usable in custom response functions anyway // 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 (req, res) = res.into_parts();
let (res, body) = res.into_parts(); let (res, body) = res.into_parts();
let temp_res = ServiceResponse::new(req, res.map_into_boxed_body()); let temp_res = ServiceResponse::new(req, res.map_into_boxed_body());
for unit in &mut format.0 { for unit in &mut format.0 {
unit.render_response(&temp_res); 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
} }
// re-construct original service response
let (req, res) = temp_res.into_parts();
ServiceResponse::new(req, res.set_body(body))
} else { } else {
res res
}; };