diff --git a/core/src/http/request.rs b/core/src/http/request.rs index cf07244e2..45bf3b640 100644 --- a/core/src/http/request.rs +++ b/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::().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(&mut self, name: N, value: V, overwrite: bool) -> crate::Result<&mut Self> + where + N: IntoHeaderName, + V: TryInto, + { + 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(&mut self, name: N, value: V, overwrite: bool) -> crate::Result<()> + where + N: IntoHeaderName, + V: TryInto, + { + 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,7 +743,7 @@ 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::("money").await.unwrap(), "sh*t"); @@ -718,9 +751,10 @@ mod tests { assert_eq!(req.form_or_query::("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\ diff --git a/core/src/http/response.rs b/core/src/http/response.rs index fe5812930..991796716 100644 --- a/core/src/http/response.rs +++ b/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(&mut self, name: N, value: V, overwrite: bool) -> crate::Result<&mut Self> + where + N: IntoHeaderName, + V: TryInto, + { + 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(&mut self, name: N, value: V, overwrite: bool) -> crate::Result<()> + where + N: IntoHeaderName, + V: TryInto, + { + 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(&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 { @@ -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]