actix_web/response/
customize_responder.rs1use actix_http::{
2 body::EitherBody,
3 error::HttpError,
4 header::{HeaderMap, TryIntoHeaderPair},
5 StatusCode,
6};
7
8use crate::{HttpRequest, HttpResponse, Responder};
9
10pub 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 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 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 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 #[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}