Add with_ functions for chaining support

core/src/http/request.rs

@@ -6,7 +6,7 @@ use std::str::FromStr;
 
 #[cfg(feature = "cookie")]
 use cookie::{Cookie, CookieJar};
-use http::header::{self, HeaderMap};
+use http::header::{self, HeaderMap, HeaderValue, IntoHeaderName};
 use http::method::Method;
 pub use http::request::Parts;
 use http::version::Version;
@@ -19,10 +19,9 @@ use serde::de::Deserialize;
 use crate::addr::SocketAddr;
 use crate::extract::{Extractible, Metadata};
 use crate::http::form::{FilePart, FormData};
-use crate::http::header::HeaderValue;
-use crate::http::Mime;
-use crate::http::ParseError;
+use crate::http::{Mime, ParseError};
 use crate::serde::{from_request, from_str_map, from_str_multi_map, from_str_multi_val, from_str_val};
+use crate::Error;
 
 /// Represents an HTTP request.
 ///
@@ -262,6 +261,40 @@ impl Request {
             .and_then(|s| s.parse::<T>().ok())
     }
 
+    /// Modify a header for this request.
+    ///
+    /// When `overwrite` is set to `true`, If the header is already present, the value will be replaced.
+    /// When `overwrite` is set to `false`, The new header is always appended to the request, even if the header already exists.
+    #[inline]
+    pub fn with_header<N, V>(&mut self, name: N, value: V, overwrite: bool) -> crate::Result<&mut Self>
+    where
+        N: IntoHeaderName,
+        V: TryInto<HeaderValue>,
+    {
+        self.add_header(name, value, overwrite)?;
+        Ok(self)
+    }
+
+    /// Modify a header for this request.
+    ///
+    /// When `overwrite` is set to `true`, If the header is already present, the value will be replaced.
+    /// When `overwrite` is set to `false`, The new header is always appended to the request, even if the header already exists.
+    pub fn add_header<N, V>(&mut self, name: N, value: V, overwrite: bool) -> crate::Result<()>
+    where
+        N: IntoHeaderName,
+        V: TryInto<HeaderValue>,
+    {
+        let value = value
+            .try_into()
+            .map_err(|_| Error::Other("invalid header value".into()))?;
+        if overwrite {
+            self.headers.insert(name, value);
+        } else {
+            self.headers.append(name, value);
+        }
+        Ok(())
+    }
+
     /// Returns a reference to the associated HTTP body.
     ///
     /// # Examples
@@ -710,17 +743,18 @@ mod tests {
     #[tokio::test]
     async fn test_form() {
         let mut req = TestClient::post("http://127.0.0.1:7878/hello?q=rust")
-            .insert_header("content-type", "application/x-www-form-urlencoded")
+            .add_header("content-type", "application/x-www-form-urlencoded", true)
             .raw_form("lover=dog&money=sh*t&q=firefox")
             .build();
         assert_eq!(req.form::<String>("money").await.unwrap(), "sh*t");
         assert_eq!(req.query_or_form::<String>("q").await.unwrap(), "rust");
         assert_eq!(req.form_or_query::<String>("q").await.unwrap(), "firefox");
 
         let mut req: Request = TestClient::post("http://127.0.0.1:7878/hello?q=rust")
-            .insert_header(
+            .add_header(
                 "content-type",
                 "multipart/form-data; boundary=----WebKitFormBoundary0mkL0yrNNupCojyz",
+                true,
             )
             .body(
                 "------WebKitFormBoundary0mkL0yrNNupCojyz\r\n\

core/src/http/response.rs

@@ -1,10 +1,5 @@
 //! Http response.
 
-#[cfg(feature = "cookie")]
-use cookie::{Cookie, CookieJar};
-use futures_util::stream::{Stream, TryStreamExt};
-use http::version::Version;
-use mime::Mime;
 #[cfg(feature = "cookie")]
 use std::borrow::Cow;
 use std::collections::VecDeque;
@@ -13,10 +8,15 @@ use std::fmt::{self, Display, Formatter};
 use std::pin::Pin;
 use std::task::{self, Poll};
 
+#[cfg(feature = "cookie")]
+use cookie::{Cookie, CookieJar};
+use futures_util::stream::{Stream, TryStreamExt};
+use http::header::{self, HeaderMap, IntoHeaderName, HeaderValue};
 pub use http::response::Parts;
+use http::version::Version;
+use mime::Mime;
 
 use super::errors::*;
-use super::header::{self, HeaderMap};
 use crate::http::StatusCode;
 use crate::{Error, Piece};
 use bytes::Bytes;
@@ -171,6 +171,40 @@ impl Response {
         self.headers = headers
     }
 
+    /// Modify a header for this response.
+    ///
+    /// When `overwrite` is set to `true`, If the header is already present, the value will be replaced.
+    /// When `overwrite` is set to `false`, The new header is always appended to the request, even if the header already exists.
+    #[inline]
+    pub fn with_header<N, V>(&mut self, name: N, value: V, overwrite: bool) -> crate::Result<&mut Self>
+    where
+        N: IntoHeaderName,
+        V: TryInto<HeaderValue>,
+    {
+        self.add_header(name, value, overwrite)?;
+        Ok(self)
+    }
+
+    /// Modify a header for this response.
+    ///
+    /// When `overwrite` is set to `true`, If the header is already present, the value will be replaced.
+    /// When `overwrite` is set to `false`, The new header is always appended to the request, even if the header already exists.
+    pub fn add_header<N, V>(&mut self, name: N, value: V, overwrite: bool) -> crate::Result<()>
+    where
+        N: IntoHeaderName,
+        V: TryInto<HeaderValue>,
+    {
+        let value = value
+            .try_into()
+            .map_err(|_| Error::Other("invalid header value".into()))?;
+        if overwrite {
+            self.headers.insert(name, value);
+        } else {
+            self.headers.append(name, value);
+        }
+        Ok(())
+    }
+
     /// Get version.
     #[inline]
     pub fn version(&self) -> Version {
@@ -198,6 +232,13 @@ impl Response {
         self.body = body
     }
 
+    /// Set body.
+    #[inline]
+    pub fn with_body(&mut self, body: Body) -> &mut Self {
+        self.body = body;
+        self
+    }
+
     /// Set body to a new value and returns old value.
     #[inline]
     pub fn replace_body(&mut self, body: Body) -> Body {
@@ -283,6 +324,14 @@ impl Response {
         pub fn add_cookie(&mut self, cookie: Cookie<'static>) {
             self.cookies.add(cookie);
         }
+
+        /// Helper function for add cookie.
+        #[inline]
+        pub fn with_cookie(&mut self, cookie: Cookie<'static>) -> &mut Self {
+            self.add_cookie(cookie);
+            self
+        }
+
         /// Helper function for remove cookie.
         #[inline]
         pub fn remove_cookie<T>(&mut self, name: T)
@@ -308,6 +357,13 @@ impl Response {
         }
     }
 
+    /// Set status code.
+    #[inline]
+    pub fn with_status_code(&mut self, code: StatusCode) -> &mut Self {
+        self.set_status_code(code);
+        self
+    }
+
     /// Get content type.
     #[inline]
     pub fn content_type(&self) -> Option<Mime> {
@@ -328,6 +384,12 @@ impl Response {
         self.status_code = Some(err.code);
         self.status_error = Some(err);
     }
+    /// Set http error.
+    #[inline]
+    pub fn with_status_error(&mut self, err: StatusError) -> &mut Self {
+        self.set_status_error(err);
+        self
+    }
 
     /// Render content.
     #[inline]