fix wrong insertion/appending of headers

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