Skip to main content

api_tools/server/axum/layers/
logger.rs

1//! Logger layer
2
3use super::header_value_to_str;
4use axum::body::HttpBody;
5use axum::http::{Method, StatusCode};
6use axum::{body::Body, http::Request, response::Response};
7use bytesize::ByteSize;
8use futures::future::BoxFuture;
9use std::{
10    fmt::Display,
11    task::{Context, Poll},
12    time::{Duration, Instant},
13};
14use tower::{Layer, Service};
15
16#[derive(Debug, Default)]
17struct LoggerMessage {
18    method: String,
19    request_id: String,
20    host: String,
21    path: String,
22    uri: String,
23    user_agent: String,
24    status_code: u16,
25    version: String,
26    latency: Duration,
27    body_size: u64,
28}
29
30impl Display for LoggerMessage {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        write!(
33            f,
34            "status_code: {}, method: {}, path: {}, uri: {}, host: {}, request_id: {}, user_agent: {}, version: {}, latency: {:?}, body_size: {}",
35            self.status_code,
36            self.method,
37            self.path,
38            self.uri,
39            self.host,
40            self.request_id,
41            self.user_agent,
42            self.version,
43            self.latency,
44            ByteSize::b(self.body_size),
45        )
46    }
47}
48
49#[derive(Clone)]
50pub struct LoggerLayer;
51
52impl<S> Layer<S> for LoggerLayer {
53    type Service = LoggerMiddleware<S>;
54
55    fn layer(&self, inner: S) -> Self::Service {
56        LoggerMiddleware { inner }
57    }
58}
59
60#[derive(Clone)]
61pub struct LoggerMiddleware<S> {
62    inner: S,
63}
64
65impl<S> Service<Request<Body>> for LoggerMiddleware<S>
66where
67    S: Service<Request<Body>, Response = Response> + Send + 'static,
68    S::Future: Send + 'static,
69{
70    type Response = S::Response;
71    type Error = S::Error;
72    // `BoxFuture` is a type alias for `Pin<Box<dyn Future + Send + 'a>>`
73    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
74
75    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
76        self.inner.poll_ready(cx)
77    }
78
79    fn call(&mut self, request: Request<Body>) -> Self::Future {
80        let now = Instant::now();
81        let request_headers = request.headers();
82
83        let message = LoggerMessage {
84            method: request.method().to_string(),
85            path: request.uri().path().to_string(),
86            uri: request.uri().to_string(),
87            host: header_value_to_str(request_headers.get("host")).to_string(),
88            request_id: header_value_to_str(request_headers.get("x-request-id")).to_string(),
89            user_agent: header_value_to_str(request_headers.get("user-agent")).to_string(),
90            ..Default::default()
91        };
92
93        let future = self.inner.call(request);
94        Box::pin(async move {
95            let response: Response = future.await?;
96
97            // Skip logging for OPTIONS requests
98            if message.method == Method::OPTIONS.to_string() {
99                return Ok(response);
100            }
101
102            let status_code = response.status().as_u16();
103            let version = format!("{:?}", response.version());
104            let latency = now.elapsed();
105            let body_size = response.body().size_hint().lower();
106
107            macro_rules! log_request {
108                ($level:ident) => {
109                    $level!(
110                        status_code = %status_code,
111                        method = %message.method,
112                        path = %message.path,
113                        uri = %message.uri,
114                        host = %message.host,
115                        request_id = %message.request_id,
116                        user_agent = %message.user_agent,
117                        version = %version,
118                        latency = %format!("{:?}", latency),
119                        body_size = %ByteSize::b(body_size),
120                    );
121                };
122            }
123
124            if response.status() >= StatusCode::INTERNAL_SERVER_ERROR
125                && response.status() != StatusCode::SERVICE_UNAVAILABLE
126            {
127                log_request!(error);
128            } else if !message.path.starts_with("/metrics") {
129                log_request!(info);
130            }
131
132            Ok(response)
133        })
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140    use std::convert::Infallible;
141    use std::time::Duration;
142    use tower::{ServiceBuilder, ServiceExt};
143
144    fn make_svc(
145        status: StatusCode,
146    ) -> impl tower::Service<
147        Request<Body>,
148        Response = Response,
149        Error = Infallible,
150        Future = impl std::future::Future<Output = Result<Response, Infallible>> + Send,
151    > + Clone
152    + Send
153    + 'static {
154        let layer = LoggerLayer;
155        ServiceBuilder::new()
156            .layer(layer)
157            .service(tower::service_fn(move |_req: Request<Body>| async move {
158                Ok::<_, Infallible>(Response::builder().status(status).body(Body::from("ok")).unwrap())
159            }))
160    }
161
162    #[tokio::test]
163    async fn middleware_passes_through_2xx_response() {
164        let svc = make_svc(StatusCode::OK);
165        let req = Request::builder()
166            .method(Method::GET)
167            .uri("/foo")
168            .header("host", "localhost")
169            .header("user-agent", "test/1.0")
170            .header("x-request-id", "abc")
171            .body(Body::empty())
172            .unwrap();
173
174        let resp = svc.oneshot(req).await.unwrap();
175        assert_eq!(resp.status(), StatusCode::OK);
176        let body = axum::body::to_bytes(resp.into_body(), 1024).await.unwrap();
177        assert_eq!(&body[..], b"ok");
178    }
179
180    #[tokio::test]
181    async fn middleware_passes_through_5xx_response() {
182        // 5xx (except 503) takes the ERROR-level branch.
183        let svc = make_svc(StatusCode::INTERNAL_SERVER_ERROR);
184        let req = Request::builder().uri("/boom").body(Body::empty()).unwrap();
185        let resp = svc.oneshot(req).await.unwrap();
186        assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
187    }
188
189    /// SERVICE_UNAVAILABLE is the only 5xx still logged at info level — it is
190    /// the legitimate "outside time slot" response from `time_limiter`, not a
191    /// real failure.
192    #[tokio::test]
193    async fn middleware_treats_503_as_info_branch() {
194        let svc = make_svc(StatusCode::SERVICE_UNAVAILABLE);
195        let req = Request::builder().uri("/").body(Body::empty()).unwrap();
196        let resp = svc.oneshot(req).await.unwrap();
197        assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
198    }
199
200    /// Sentinel: an OPTIONS request must short-circuit the logger before any
201    /// status/latency computation. Pins the regression added in commit
202    /// 76f81c7.
203    #[tokio::test]
204    async fn middleware_short_circuits_on_options_method() {
205        let svc = make_svc(StatusCode::NO_CONTENT);
206        let req = Request::builder()
207            .method(Method::OPTIONS)
208            .uri("/foo")
209            .body(Body::empty())
210            .unwrap();
211        let resp = svc.oneshot(req).await.unwrap();
212        assert_eq!(resp.status(), StatusCode::NO_CONTENT);
213    }
214
215    /// `/metrics` requests are passed through but skip the info log to avoid
216    /// a feedback loop with the Prometheus scraper.
217    #[tokio::test]
218    async fn middleware_passes_through_metrics_path() {
219        let svc = make_svc(StatusCode::OK);
220        let req = Request::builder().uri("/metrics").body(Body::empty()).unwrap();
221        let resp = svc.oneshot(req).await.unwrap();
222        assert_eq!(resp.status(), StatusCode::OK);
223    }
224
225    /// Missing optional headers (host, user-agent, x-request-id) must default
226    /// to empty strings without panicking.
227    #[tokio::test]
228    async fn middleware_handles_missing_optional_headers() {
229        let svc = make_svc(StatusCode::OK);
230        let req = Request::builder().uri("/").body(Body::empty()).unwrap();
231        let resp = svc.oneshot(req).await.unwrap();
232        assert_eq!(resp.status(), StatusCode::OK);
233    }
234
235    #[test]
236    fn test_logger_message_fmt() {
237        let message = LoggerMessage {
238            method: "GET".to_string(),
239            request_id: "abc-123".to_string(),
240            host: "localhost".to_string(),
241            path: "/test".to_string(),
242            uri: "/test?query=1".to_string(),
243            user_agent: "TestAgent/1.0".to_string(),
244            status_code: 200,
245            version: "HTTP/1.1".to_string(),
246            latency: Duration::from_millis(42),
247            body_size: 1_524,
248        };
249        let expected = String::from(
250            "status_code: 200, method: GET, path: /test, uri: /test?query=1, host: localhost, request_id: abc-123, user_agent: TestAgent/1.0, version: HTTP/1.1, latency: 42ms, body_size: 1.5 KiB",
251        );
252
253        assert_eq!(message.to_string(), expected);
254    }
255}