diff --git a/src/body.rs b/src/body.rs index 8f5d0979..b3b19350 100644 --- a/src/body.rs +++ b/src/body.rs @@ -52,7 +52,7 @@ pin_project_lite::pin_project! { pub struct Body { #[pin] reader: Box, - mime: Option, + mime: Mime, length: Option, } } @@ -74,7 +74,7 @@ impl Body { pub fn empty() -> Self { Self { reader: Box::new(io::empty()), - mime: Some(mime::BYTE_STREAM), + mime: mime::BYTE_STREAM, length: Some(0), } } @@ -104,7 +104,7 @@ impl Body { ) -> Self { Self { reader: Box::new(reader), - mime: Some(mime::BYTE_STREAM), + mime: mime::BYTE_STREAM, length: len, } } @@ -147,9 +147,8 @@ impl Body { self.reader } - /// Return the mime type. - pub(crate) fn take_mime(&mut self) -> Mime { - self.mime.take().expect("internal error: missing mime") + pub(crate) fn mime(&self) -> &Mime { + &self.mime } } @@ -167,7 +166,7 @@ impl From for Body { Self { length: Some(s.len()), reader: Box::new(io::Cursor::new(s.into_bytes())), - mime: Some(string_mime()), + mime: string_mime(), } } } @@ -177,7 +176,7 @@ impl<'a> From<&'a str> for Body { Self { length: Some(s.len()), reader: Box::new(io::Cursor::new(s.to_owned().into_bytes())), - mime: Some(string_mime()), + mime: string_mime(), } } } @@ -195,7 +194,7 @@ impl From> for Body { Self { length: Some(b.len()), reader: Box::new(io::Cursor::new(b)), - mime: Some(mime::BYTE_STREAM), + mime: mime::BYTE_STREAM, } } } diff --git a/src/request.rs b/src/request.rs index 044ae942..ac229599 100644 --- a/src/request.rs +++ b/src/request.rs @@ -111,11 +111,7 @@ impl Request { /// req.set_body("Hello, Nori!"); /// ``` pub fn set_body(&mut self, body: impl Into) { - self.body = body.into(); - if self.header(&CONTENT_TYPE).is_none() { - let mime = self.body.take_mime(); - self.set_content_type(mime); - } + self.replace_body(body); } /// Swaps the value of the body with another body, without deinitializing @@ -141,7 +137,9 @@ impl Request { /// # Ok(()) }) } /// ``` pub fn replace_body(&mut self, body: impl Into) -> Body { - mem::replace(&mut self.body, body.into()) + let body = mem::replace(&mut self.body, body.into()); + self.copy_content_type_from_body(); + body } /// Replace the request body with a new body, and return the old body. @@ -168,6 +166,7 @@ impl Request { /// ``` pub fn swap_body(&mut self, body: &mut Body) { mem::swap(&mut self.body, body); + self.copy_content_type_from_body(); } /// Take the request body, replacing it with an empty body. @@ -270,6 +269,13 @@ impl Request { self.insert_header(CONTENT_TYPE, value).unwrap() } + /// Copy MIME data from the body. + fn copy_content_type_from_body(&mut self) { + if self.header(&CONTENT_TYPE).is_none() { + self.set_content_type(self.body.mime().clone()); + } + } + /// Get the current content type pub fn content_type(&self) -> Option { self.header(&CONTENT_TYPE)?.last()?.as_str().parse().ok() @@ -446,7 +452,7 @@ impl Request { /// assert_eq!(req.local().get(), Some(&"hello from the extension")); /// # /// # Ok(()) } - /// ``` + /// ``` pub fn local_mut(&mut self) -> &mut TypeMap { &mut self.local } diff --git a/src/response.rs b/src/response.rs index 25f659cb..a2a19941 100644 --- a/src/response.rs +++ b/src/response.rs @@ -138,11 +138,7 @@ impl Response { /// # Ok(()) } /// ``` pub fn set_body(&mut self, body: impl Into) { - self.body = body.into(); - if self.header(&CONTENT_TYPE).is_none() { - let mime = self.body.take_mime(); - self.set_content_type(mime); - } + self.replace_body(body); } /// Replace the response body with a new body, returning the old body. @@ -168,7 +164,9 @@ impl Response { /// # Ok(()) }) } /// ``` pub fn replace_body(&mut self, body: impl Into) -> Body { - mem::replace(&mut self.body, body.into()) + let body = mem::replace(&mut self.body, body.into()); + self.copy_content_type_from_body(); + body } /// Swaps the value of the body with another body, without deinitializing @@ -197,6 +195,7 @@ impl Response { /// ``` pub fn swap_body(&mut self, body: &mut Body) { mem::swap(&mut self.body, body); + self.copy_content_type_from_body(); } /// Take the response body, replacing it with an empty body. @@ -236,6 +235,13 @@ impl Response { self.insert_header(CONTENT_TYPE, value).unwrap() } + /// Copy MIME data from the body. + fn copy_content_type_from_body(&mut self) { + if self.header(&CONTENT_TYPE).is_none() { + self.set_content_type(self.body.mime().clone()); + } + } + /// Get the length of the body stream, if it has been set. /// /// This value is set when passing a fixed-size object into as the body. E.g. a string, or a @@ -412,7 +418,7 @@ impl Response { /// assert_eq!(res.local().get(), Some(&"hello from the extension")); /// # /// # Ok(()) } - /// ``` + /// ``` pub fn local_mut(&mut self) -> &mut TypeMap { &mut self.local } diff --git a/tests/req_res_body.rs b/tests/req_res_body.rs new file mode 100644 index 00000000..df884c27 --- /dev/null +++ b/tests/req_res_body.rs @@ -0,0 +1,30 @@ +use async_std::io::prelude::*; +use http_types::{Body, Method, Request, Response, StatusCode, Url}; + +#[test] +fn test_req_res_set_body() { + let mut req = Request::new(Method::Get, Url::parse("http://example.com/").unwrap()); + req.set_body(Body::empty()); + let mut res = Response::new(StatusCode::Ok); + res.set_body(req); + let body = async_std::task::block_on(async move { + let mut body = Vec::new(); + res.read_to_end(&mut body).await.unwrap(); + body + }); + assert!(body.is_empty()); +} + +#[test] +fn test_req_res_take_replace_body() { + let mut req = Request::new(Method::Get, Url::parse("http://example.com/").unwrap()); + req.take_body(); + let mut res = Response::new(StatusCode::Ok); + res.replace_body(req); + let body = async_std::task::block_on(async move { + let mut body = Vec::new(); + res.read_to_end(&mut body).await.unwrap(); + body + }); + assert!(body.is_empty()); +}