Skip to main content

api_tools/server/axum/
extractors.rs

1//! Extractor modules for Axum
2
3use crate::server::axum::layers::request_id::REQUEST_ID_HEADER;
4use crate::server::axum::response::ApiError;
5use axum::extract::FromRequestParts;
6use axum::extract::path::ErrorKind;
7use axum::extract::rejection::PathRejection;
8use axum::http::request::Parts;
9use axum::http::{HeaderValue, StatusCode};
10use serde::de::DeserializeOwned;
11
12/// Request ID extractor from HTTP headers
13pub struct RequestId(pub HeaderValue);
14
15impl<S> FromRequestParts<S> for RequestId
16where
17    S: Send + Sync,
18{
19    type Rejection = ();
20
21    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
22        match parts.headers.get(REQUEST_ID_HEADER.clone()) {
23            Some(id) => Ok(RequestId(id.clone())),
24            _ => Ok(RequestId(HeaderValue::from_static(""))),
25        }
26    }
27}
28
29/// `Path` extractor customizes the error from `axum::extract::Path`
30pub struct Path<T>(pub T);
31
32impl<S, T> FromRequestParts<S> for Path<T>
33where
34    T: DeserializeOwned + Send,
35    S: Send + Sync,
36{
37    type Rejection = (StatusCode, ApiError);
38
39    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
40        match axum::extract::Path::<T>::from_request_parts(parts, state).await {
41            Ok(value) => Ok(Self(value.0)),
42            Err(rejection) => {
43                let (status, body) = match rejection {
44                    PathRejection::FailedToDeserializePathParams(inner) => {
45                        let mut status = StatusCode::BAD_REQUEST;
46
47                        let kind = inner.into_kind();
48                        let body = match &kind {
49                            ErrorKind::WrongNumberOfParameters { .. } => ApiError::BadRequest(kind.to_string()),
50                            ErrorKind::ParseErrorAtKey { .. } => ApiError::BadRequest(kind.to_string()),
51                            ErrorKind::ParseErrorAtIndex { .. } => ApiError::BadRequest(kind.to_string()),
52                            ErrorKind::ParseError { .. } => ApiError::BadRequest(kind.to_string()),
53                            ErrorKind::InvalidUtf8InPathParam { .. } => ApiError::BadRequest(kind.to_string()),
54                            ErrorKind::UnsupportedType { .. } => {
55                                // this error is caused by the programmer using an unsupported type
56                                // (such as nested maps) so respond with `500` instead
57                                status = StatusCode::INTERNAL_SERVER_ERROR;
58                                ApiError::InternalServerError(kind.to_string())
59                            }
60                            ErrorKind::Message(msg) => ApiError::BadRequest(msg.clone()),
61                            _ => ApiError::BadRequest(format!("Unhandled deserialization error: {kind}")),
62                        };
63
64                        (status, body)
65                    }
66                    PathRejection::MissingPathParams(error) => (
67                        StatusCode::INTERNAL_SERVER_ERROR,
68                        ApiError::InternalServerError(error.to_string()),
69                    ),
70                    _ => (
71                        StatusCode::INTERNAL_SERVER_ERROR,
72                        ApiError::InternalServerError(format!("Unhandled path rejection: {rejection}")),
73                    ),
74                };
75
76                Err((status, body))
77            }
78        }
79    }
80}
81
82/// `Query` extractor customizes the error from `axum::extract::Query`
83pub struct Query<T>(pub T);
84
85impl<T, S> FromRequestParts<S> for Query<T>
86where
87    T: DeserializeOwned,
88    S: Send + Sync,
89{
90    type Rejection = (StatusCode, ApiError);
91
92    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
93        let query = parts.uri.query().unwrap_or_default();
94        let value = serde_urlencoded::from_str(query)
95            .map_err(|err| (StatusCode::BAD_REQUEST, ApiError::BadRequest(err.to_string())))?;
96
97        Ok(Query(value))
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use axum::Router;
105    use axum::body::Body;
106    use axum::http::Request;
107    use axum::routing::get;
108    use serde::Deserialize;
109    use tower::ServiceExt;
110
111    async fn read_body(response: axum::response::Response) -> String {
112        let body = axum::body::to_bytes(response.into_body(), 4096).await.unwrap();
113        String::from_utf8(body.to_vec()).unwrap()
114    }
115
116    // ---------------- RequestId ----------------
117
118    #[tokio::test]
119    async fn request_id_returns_header_value_when_present() {
120        let app: Router = Router::new().route(
121            "/",
122            get(|RequestId(value): RequestId| async move { value.to_str().unwrap_or_default().to_owned() }),
123        );
124
125        let response = app
126            .oneshot(
127                Request::builder()
128                    .uri("/")
129                    .header("x-request-id", "abc-123")
130                    .body(Body::empty())
131                    .unwrap(),
132            )
133            .await
134            .unwrap();
135
136        assert_eq!(response.status(), StatusCode::OK);
137        assert_eq!(read_body(response).await, "abc-123");
138    }
139
140    #[tokio::test]
141    async fn request_id_returns_empty_value_when_header_absent() {
142        let app: Router = Router::new().route(
143            "/",
144            get(|RequestId(value): RequestId| async move { value.to_str().unwrap_or_default().to_owned() }),
145        );
146
147        let response = app
148            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
149            .await
150            .unwrap();
151
152        assert_eq!(response.status(), StatusCode::OK);
153        assert_eq!(read_body(response).await, "");
154    }
155
156    // ---------------- Path ----------------
157
158    #[derive(Deserialize)]
159    struct PathArgs {
160        id: u64,
161    }
162
163    async fn path_handler(Path(args): Path<PathArgs>) -> String {
164        args.id.to_string()
165    }
166
167    #[tokio::test]
168    async fn path_extractor_deserializes_valid_param() {
169        let app: Router = Router::new().route("/users/{id}", get(path_handler));
170
171        let response = app
172            .oneshot(Request::builder().uri("/users/42").body(Body::empty()).unwrap())
173            .await
174            .unwrap();
175
176        assert_eq!(response.status(), StatusCode::OK);
177        assert_eq!(read_body(response).await, "42");
178    }
179
180    #[tokio::test]
181    async fn path_extractor_returns_400_with_api_error_on_type_mismatch() {
182        let app: Router = Router::new().route("/users/{id}", get(path_handler));
183
184        let response = app
185            .oneshot(Request::builder().uri("/users/abc").body(Body::empty()).unwrap())
186            .await
187            .unwrap();
188
189        assert_eq!(response.status(), StatusCode::BAD_REQUEST);
190        let body = read_body(response).await;
191        assert!(body.contains("\"code\":400"), "body was: {body}");
192    }
193
194    // ---------------- Query ----------------
195
196    #[derive(Deserialize)]
197    struct QueryArgs {
198        page: u32,
199        limit: u32,
200    }
201
202    async fn query_handler(Query(args): Query<QueryArgs>) -> String {
203        format!("{}-{}", args.page, args.limit)
204    }
205
206    #[tokio::test]
207    async fn query_extractor_deserializes_valid_query_string() {
208        let app: Router = Router::new().route("/", get(query_handler));
209
210        let response = app
211            .oneshot(Request::builder().uri("/?page=1&limit=20").body(Body::empty()).unwrap())
212            .await
213            .unwrap();
214
215        assert_eq!(response.status(), StatusCode::OK);
216        assert_eq!(read_body(response).await, "1-20");
217    }
218
219    #[tokio::test]
220    async fn query_extractor_returns_400_when_required_keys_missing() {
221        let app: Router = Router::new().route("/", get(query_handler));
222
223        let response = app
224            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
225            .await
226            .unwrap();
227
228        assert_eq!(response.status(), StatusCode::BAD_REQUEST);
229        let body = read_body(response).await;
230        assert!(body.contains("\"code\":400"), "body was: {body}");
231    }
232
233    #[tokio::test]
234    async fn query_extractor_returns_400_on_type_mismatch() {
235        let app: Router = Router::new().route("/", get(query_handler));
236
237        let response = app
238            .oneshot(
239                Request::builder()
240                    .uri("/?page=abc&limit=20")
241                    .body(Body::empty())
242                    .unwrap(),
243            )
244            .await
245            .unwrap();
246
247        assert_eq!(response.status(), StatusCode::BAD_REQUEST);
248    }
249}