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