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 = 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 = 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> { + 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")); } }