diff options
author | Martin Fischer <martin@push-f.com> | 2021-01-26 10:56:58 +0100 |
---|---|---|
committer | Martin Fischer <martin@push-f.com> | 2021-01-26 11:01:52 +0100 |
commit | 18b1b875c80a244f6ea29894672db18b90ec4eea (patch) | |
tree | 03bd90c620b4f17d59fcbc3a4843d00adf4368bb /src | |
parent | 3fd065757f02ebf1d055912e6809fac15c8bf058 (diff) |
fix wrong insertion/appending of headers
Diffstat (limited to 'src')
-rw-r--r-- | src/response.rs | 53 |
1 files changed, 47 insertions, 6 deletions
diff --git a/src/response.rs b/src/response.rs index 9c92e4c..78e9a70 100644 --- a/src/response.rs +++ b/src/response.rs @@ -1,15 +1,17 @@ //! Provides the [`SputnikBuilder`] trait. +use std::convert::TryInto; + use cookie::Cookie; -use hyper::{StatusCode, header::{self, HeaderValue}, http}; +use hyper::{StatusCode, header, http}; use time::{Duration, OffsetDateTime}; use hyper::http::response::Builder; pub trait SputnikBuilder { - /// Adds a Set-Cookie header. + /// Appends a Set-Cookie header. fn set_cookie(&mut self, cookie: Cookie); - /// Adds a Set-Cookie header to delete a cookie. + /// Appends a Set-Cookie header to delete a cookie. fn delete_cookie(&mut self, name: &str); /// Sets the Content-Type. @@ -23,7 +25,9 @@ pub fn redirect(location: &str, code: StatusCode) -> Builder { impl SputnikBuilder for Builder { fn set_cookie(&mut self, cookie: Cookie) { - self.headers_mut().unwrap().insert(header::SET_COOKIE, HeaderValue::from_str(&cookie.encoded().to_string()).unwrap()); + if let Some(headers) = self.headers_mut() { + headers.append(header::SET_COOKIE, cookie.encoded().to_string().try_into().unwrap()); + } } fn delete_cookie(&mut self, name: &str) { @@ -33,8 +37,11 @@ impl SputnikBuilder for Builder { self.set_cookie(cookie); } - fn content_type(self, mime: mime::Mime) -> Self { - self.header(header::CONTENT_TYPE, mime.to_string()) + fn content_type(mut self, mime: mime::Mime) -> Self { + if let Some(headers) = self.headers_mut() { + headers.insert(header::CONTENT_TYPE, mime.to_string().try_into().unwrap()); + } + self } } @@ -47,4 +54,38 @@ impl EmptyBuilder<hyper::Body> for Builder { fn empty(self) -> http::Result<http::response::Response<hyper::Body>> { self.body(hyper::Body::empty()) } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_set_cookie() { + let mut builder = Builder::new(); + builder.set_cookie(Cookie::new("some", "cookie")); + builder.set_cookie(Cookie::new("some", "cookie")); + let resp = builder.body(hyper::Body::empty()).unwrap(); + assert_eq!(resp.headers().len(), 2); + + let mut builder = Builder::new() + .header("foo", "invalid\r\n"); + // doesn't panic after invalid header + builder.set_cookie(Cookie::new("some", "cookie")); + } + + #[test] + fn test_content_type() { + let resp = Builder::new() + .content_type(mime::TEXT_PLAIN) + .content_type(mime::TEXT_HTML) + .body(hyper::Body::empty()).unwrap(); + assert_eq!(resp.headers().len(), 1); + assert_eq!(resp.headers().get(header::CONTENT_TYPE).unwrap(), "text/html"); + + // doesn't panic after invalid header + Builder::new() + .header("foo", "invalid\r\n") + .content_type(mime::TEXT_HTML); + } }
\ No newline at end of file |