From d8e1fd102de6c1014f9dad688344da4f9d0b6009 Mon Sep 17 00:00:00 2001
From: axon-q <axon-q@axon-q.cc>
Date: Tue, 12 Jun 2018 13:49:07 +0000
Subject: [PATCH] add cookie methods to HttpResponse

---
 CHANGES.md          |   5 +++
 src/httpresponse.rs | 104 ++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 106 insertions(+), 3 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 29046f8f8..f7a753543 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -4,6 +4,8 @@
 
 ### Added
 
+* Add methods to `HttpResponse` to retrieve, add, and delete cookies
+
 * Add `.set_content_type()` and `.set_content_disposition()` methods
   to `fs::NamedFile` to allow overriding the values inferred by default
 
@@ -22,6 +24,9 @@
 
 * Min rustc version is 1.26
 
+* `HttpResponse::into_builder()` now moves cookies into the builder
+  instead of dropping them
+
 * Use tokio instead of tokio-core
 
 * Use `&mut self` instead of `&self` for Middleware trait
diff --git a/src/httpresponse.rs b/src/httpresponse.rs
index f42be2b9f..40a8cc61e 100644
--- a/src/httpresponse.rs
+++ b/src/httpresponse.rs
@@ -97,14 +97,25 @@ impl HttpResponse {
     /// Convert `HttpResponse` to a `HttpResponseBuilder`
     #[inline]
     pub fn into_builder(mut self) -> HttpResponseBuilder {
+        // If this response has cookies, load them into a jar
+        let mut jar: Option<CookieJar> = None;
+        for c in self.cookies() {
+            if let Some(ref mut j) = jar {
+                j.add_original(c.into_owned());
+            } else {
+                let mut j = CookieJar::new();
+                j.add_original(c.into_owned());
+                jar = Some(j);
+            }
+        }
+
         let response = self.0.take();
         let pool = Some(Rc::clone(&self.1));
-
         HttpResponseBuilder {
             response,
             pool,
             err: None,
-            cookies: None, // TODO: convert set-cookie headers
+            cookies: jar,
         }
     }
 
@@ -132,6 +143,49 @@ impl HttpResponse {
         &mut self.get_mut().headers
     }
 
+    /// Get an iterator for the cookies set by this response
+    #[inline]
+    pub fn cookies(&self) -> CookieIter {
+        CookieIter {
+            iter: self.get_ref().headers.get_all(header::SET_COOKIE).iter()
+        }
+    }
+
+    /// Add a cookie to this response
+    #[inline]
+    pub fn add_cookie(&mut self, cookie: Cookie) -> Result<(), HttpError> {
+        let h = &mut self.get_mut().headers;
+        HeaderValue::from_str(&cookie.to_string())
+            .map(|c| { h.append(header::SET_COOKIE, c); })
+            .map_err(|e| e.into())
+    }
+
+    /// Remove all cookies with the given name from this response. Returns
+    /// the number of cookies removed.
+    #[inline]
+    pub fn del_cookie(&mut self, name: &str) -> usize {
+        let h = &mut self.get_mut().headers;
+        let vals: Vec<HeaderValue> = h.get_all(header::SET_COOKIE)
+            .iter()
+            .map(|v| v.to_owned())
+            .collect();
+        h.remove(header::SET_COOKIE);
+
+        let mut count: usize = 0;
+        for v in vals {
+            if let Ok(s) = v.to_str() {
+                if let Ok(c) = Cookie::parse(s) {
+                    if c.name() == name {
+                        count += 1;
+                        continue;
+                    }
+                }
+            }
+            h.append(header::SET_COOKIE, v);
+        }
+        return count;
+    }
+
     /// Get the response status code
     #[inline]
     pub fn status(&self) -> StatusCode {
@@ -269,6 +323,24 @@ impl fmt::Debug for HttpResponse {
     }
 }
 
+pub struct CookieIter<'a> {
+    iter: header::ValueIter<'a, HeaderValue>,
+}
+
+impl<'a> Iterator for CookieIter<'a> {
+    type Item = Cookie<'a>;
+
+    #[inline]
+    fn next(&mut self) -> Option<Cookie<'a>> {
+        for v in self.iter.by_ref() {
+            if let Some(c) = (|| Cookie::parse(v.to_str().ok()?).ok())() {
+                return Some(c);
+            }
+        }
+        None
+    }
+}
+
 /// An HTTP response builder
 ///
 /// This type can be used to construct an instance of `HttpResponse` through a
@@ -984,6 +1056,27 @@ mod tests {
         );
     }
 
+
+    #[test]
+    fn test_update_response_cookies() {
+        let mut r = HttpResponse::Ok()
+            .cookie(http::Cookie::new("original", "val100"))
+            .finish();
+
+        r.add_cookie(http::Cookie::new("cookie2", "val200")).unwrap();
+        r.add_cookie(http::Cookie::new("cookie2", "val250")).unwrap();
+        r.add_cookie(http::Cookie::new("cookie3", "val300")).unwrap();
+
+        assert_eq!(r.cookies().count(), 4);
+        r.del_cookie("cookie2");
+
+        let mut iter = r.cookies();
+        let v = iter.next().unwrap();
+        assert_eq!((v.name(), v.value()), ("original", "val100"));
+        let v = iter.next().unwrap();
+        assert_eq!((v.name(), v.value()), ("cookie3", "val300"));
+    }
+
     #[test]
     fn test_basic_builder() {
         let resp = HttpResponse::Ok()
@@ -1191,11 +1284,16 @@ mod tests {
 
     #[test]
     fn test_into_builder() {
-        let resp: HttpResponse = "test".into();
+        let mut resp: HttpResponse = "test".into();
         assert_eq!(resp.status(), StatusCode::OK);
 
+        resp.add_cookie(http::Cookie::new("cookie1", "val100")).unwrap();
+
         let mut builder = resp.into_builder();
         let resp = builder.status(StatusCode::BAD_REQUEST).finish();
         assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
+
+        let cookie = resp.cookies().next().unwrap();
+        assert_eq!((cookie.name(), cookie.value()), ("cookie1", "val100"));
     }
 }