actix_web/response/
customize_responder.rs

1use actix_http::{
2    body::EitherBody,
3    error::HttpError,
4    header::{HeaderMap, TryIntoHeaderPair},
5    StatusCode,
6};
7
8use crate::{HttpRequest, HttpResponse, Responder};
9
10/// Allows overriding status code and headers (including cookies) for a [`Responder`].
11///
12/// Created by calling the [`customize`](Responder::customize) method on a [`Responder`] type.
13pub struct CustomizeResponder<R> {
14    inner: CustomizeResponderInner<R>,
15    error: Option<HttpError>,
16}
17
18struct CustomizeResponderInner<R> {
19    responder: R,
20    status: Option<StatusCode>,
21    override_headers: HeaderMap,
22    append_headers: HeaderMap,
23}
24
25impl<R: Responder> CustomizeResponder<R> {
26    pub(crate) fn new(responder: R) -> Self {
27        CustomizeResponder {
28            inner: CustomizeResponderInner {
29                responder,
30                status: None,
31                override_headers: HeaderMap::new(),
32                append_headers: HeaderMap::new(),
33            },
34            error: None,
35        }
36    }
37
38    /// Override a status code for the Responder's response.
39    ///
40    /// # Examples
41    /// ```
42    /// use actix_web::{Responder, http::StatusCode, test::TestRequest};
43    ///
44    /// let responder = "Welcome!".customize().with_status(StatusCode::ACCEPTED);
45    ///
46    /// let request = TestRequest::default().to_http_request();
47    /// let response = responder.respond_to(&request);
48    /// assert_eq!(response.status(), StatusCode::ACCEPTED);
49    /// ```
50    pub fn with_status(mut self, status: StatusCode) -> Self {
51        if let Some(inner) = self.inner() {
52            inner.status = Some(status);
53        }
54
55        self
56    }
57
58    /// Insert (override) header in the final response.
59    ///
60    /// Overrides other headers with the same name.
61    /// See [`HeaderMap::insert`](crate::http::header::HeaderMap::insert).
62    ///
63    /// Headers added with this method will be inserted before those added
64    /// with [`append_header`](Self::append_header). As such, header(s) can be overridden with more
65    /// than one new header by first calling `insert_header` followed by `append_header`.
66    ///
67    /// # Examples
68    /// ```
69    /// use actix_web::{Responder, test::TestRequest};
70    ///
71    /// let responder = "Hello world!"
72    ///     .customize()
73    ///     .insert_header(("x-version", "1.2.3"));
74    ///
75    /// let request = TestRequest::default().to_http_request();
76    /// let response = responder.respond_to(&request);
77    /// assert_eq!(response.headers().get("x-version").unwrap(), "1.2.3");
78    /// ```
79    pub fn insert_header(mut self, header: impl TryIntoHeaderPair) -> Self {
80        if let Some(inner) = self.inner() {
81            match header.try_into_pair() {
82                Ok((key, value)) => {
83                    inner.override_headers.insert(key, value);
84                }
85                Err(err) => self.error = Some(err.into()),
86            };
87        }
88
89        self
90    }
91
92    /// Append header to the final response.
93    ///
94    /// Unlike [`insert_header`](Self::insert_header), this will not override existing headers.
95    /// See [`HeaderMap::append`](crate::http::header::HeaderMap::append).
96    ///
97    /// Headers added here are appended _after_ additions/overrides from `insert_header`.
98    ///
99    /// # Examples
100    /// ```
101    /// use actix_web::{Responder, test::TestRequest};
102    ///
103    /// let responder = "Hello world!"
104    ///     .customize()
105    ///     .append_header(("x-version", "1.2.3"));
106    ///
107    /// let request = TestRequest::default().to_http_request();
108    /// let response = responder.respond_to(&request);
109    /// assert_eq!(response.headers().get("x-version").unwrap(), "1.2.3");
110    /// ```
111    pub fn append_header(mut self, header: impl TryIntoHeaderPair) -> Self {
112        if let Some(inner) = self.inner() {
113            match header.try_into_pair() {
114                Ok((key, value)) => {
115                    inner.append_headers.append(key, value);
116                }
117                Err(err) => self.error = Some(err.into()),
118            };
119        }
120
121        self
122    }
123
124    #[doc(hidden)]
125    #[deprecated(since = "4.0.0", note = "Renamed to `insert_header`.")]
126    pub fn with_header(self, header: impl TryIntoHeaderPair) -> Self
127    where
128        Self: Sized,
129    {
130        self.insert_header(header)
131    }
132
133    fn inner(&mut self) -> Option<&mut CustomizeResponderInner<R>> {
134        if self.error.is_some() {
135            None
136        } else {
137            Some(&mut self.inner)
138        }
139    }
140
141    /// Appends a `cookie` to the final response.
142    ///
143    /// # Errors
144    ///
145    /// Final response will be an error if `cookie` cannot be converted into a valid header value.
146    #[cfg(feature = "cookies")]
147    pub fn add_cookie(mut self, cookie: &crate::cookie::Cookie<'_>) -> Self {
148        use actix_http::header::{TryIntoHeaderValue as _, SET_COOKIE};
149
150        if let Some(inner) = self.inner() {
151            match cookie.to_string().try_into_value() {
152                Ok(val) => {
153                    inner.append_headers.append(SET_COOKIE, val);
154                }
155                Err(err) => {
156                    self.error = Some(err.into());
157                }
158            }
159        }
160
161        self
162    }
163}
164
165impl<T> Responder for CustomizeResponder<T>
166where
167    T: Responder,
168{
169    type Body = EitherBody<T::Body>;
170
171    fn respond_to(self, req: &HttpRequest) -> HttpResponse<Self::Body> {
172        if let Some(err) = self.error {
173            return HttpResponse::from_error(err).map_into_right_body();
174        }
175
176        let mut res = self.inner.responder.respond_to(req);
177
178        if let Some(status) = self.inner.status {
179            *res.status_mut() = status;
180        }
181
182        for (k, v) in self.inner.override_headers {
183            res.headers_mut().insert(k, v);
184        }
185
186        for (k, v) in self.inner.append_headers {
187            res.headers_mut().append(k, v);
188        }
189
190        res.map_into_left_body()
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use actix_http::body::to_bytes;
197    use bytes::Bytes;
198
199    use super::*;
200    use crate::{
201        cookie::Cookie,
202        http::header::{HeaderValue, CONTENT_TYPE},
203        test::TestRequest,
204    };
205
206    #[actix_rt::test]
207    async fn customize_responder() {
208        let req = TestRequest::default().to_http_request();
209        let res = "test"
210            .to_string()
211            .customize()
212            .with_status(StatusCode::BAD_REQUEST)
213            .respond_to(&req);
214
215        assert_eq!(res.status(), StatusCode::BAD_REQUEST);
216        assert_eq!(
217            to_bytes(res.into_body()).await.unwrap(),
218            Bytes::from_static(b"test"),
219        );
220
221        let res = "test"
222            .to_string()
223            .customize()
224            .insert_header(("content-type", "json"))
225            .respond_to(&req);
226
227        assert_eq!(res.status(), StatusCode::OK);
228        assert_eq!(
229            res.headers().get(CONTENT_TYPE).unwrap(),
230            HeaderValue::from_static("json")
231        );
232        assert_eq!(
233            to_bytes(res.into_body()).await.unwrap(),
234            Bytes::from_static(b"test"),
235        );
236
237        let res = "test"
238            .to_string()
239            .customize()
240            .add_cookie(&Cookie::new("name", "value"))
241            .respond_to(&req);
242
243        assert!(res.status().is_success());
244        assert_eq!(
245            res.cookies().collect::<Vec<Cookie<'_>>>(),
246            vec![Cookie::new("name", "value")],
247        );
248        assert_eq!(
249            to_bytes(res.into_body()).await.unwrap(),
250            Bytes::from_static(b"test"),
251        );
252    }
253
254    #[actix_rt::test]
255    async fn tuple_responder_with_status_code() {
256        let req = TestRequest::default().to_http_request();
257        let res = ("test".to_string(), StatusCode::BAD_REQUEST).respond_to(&req);
258        assert_eq!(res.status(), StatusCode::BAD_REQUEST);
259        assert_eq!(
260            to_bytes(res.into_body()).await.unwrap(),
261            Bytes::from_static(b"test"),
262        );
263
264        let req = TestRequest::default().to_http_request();
265        let res = ("test".to_string(), StatusCode::OK)
266            .customize()
267            .insert_header((CONTENT_TYPE, mime::APPLICATION_JSON))
268            .respond_to(&req);
269        assert_eq!(res.status(), StatusCode::OK);
270        assert_eq!(
271            res.headers().get(CONTENT_TYPE).unwrap(),
272            HeaderValue::from_static("application/json")
273        );
274        assert_eq!(
275            to_bytes(res.into_body()).await.unwrap(),
276            Bytes::from_static(b"test"),
277        );
278    }
279}