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