api_tools/server/axum/layers/
logger.rs1use 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 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 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 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 #[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 #[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 #[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 #[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}