From 69d710dbce90e9d00539334c063e6d4f70ba45e3 Mon Sep 17 00:00:00 2001 From: Kornel Date: Wed, 27 Feb 2019 12:52:42 +0000 Subject: [PATCH] Add insert and remove() to response builder (#707) --- src/httpresponse.rs | 97 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 95 insertions(+), 2 deletions(-) diff --git a/src/httpresponse.rs b/src/httpresponse.rs index 168e9bf64..226c847f3 100644 --- a/src/httpresponse.rs +++ b/src/httpresponse.rs @@ -366,7 +366,7 @@ impl HttpResponseBuilder { self } - /// Set a header. + /// Append a header. /// /// ```rust /// # extern crate actix_web; @@ -394,7 +394,7 @@ impl HttpResponseBuilder { self } - /// Set a header. + /// Append a header. /// /// ```rust /// # extern crate actix_web; @@ -426,6 +426,65 @@ impl HttpResponseBuilder { } self } + /// Set or replace a header with a single value. + /// + /// ```rust + /// # extern crate actix_web; + /// use actix_web::{http, HttpRequest, HttpResponse}; + /// + /// fn index(req: HttpRequest) -> HttpResponse { + /// HttpResponse::Ok() + /// .insert("X-TEST", "value") + /// .insert(http::header::CONTENT_TYPE, "application/json") + /// .finish() + /// } + /// fn main() {} + /// ``` + pub fn insert(&mut self, key: K, value: V) -> &mut Self + where + HeaderName: HttpTryFrom, + V: IntoHeaderValue, + { + if let Some(parts) = parts(&mut self.response, &self.err) { + match HeaderName::try_from(key) { + Ok(key) => match value.try_into() { + Ok(value) => { + parts.headers.insert(key, value); + } + Err(e) => self.err = Some(e.into()), + }, + Err(e) => self.err = Some(e.into()), + }; + } + self + } + + /// Remove all instances of a header already set on this `HttpResponseBuilder`. + /// + /// ```rust + /// # extern crate actix_web; + /// use actix_web::{http, HttpRequest, HttpResponse}; + /// + /// fn index(req: HttpRequest) -> HttpResponse { + /// HttpResponse::Ok() + /// .header(http::header::CONTENT_TYPE, "nevermind") // won't be used + /// .remove(http::header::CONTENT_TYPE) + /// .finish() + /// } + /// ``` + pub fn remove(&mut self, key: K) -> &mut Self + where HeaderName: HttpTryFrom + { + if let Some(parts) = parts(&mut self.response, &self.err) { + match HeaderName::try_from(key) { + Ok(key) => { + parts.headers.remove(key); + }, + Err(e) => self.err = Some(e.into()), + }; + } + self + } /// Set the custom reason for the response. #[inline] @@ -1128,6 +1187,40 @@ mod tests { assert_eq!(resp.status(), StatusCode::OK); } + #[test] + fn test_insert() { + let resp = HttpResponse::Ok() + .insert("deleteme", "old value") + .insert("deleteme", "new value") + .finish(); + assert_eq!("new value", resp.headers().get("deleteme").expect("new value")); + } + + #[test] + fn test_remove() { + let resp = HttpResponse::Ok() + .header("deleteme", "value") + .remove("deleteme") + .finish(); + assert!(resp.headers().get("deleteme").is_none()) + } + + #[test] + fn test_remove_replace() { + let resp = HttpResponse::Ok() + .header("some-header", "old_value1") + .header("some-header", "old_value2") + .remove("some-header") + .header("some-header", "new_value1") + .header("some-header", "new_value2") + .remove("unrelated-header") + .finish(); + let mut v = resp.headers().get_all("some-header").into_iter(); + assert_eq!("new_value1", v.next().unwrap()); + assert_eq!("new_value2", v.next().unwrap()); + assert_eq!(None, v.next()); + } + #[test] fn test_upgrade() { let resp = HttpResponse::build(StatusCode::OK).upgrade().finish();