salvo_core/http/
response.rs

1//! HTTP response.
2use std::collections::VecDeque;
3use std::fmt::{self, Debug, Display, Formatter};
4use std::path::PathBuf;
5
6#[cfg(feature = "cookie")]
7use cookie::{Cookie, CookieJar};
8use futures_util::stream::Stream;
9use http::header::{HeaderMap, HeaderValue, IntoHeaderName};
10pub use http::response::Parts;
11use http::{Extensions, version::Version};
12use mime::Mime;
13
14use crate::fs::NamedFile;
15use crate::fuse::TransProto;
16use crate::http::{StatusCode, StatusError};
17use crate::{BoxedError, Error, Scribe};
18use bytes::Bytes;
19
20pub use crate::http::body::{BodySender, BytesFrame, ResBody};
21
22/// Represents an HTTP response.
23#[non_exhaustive]
24pub struct Response {
25    /// The HTTP status code.WebTransportSession
26    pub status_code: Option<StatusCode>,
27    /// The HTTP headers.
28    pub headers: HeaderMap,
29    /// The HTTP version.
30    pub version: Version,
31    /// The HTTP cookies.
32    #[cfg(feature = "cookie")]
33    pub cookies: CookieJar,
34    /// The HTTP body.
35    pub body: ResBody,
36    /// Used to store extra data derived from the underlying protocol.
37    pub extensions: Extensions,
38}
39impl Default for Response {
40    #[inline]
41    fn default() -> Self {
42        Self::new()
43    }
44}
45impl<B> From<hyper::Response<B>> for Response
46where
47    B: Into<ResBody>,
48{
49    #[inline]
50    fn from(res: hyper::Response<B>) -> Self {
51        let (
52            http::response::Parts {
53                status,
54                version,
55                headers,
56                // extensions,
57                ..
58            },
59            body,
60        ) = res.into_parts();
61        #[cfg(feature = "cookie")]
62        // Set the request cookies, if they exist.
63        let cookies = if let Some(header) = headers.get(http::header::SET_COOKIE) {
64            let mut cookie_jar = CookieJar::new();
65            if let Ok(header) = header.to_str() {
66                for cookie_str in header.split(';').map(|s| s.trim()) {
67                    if let Ok(cookie) = Cookie::parse_encoded(cookie_str).map(|c| c.into_owned()) {
68                        cookie_jar.add(cookie);
69                    }
70                }
71            }
72            cookie_jar
73        } else {
74            CookieJar::new()
75        };
76
77        Self {
78            status_code: Some(status),
79            body: body.into(),
80            version,
81            headers,
82            #[cfg(feature = "cookie")]
83            cookies,
84            extensions: Extensions::new(),
85        }
86    }
87}
88
89impl Response {
90    /// Creates a new blank `Response`.
91    #[inline]
92    #[must_use]
93    pub fn new() -> Self {
94        Self {
95            status_code: None,
96            body: ResBody::None,
97            version: Version::default(),
98            headers: HeaderMap::new(),
99            #[cfg(feature = "cookie")]
100            cookies: CookieJar::default(),
101            extensions: Extensions::new(),
102        }
103    }
104
105    /// Creates a new blank `Response`.
106    #[cfg(feature = "cookie")]
107    #[inline]
108    #[must_use]
109    pub fn with_cookies(cookies: CookieJar) -> Self {
110        Self {
111            status_code: None,
112            body: ResBody::None,
113            version: Version::default(),
114            headers: HeaderMap::new(),
115            cookies,
116            extensions: Extensions::new(),
117        }
118    }
119
120    /// Get headers reference.
121    #[inline]
122    pub fn headers(&self) -> &HeaderMap {
123        &self.headers
124    }
125    /// Get mutable headers reference.
126    #[inline]
127    pub fn headers_mut(&mut self) -> &mut HeaderMap {
128        &mut self.headers
129    }
130    /// Sets headers.
131    #[inline]
132    pub fn set_headers(&mut self, headers: HeaderMap) {
133        self.headers = headers
134    }
135
136    /// Modify a header for this response.
137    ///
138    /// When `overwrite` is set to `true`, If the header is already present, the value will be replaced.
139    /// When `overwrite` is set to `false`, The new header is always appended to the request, even if the header already exists.
140    pub fn add_header<N, V>(
141        &mut self,
142        name: N,
143        value: V,
144        overwrite: bool,
145    ) -> crate::Result<&mut Self>
146    where
147        N: IntoHeaderName,
148        V: TryInto<HeaderValue>,
149    {
150        let value = value
151            .try_into()
152            .map_err(|_| Error::Other("invalid header value".into()))?;
153        if overwrite {
154            self.headers.insert(name, value);
155        } else {
156            self.headers.append(name, value);
157        }
158        Ok(self)
159    }
160
161    /// Get version.
162    #[inline]
163    pub fn version(&self) -> Version {
164        self.version
165    }
166    /// Get mutable version reference.
167    #[inline]
168    pub fn version_mut(&mut self) -> &mut Version {
169        &mut self.version
170    }
171    #[doc(hidden)]
172    pub fn trans_proto(&self) -> TransProto {
173        if self.version == Version::HTTP_3 {
174            TransProto::Quic
175        } else {
176            TransProto::Tcp
177        }
178    }
179
180    /// Get mutable body reference.
181    #[inline]
182    pub fn body_mut(&mut self) -> &mut ResBody {
183        &mut self.body
184    }
185    /// Sets body.
186    #[inline]
187    pub fn body(&mut self, body: impl Into<ResBody>) -> &mut Self {
188        self.body = body.into();
189        self
190    }
191
192    /// Sets body to a new value and returns old value.
193    #[inline]
194    pub fn replace_body(&mut self, body: ResBody) -> ResBody {
195        std::mem::replace(&mut self.body, body)
196    }
197
198    /// Take body from response.
199    #[inline]
200    pub fn take_body(&mut self) -> ResBody {
201        self.replace_body(ResBody::None)
202    }
203
204    /// If returns `true`, it means this response is ready for write back and the reset handlers should be skipped.
205    #[inline]
206    pub fn is_stamped(&mut self) -> bool {
207        if let Some(code) = self.status_code {
208            if code.is_client_error() || code.is_server_error() || code.is_redirection() {
209                return true;
210            }
211        }
212        false
213    }
214
215    /// Convert to hyper response.
216    #[doc(hidden)]
217    #[inline]
218    pub fn into_hyper(self) -> hyper::Response<ResBody> {
219        let Self {
220            status_code,
221            #[cfg(feature = "cookie")]
222            mut headers,
223            #[cfg(feature = "cookie")]
224            cookies,
225            #[cfg(not(feature = "cookie"))]
226            headers,
227            body,
228            extensions,
229            ..
230        } = self;
231
232        #[cfg(feature = "cookie")]
233        for cookie in cookies.delta() {
234            if let Ok(hv) = cookie.encoded().to_string().parse() {
235                headers.append(http::header::SET_COOKIE, hv);
236            }
237        }
238
239        let status_code = status_code.unwrap_or(match &body {
240            ResBody::None => StatusCode::NOT_FOUND,
241            ResBody::Error(e) => e.code,
242            _ => StatusCode::OK,
243        });
244        let mut res = hyper::Response::new(body);
245        *res.extensions_mut() = extensions;
246        *res.headers_mut() = headers;
247        *res.status_mut() = status_code;
248
249        res
250    }
251
252    /// Strip the response to [`hyper::Response`].
253    #[doc(hidden)]
254    #[inline]
255    pub fn strip_to_hyper(&mut self) -> hyper::Response<ResBody> {
256        let mut res = hyper::Response::new(std::mem::take(&mut self.body));
257        *res.extensions_mut() = std::mem::take(&mut self.extensions);
258        *res.headers_mut() = std::mem::take(&mut self.headers);
259        if let Some(status) = self.status_code {
260            // Default to a 404 if no response code was set
261            *res.status_mut() = status;
262        }
263
264        res
265    }
266
267    /// Merge data from [`hyper::Response`].
268    #[doc(hidden)]
269    #[inline]
270    pub fn merge_hyper<B>(&mut self, hyper_res: hyper::Response<B>)
271    where
272        B: Into<ResBody>,
273    {
274        let (
275            http::response::Parts {
276                status,
277                version,
278                headers,
279                extensions,
280                ..
281            },
282            body,
283        ) = hyper_res.into_parts();
284
285        self.status_code = Some(status);
286        self.version = version;
287        self.headers = headers;
288        self.extensions = extensions;
289        self.body = body.into();
290    }
291
292    cfg_feature! {
293        #![feature = "cookie"]
294        /// Get cookies reference.
295        #[inline]
296        pub fn cookies(&self) -> &CookieJar {
297            &self.cookies
298        }
299        /// Get mutable cookies reference.
300        #[inline]
301        pub fn cookies_mut(&mut self) -> &mut CookieJar {
302            &mut self.cookies
303        }
304        /// Helper function for get cookie.
305        #[inline]
306        pub fn cookie<T>(&self, name: T) -> Option<&Cookie<'static>>
307        where
308            T: AsRef<str>,
309        {
310            self.cookies.get(name.as_ref())
311        }
312        /// Helper function for add cookie.
313        #[inline]
314        pub fn add_cookie(&mut self, cookie: Cookie<'static>)-> &mut Self {
315            self.cookies.add(cookie);
316            self
317        }
318
319        /// Helper function for remove cookie.
320        ///
321        /// Removes `cookie` from this [`CookieJar`]. If an _original_ cookie with the same
322        /// name as `cookie` is present in the jar, a _removal_ cookie will be
323        /// present in the `delta` computation. **To properly generate the removal
324        /// cookie, `cookie` must contain the same `path` and `domain` as the cookie
325        /// that was initially set.**
326        ///
327        /// A "removal" cookie is a cookie that has the same name as the original
328        /// cookie but has an empty value, a max-age of 0, and an expiration date
329        /// far in the past.
330        ///
331        /// Read more about [removal cookies](https://docs.rs/cookie/0.18.0/cookie/struct.CookieJar.html#method.remove).
332        #[inline]
333        pub fn remove_cookie(&mut self, name: &str) -> &mut Self
334        {
335            if let Some(cookie) = self.cookies.get(name).cloned() {
336                self.cookies.remove(cookie);
337            }
338            self
339        }
340    }
341
342    /// Get content type..
343    ///
344    /// # Example
345    ///
346    /// ```
347    /// use salvo_core::http::{Response, StatusCode};
348    ///
349    /// let mut res = Response::new();
350    /// assert_eq!(None, res.content_type());
351    /// res.headers_mut().insert("content-type", "text/plain".parse().unwrap());
352    /// assert_eq!(Some(mime::TEXT_PLAIN), res.content_type());
353    ///
354    #[inline]
355    pub fn content_type(&self) -> Option<Mime> {
356        self.headers
357            .get("content-type")
358            .and_then(|h| h.to_str().ok())
359            .and_then(|v| v.parse().ok())
360    }
361
362    /// Sets status code and returns `&mut Self`.
363    ///
364    /// # Example
365    ///
366    /// ```
367    /// use salvo_core::http::StatusCode;
368    /// use salvo_core::http::response::Response;
369    ///
370    /// let mut res = Response::new();
371    /// res.status_code(StatusCode::OK);
372    /// ```
373    #[inline]
374    pub fn status_code(&mut self, code: StatusCode) -> &mut Self {
375        self.status_code = Some(code);
376        self
377    }
378
379    /// Render content.
380    ///
381    /// # Example
382    ///
383    /// ```
384    /// use salvo_core::http::{Response, StatusCode};
385    ///
386    /// let mut res = Response::new();
387    /// res.render("hello world");
388    /// ```
389    pub fn render<P>(&mut self, scribe: P)
390    where
391        P: Scribe,
392    {
393        scribe.render(self);
394    }
395
396    /// Render content with status code.
397    #[inline]
398    pub fn stuff<P>(&mut self, code: StatusCode, scribe: P)
399    where
400        P: Scribe,
401    {
402        self.status_code = Some(code);
403        scribe.render(self);
404    }
405
406    /// Attempts to send a file. If file not exists, not found error will occur.
407    ///
408    /// If you want more settings, you can use `NamedFile::builder` to create a new [`NamedFileBuilder`](crate::fs::NamedFileBuilder).
409    pub async fn send_file<P>(&mut self, path: P, req_headers: &HeaderMap)
410    where
411        P: Into<PathBuf> + Send,
412    {
413        let path = path.into();
414        if !path.exists() {
415            self.render(StatusError::not_found());
416        } else {
417            match NamedFile::builder(path).build().await {
418                Ok(file) => file.send(req_headers, self).await,
419                Err(_) => self.render(StatusError::internal_server_error()),
420            }
421        }
422    }
423
424    /// Write bytes data to body. If body is none, a new `ResBody` will created.
425    pub fn write_body(&mut self, data: impl Into<Bytes>) -> crate::Result<()> {
426        match self.body_mut() {
427            ResBody::Once(bytes) => {
428                let mut chunks = VecDeque::new();
429                chunks.push_back(bytes.clone());
430                chunks.push_back(data.into());
431                self.body = ResBody::Chunks(chunks);
432            }
433            ResBody::Chunks(chunks) => {
434                chunks.push_back(data.into());
435            }
436            ResBody::Hyper(_) => {
437                tracing::error!(
438                    "current body's kind is `ResBody::Hyper`, it is not allowed to write bytes"
439                );
440                return Err(Error::other(
441                    "current body's kind is `ResBody::Hyper`, it is not allowed to write bytes",
442                ));
443            }
444            ResBody::Boxed(_) => {
445                tracing::error!(
446                    "current body's kind is `ResBody::Boxed`, it is not allowed to write bytes"
447                );
448                return Err(Error::other(
449                    "current body's kind is `ResBody::Boxed`, it is not allowed to write bytes",
450                ));
451            }
452            ResBody::Stream(_) => {
453                tracing::error!(
454                    "current body's kind is `ResBody::Stream`, it is not allowed to write bytes"
455                );
456                return Err(Error::other(
457                    "current body's kind is `ResBody::Stream`, it is not allowed to write bytes",
458                ));
459            }
460            ResBody::Channel { .. } => {
461                tracing::error!(
462                    "current body's kind is `ResBody::Channel`, it is not allowed to write bytes"
463                );
464                return Err(Error::other(
465                    "current body's kind is `ResBody::Channel`, it is not allowed to write bytes",
466                ));
467            }
468            ResBody::None | ResBody::Error(_) => {
469                self.body = ResBody::Once(data.into());
470            }
471        }
472        Ok(())
473    }
474
475    /// Set response's body to stream.
476    #[inline]
477    pub fn stream<S, O, E>(&mut self, stream: S)
478    where
479        S: Stream<Item = Result<O, E>> + Send + 'static,
480        O: Into<BytesFrame> + 'static,
481        E: Into<BoxedError> + 'static,
482    {
483        self.body = ResBody::stream(stream);
484    }
485
486    /// Create a `Body` stream with an associated sender half.
487    ///
488    /// Useful when wanting to stream chunks from another thread.
489    ///
490    /// # Example
491    ///
492    /// ```
493    /// use salvo_core::prelude::*;
494    /// #[handler]
495    /// async fn hello(res: &mut Response) {
496    ///     res.add_header("content-type", "text/plain", true).unwrap();
497    ///     let mut tx = res.channel();
498    ///     tokio::spawn(async move {
499    ///         tx.send_data("Hello world").await.unwrap();
500    ///     });
501    /// }
502    /// ```
503    #[inline]
504    pub fn channel(&mut self) -> BodySender {
505        let (sender, body) = ResBody::channel();
506        self.body = body;
507        sender
508    }
509}
510
511impl Debug for Response {
512    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
513        f.debug_struct("Response")
514            .field("status_code", &self.status_code)
515            .field("version", &self.version)
516            .field("headers", &self.headers)
517            // omits Extensions because not useful
518            .field("body", &self.body)
519            .finish()
520    }
521}
522
523impl Display for Response {
524    #[inline]
525    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
526        fmt::Debug::fmt(self, f)
527    }
528}
529
530#[cfg(test)]
531mod test {
532    use bytes::BytesMut;
533    use futures_util::stream::{StreamExt, iter};
534    use std::error::Error;
535
536    use super::*;
537
538    #[test]
539    fn test_body_empty() {
540        let body = ResBody::Once(Bytes::from("hello"));
541        assert!(!body.is_none());
542        let body = ResBody::None;
543        assert!(body.is_none());
544    }
545
546    #[tokio::test]
547    async fn test_body_stream1() {
548        let mut body = ResBody::Once(Bytes::from("hello"));
549
550        let mut result = bytes::BytesMut::new();
551        while let Some(Ok(data)) = body.next().await {
552            result.extend_from_slice(&data.into_data().unwrap_or_default())
553        }
554
555        assert_eq!("hello", &result)
556    }
557
558    #[tokio::test]
559    async fn test_body_stream2() {
560        let mut body = ResBody::stream(iter(vec![
561            Result::<_, Box<dyn Error + Send + Sync>>::Ok(BytesMut::from("Hello").freeze()),
562            Result::<_, Box<dyn Error + Send + Sync>>::Ok(BytesMut::from(" World").freeze()),
563        ]));
564
565        let mut result = bytes::BytesMut::new();
566        while let Some(Ok(data)) = body.next().await {
567            result.extend_from_slice(&data.into_data().unwrap_or_default())
568        }
569
570        assert_eq!("Hello World", &result)
571    }
572}
salvo_core/http/response.rs

salvo_core/http/
response.rs
