Skip to main content

rama_http/service/web/endpoint/extract/
header_map.rs

1use super::FromOwnedRequestParts;
2use crate::{HeaderMap, request::Parts};
3use std::convert::Infallible;
4
5/// Move all headers out of request parts without cloning the map.
6impl FromOwnedRequestParts for HeaderMap {
7    type Rejection = Infallible;
8
9    fn from_owned_request_parts(
10        parts: Parts,
11    ) -> impl Future<Output = Result<Self, Self::Rejection>> + Send + 'static {
12        std::future::ready(Ok(parts.headers))
13    }
14}
15
16#[cfg(test)]
17mod tests {
18    use super::*;
19    use crate::{
20        Body, Request, StatusCode,
21        body::util::BodyExt,
22        header,
23        service::web::{
24            WebService,
25            extract::{Body as ExtractBody, FromRequest, Json, State},
26        },
27    };
28    use rama_core::Service;
29
30    #[tokio::test]
31    async fn extracts_all_request_headers_as_terminal_extractor() {
32        let request = Request::builder()
33            .header(header::ACCEPT, "application/json")
34            .header(header::COOKIE, "session=secret")
35            .header("x-custom-header", "custom-value")
36            .body(Body::empty())
37            .unwrap();
38
39        let headers = HeaderMap::from_request(request).await.unwrap();
40
41        assert_eq!(headers.len(), 3);
42        assert_eq!(headers.get(header::ACCEPT).unwrap(), "application/json");
43        assert_eq!(headers.get(header::COOKIE).unwrap(), "session=secret");
44        assert_eq!(headers.get("x-custom-header").unwrap(), "custom-value");
45    }
46
47    #[tokio::test]
48    async fn extracts_owned_headers_from_consumed_request_parts() {
49        let request = Request::builder()
50            .header(header::AUTHORIZATION, "Bearer secret")
51            .body(Body::empty())
52            .unwrap();
53        let (parts, _) = request.into_parts();
54
55        let headers = HeaderMap::from_owned_request_parts(parts).await.unwrap();
56
57        assert_eq!(headers.get(header::AUTHORIZATION).unwrap(), "Bearer secret");
58    }
59
60    #[tokio::test]
61    async fn composes_after_regular_parts_and_before_json_body_extractor() {
62        #[derive(serde::Deserialize)]
63        struct Input {
64            message: String,
65        }
66
67        let service = WebService::new_with_state("prefix".to_owned()).with_post(
68            "/",
69            async |State(prefix): State<String>, headers: HeaderMap, Json(input): Json<Input>| {
70                assert_eq!(headers.get(header::AUTHORIZATION).unwrap(), "Bearer secret");
71                assert_eq!(headers.get(header::ORIGIN).unwrap(), "https://example.com");
72                format!("{prefix}: {}", input.message)
73            },
74        );
75        let request = Request::builder()
76            .method(crate::Method::POST)
77            .uri("/")
78            .header(header::AUTHORIZATION, "Bearer secret")
79            .header(header::ORIGIN, "https://example.com")
80            .header(header::CONTENT_TYPE, "application/json")
81            .body(Body::from(r#"{"message":"body remains available"}"#))
82            .unwrap();
83
84        let response = service.serve(request).await.unwrap();
85
86        assert_eq!(response.status(), StatusCode::OK);
87        let body = response.into_body().collect().await.unwrap().to_bytes();
88        assert_eq!(body, "prefix: body remains available");
89    }
90
91    #[tokio::test]
92    async fn composes_with_optional_body_extractor() {
93        #[derive(serde::Deserialize)]
94        struct Input {
95            _message: String,
96        }
97
98        let service = WebService::default().with_post(
99            "/",
100            async |headers: HeaderMap, body: Option<Json<Input>>| {
101                assert_eq!(headers.get(header::AUTHORIZATION).unwrap(), "Bearer secret");
102                assert!(body.is_none());
103            },
104        );
105        let request = Request::builder()
106            .method(crate::Method::POST)
107            .uri("/")
108            .header(header::AUTHORIZATION, "Bearer secret")
109            .body(Body::empty())
110            .unwrap();
111
112        let response = service.serve(request).await.unwrap();
113
114        assert_eq!(response.status(), StatusCode::OK);
115    }
116
117    #[tokio::test]
118    async fn preserves_body_extractor_rejections() {
119        #[derive(serde::Deserialize)]
120        struct Input {
121            _message: String,
122        }
123
124        let service = WebService::default()
125            .with_post("/", async |_headers: HeaderMap, Json(_): Json<Input>| {
126                StatusCode::OK
127            });
128        let request = Request::builder()
129            .method(crate::Method::POST)
130            .uri("/")
131            .header(header::CONTENT_TYPE, "text/plain")
132            .body(Body::from(r#"{"message":"invalid content type"}"#))
133            .unwrap();
134
135        let response = service.serve(request).await.unwrap();
136
137        assert_eq!(response.status(), StatusCode::UNSUPPORTED_MEDIA_TYPE);
138    }
139
140    #[tokio::test]
141    async fn owned_parts_rejection_precedes_body_rejection() {
142        struct RejectingHeaders;
143
144        impl FromOwnedRequestParts for RejectingHeaders {
145            type Rejection = StatusCode;
146
147            fn from_owned_request_parts(
148                _parts: Parts,
149            ) -> impl Future<Output = Result<Self, Self::Rejection>> + Send + 'static {
150                std::future::ready(Err(StatusCode::UNAUTHORIZED))
151            }
152        }
153
154        #[derive(serde::Deserialize)]
155        struct Input {
156            _message: String,
157        }
158
159        let service = WebService::default().with_post(
160            "/",
161            async |_headers: RejectingHeaders, Json(_): Json<Input>| StatusCode::OK,
162        );
163        let request = Request::builder()
164            .method(crate::Method::POST)
165            .uri("/")
166            .header(header::CONTENT_TYPE, "text/plain")
167            .body(Body::from(r#"{"message":"invalid content type"}"#))
168            .unwrap();
169
170        let response = service.serve(request).await.unwrap();
171
172        assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
173    }
174
175    #[tokio::test]
176    async fn composes_with_streaming_body_extractor() {
177        let service = WebService::default().with_post(
178            "/",
179            async |headers: HeaderMap, ExtractBody(body): ExtractBody| {
180                assert_eq!(headers.get("x-request-id").unwrap(), "123");
181                body.collect().await.unwrap().to_bytes()
182            },
183        );
184        let request = Request::builder()
185            .method(crate::Method::POST)
186            .uri("/")
187            .header("x-request-id", "123")
188            .body(Body::from("stream me"))
189            .unwrap();
190
191        let response = service.serve(request).await.unwrap();
192
193        assert_eq!(response.status(), StatusCode::OK);
194        let body = response.into_body().collect().await.unwrap().to_bytes();
195        assert_eq!(body, "stream me");
196    }
197}