1 files changed, 22 insertions, 27 deletions

diff --git a/src/response.rs b/src/response.rs
index 78e9a70..ceb0f61 100644
--- a/src/response.rs
+++ b/src/response.rs
@@ -3,31 +3,39 @@
 use std::convert::TryInto;
 
 use cookie::Cookie;
-use hyper::{StatusCode, header, http};
+use hyper::{HeaderMap, StatusCode, header, http};
 use time::{Duration, OffsetDateTime};
 use hyper::http::response::Builder;
 
 pub trait SputnikBuilder {
-    /// Appends a Set-Cookie header.
-    fn set_cookie(&mut self, cookie: Cookie);
-
-    /// Appends a Set-Cookie header to delete a cookie.
-    fn delete_cookie(&mut self, name: &str);
-
     /// Sets the Content-Type.
     fn content_type(self, mime: mime::Mime) -> Builder;
 }
 
-
 pub fn redirect(location: &str, code: StatusCode) -> Builder {
     Builder::new().status(code).header(header::LOCATION, location)
 }
 
 impl SputnikBuilder for Builder {
-    fn set_cookie(&mut self, cookie: Cookie) {
+    fn content_type(mut self, mime: mime::Mime) -> Self {
         if let Some(headers) = self.headers_mut() {
-            headers.append(header::SET_COOKIE, cookie.encoded().to_string().try_into().unwrap());
+            headers.insert(header::CONTENT_TYPE, mime.to_string().try_into().unwrap());
         }
+        self
+    }
+}
+
+pub trait SputnikHeaders {
+    /// Appends a Set-Cookie header.
+    fn set_cookie(&mut self, cookie: Cookie);
+
+    /// Appends a Set-Cookie header to delete a cookie.
+    fn delete_cookie(&mut self, name: &str);
+}
+
+impl SputnikHeaders for HeaderMap {
+    fn set_cookie(&mut self, cookie: Cookie) {
+        self.append(header::SET_COOKIE, cookie.encoded().to_string().try_into().unwrap());
     }
 
     fn delete_cookie(&mut self, name: &str) {
@@ -36,13 +44,6 @@ impl SputnikBuilder for Builder {
         cookie.set_expires(OffsetDateTime::now_utc() - Duration::days(365));
         self.set_cookie(cookie);
     }
-
-    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
-    }
 }
 
 pub trait EmptyBuilder<B> {
@@ -62,16 +63,10 @@ mod tests {
 
     #[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"));
+        let mut map = HeaderMap::new();
+        map.set_cookie(Cookie::new("some", "cookie"));
+        map.set_cookie(Cookie::new("some", "cookie"));
+        assert_eq!(map.len(), 2);
     }
 
     #[test]