1//! Logs each and every request out in a format similar to that of Apache's logs.
  2
  3use std::{
  4    fmt::Debug,
  5    net::SocketAddr,
  6    task::{Context, Poll},
  7    time::Instant,
  8};
  9
 10use axum::{
 11    extract,
 12    http::{HeaderValue, Method, Request, Response},
 13};
 14use futures::future::{Future, FutureExt, Join, Map, Ready};
 15use tower_service::Service;
 16use tracing::{error, info, instrument::Instrumented, Instrument, Span};
 17use uuid::Uuid;
 18
 19use super::UnwrapInfallible;
 20
 21pub trait GenericError: std::error::Error + Debug + Send + Sync {}
 22
 23#[derive(Clone)]
 24pub struct LoggingMiddleware<S>(pub S);
 25
 26impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for LoggingMiddleware<S>
 27where
 28    S: Service<Request<ReqBody>, Response = Response<ResBody>, Error = std::convert::Infallible>
 29        + Clone
 30        + Send
 31        + 'static,
 32    S::Future: Send + 'static,
 33    S::Response: Default + Debug,
 34    ReqBody: Send + Debug + 'static,
 35    ResBody: Default + Send + 'static,
 36{
 37    type Response = S::Response;
 38    type Error = S::Error;
 39    type Future = Map<
 40        Join<Instrumented<S::Future>, Ready<PendingLogMessage>>,
 41        fn((<S::Future as Future>::Output, PendingLogMessage)) -> <S::Future as Future>::Output,
 42    >;
 43
 44    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
 45        self.0.poll_ready(cx)
 46    }
 47
 48    fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
 49        let request_id = Uuid::new_v4();
 50        let span = tracing::info_span!("web", "request_id" = request_id.to_string().as_str());
 51
 52        let log_message = PendingLogMessage {
 53            span: span.clone(),
 54            request_id,
 55            ip: req
 56                .extensions()
 57                .get::<extract::ConnectInfo<std::net::SocketAddr>>()
 58                .map_or_else(|| "0.0.0.0:0".parse().unwrap(), |v| v.0),
 59            method: req.method().clone(),
 60            uri: req.uri().path().to_string(),
 61            start: Instant::now(),
 62            user_agent: req.headers().get(axum::http::header::USER_AGENT).cloned(),
 63        };
 64
 65        futures::future::join(
 66            self.0.call(req).instrument(span),
 67            futures::future::ready(log_message),
 68        )
 69        .map(|(response, pending_log_message)| {
 70            let mut response = response.unwrap_infallible();
 71            pending_log_message.log(&response);
 72            response.headers_mut().insert(
 73                "X-Request-ID",
 74                HeaderValue::try_from(pending_log_message.request_id.to_string()).unwrap(),
 75            );
 76            Ok(response)
 77        })
 78    }
 79}
 80
 81pub struct PendingLogMessage {
 82    span: Span,
 83    request_id: Uuid,
 84    ip: SocketAddr,
 85    method: Method,
 86    uri: String,
 87    start: Instant,
 88    user_agent: Option<HeaderValue>,
 89}
 90
 91impl PendingLogMessage {
 92    pub fn log<ResBody>(&self, response: &Response<ResBody>) {
 93        let _enter = self.span.enter();
 94
 95        if response.status().is_server_error() {
 96            error!(
 97                "{ip} - \"{method} {uri}\" {status} {duration:?} \"{user_agent}\" \"{error:?}\"",
 98                ip = self.ip,
 99                method = self.method,
100                uri = self.uri,
101                status = response.status().as_u16(),
102                duration = self.start.elapsed(),
103                user_agent = self
104                    .user_agent
105                    .as_ref()
106                    .and_then(|v| v.to_str().ok())
107                    .unwrap_or("unknown"),
108                error = match response.extensions().get::<Box<dyn GenericError>>() {
109                    Some(e) => Err(e),
110                    None => Ok(()),
111                }
112            );
113        } else {
114            info!(
115                "{ip} - \"{method} {uri}\" {status} {duration:?} \"{user_agent}\" \"{error:?}\"",
116                ip = self.ip,
117                method = self.method,
118                uri = self.uri,
119                status = response.status().as_u16(),
120                duration = self.start.elapsed(),
121                user_agent = self
122                    .user_agent
123                    .as_ref()
124                    .and_then(|v| v.to_str().ok())
125                    .unwrap_or("unknown"),
126                error = match response.extensions().get::<Box<dyn GenericError>>() {
127                    Some(e) => Err(e),
128                    None => Ok(()),
129                }
130            );
131        }
132    }
133}