1
0
Fork 0
mirror of https://github.com/actix/actix-web.git synced 2024-12-21 23:56:35 +00:00
actix-web/awc/src/response.rs

466 lines
13 KiB
Rust
Raw Normal View History

use std::fmt;
2020-05-18 02:46:57 +00:00
use std::future::Future;
2019-04-02 17:53:44 +00:00
use std::marker::PhantomData;
use std::pin::Pin;
use std::task::{Context, Poll};
2021-02-13 15:08:43 +00:00
use std::{
cell::{Ref, RefMut},
mem,
};
2019-03-27 03:45:00 +00:00
use bytes::{Bytes, BytesMut};
2020-05-18 02:46:57 +00:00
use futures_core::{ready, Stream};
2021-02-13 15:08:43 +00:00
use actix_http::error::PayloadError;
use actix_http::http::header;
use actix_http::http::{HeaderMap, StatusCode, Version};
2019-03-27 06:25:24 +00:00
use actix_http::{Extensions, HttpMessage, Payload, PayloadStream, ResponseHead};
2019-04-01 18:51:18 +00:00
use serde::de::DeserializeOwned;
2021-02-13 15:08:43 +00:00
#[cfg(feature = "cookies")]
use actix_http::{cookie::Cookie, error::CookieParseError};
2019-04-01 18:51:18 +00:00
use crate::error::JsonPayloadError;
2019-03-27 04:31:18 +00:00
/// Client Response
2019-03-27 03:45:00 +00:00
pub struct ClientResponse<S = PayloadStream> {
pub(crate) head: ResponseHead,
2019-03-27 03:45:00 +00:00
pub(crate) payload: Payload<S>,
}
2019-03-27 03:45:00 +00:00
impl<S> HttpMessage for ClientResponse<S> {
type Stream = S;
fn headers(&self) -> &HeaderMap {
&self.head.headers
}
2019-12-08 06:31:16 +00:00
fn extensions(&self) -> Ref<'_, Extensions> {
self.head.extensions()
}
2019-12-08 06:31:16 +00:00
fn extensions_mut(&self) -> RefMut<'_, Extensions> {
self.head.extensions_mut()
}
2019-03-27 03:45:00 +00:00
fn take_payload(&mut self) -> Payload<S> {
2021-02-13 15:08:43 +00:00
mem::replace(&mut self.payload, Payload::None)
}
2019-03-27 04:31:18 +00:00
/// Load request cookies.
2021-02-13 15:08:43 +00:00
#[cfg(feature = "cookies")]
2019-12-08 06:31:16 +00:00
fn cookies(&self) -> Result<Ref<'_, Vec<Cookie<'static>>>, CookieParseError> {
2019-03-27 04:31:18 +00:00
struct Cookies(Vec<Cookie<'static>>);
if self.extensions().get::<Cookies>().is_none() {
let mut cookies = Vec::new();
2021-02-13 15:08:43 +00:00
for hdr in self.headers().get_all(&header::SET_COOKIE) {
2021-02-11 23:03:17 +00:00
let s = std::str::from_utf8(hdr.as_bytes()).map_err(CookieParseError::from)?;
2019-03-27 04:31:18 +00:00
cookies.push(Cookie::parse_encoded(s)?.into_owned());
}
self.extensions_mut().insert(Cookies(cookies));
}
Ok(Ref::map(self.extensions(), |ext| {
&ext.get::<Cookies>().unwrap().0
}))
}
}
2019-03-27 03:45:00 +00:00
impl<S> ClientResponse<S> {
/// Create new Request instance
2019-03-27 03:45:00 +00:00
pub(crate) fn new(head: ResponseHead, payload: Payload<S>) -> Self {
ClientResponse { head, payload }
}
#[inline]
pub(crate) fn head(&self) -> &ResponseHead {
&self.head
}
/// Read the Request Version.
#[inline]
pub fn version(&self) -> Version {
self.head().version
}
/// Get the status from the server.
#[inline]
pub fn status(&self) -> StatusCode {
self.head().status
}
#[inline]
2019-04-02 20:35:01 +00:00
/// Returns request's headers.
pub fn headers(&self) -> &HeaderMap {
&self.head().headers
}
2019-03-27 03:45:00 +00:00
/// Set a body and return previous body value
pub fn map_body<F, U>(mut self, f: F) -> ClientResponse<U>
where
F: FnOnce(&mut ResponseHead, Payload<S>) -> Payload<U>,
{
let payload = f(&mut self.head, self.payload);
ClientResponse {
payload,
head: self.head,
}
}
}
2019-03-27 03:45:00 +00:00
impl<S> ClientResponse<S>
where
S: Stream<Item = Result<Bytes, PayloadError>>,
2019-03-27 03:45:00 +00:00
{
2021-02-11 22:39:54 +00:00
/// Loads HTTP response's body.
2019-04-01 18:29:26 +00:00
pub fn body(&mut self) -> MessageBody<S> {
2019-03-27 03:45:00 +00:00
MessageBody::new(self)
}
2019-04-01 18:51:18 +00:00
/// Loads and parse `application/json` encoded body.
/// Return `JsonBody<T>` future. It resolves to a `T` value.
///
/// Returns error:
///
/// * content type is not `application/json`
/// * content length is greater than 256k
pub fn json<T: DeserializeOwned>(&mut self) -> JsonBody<S, T> {
JsonBody::new(self)
}
2019-03-27 03:45:00 +00:00
}
impl<S> Stream for ClientResponse<S>
where
S: Stream<Item = Result<Bytes, PayloadError>> + Unpin,
2019-03-27 03:45:00 +00:00
{
type Item = Result<Bytes, PayloadError>;
2021-02-11 23:03:17 +00:00
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
Pin::new(&mut self.get_mut().payload).poll_next(cx)
}
}
2019-03-27 03:45:00 +00:00
impl<S> fmt::Debug for ClientResponse<S> {
2019-12-08 06:31:16 +00:00
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
writeln!(f, "\nClientResponse {:?} {}", self.version(), self.status(),)?;
writeln!(f, " headers:")?;
for (key, val) in self.headers().iter() {
writeln!(f, " {:?}: {:?}", key, val)?;
}
Ok(())
}
}
2019-03-27 03:45:00 +00:00
2021-02-11 22:39:54 +00:00
/// Future that resolves to a complete HTTP message body.
2019-03-27 03:45:00 +00:00
pub struct MessageBody<S> {
length: Option<usize>,
err: Option<PayloadError>,
2019-04-02 17:53:44 +00:00
fut: Option<ReadBody<S>>,
2019-03-27 03:45:00 +00:00
}
impl<S> MessageBody<S>
where
S: Stream<Item = Result<Bytes, PayloadError>>,
2019-03-27 03:45:00 +00:00
{
/// Create `MessageBody` for request.
2019-04-01 18:29:26 +00:00
pub fn new(res: &mut ClientResponse<S>) -> MessageBody<S> {
2019-03-27 03:45:00 +00:00
let mut len = None;
2021-02-13 15:08:43 +00:00
if let Some(l) = res.headers().get(&header::CONTENT_LENGTH) {
2019-03-27 03:45:00 +00:00
if let Ok(s) = l.to_str() {
if let Ok(l) = s.parse::<usize>() {
len = Some(l)
} else {
return Self::err(PayloadError::UnknownLength);
}
} else {
return Self::err(PayloadError::UnknownLength);
}
}
MessageBody {
length: len,
err: None,
2019-04-02 17:53:44 +00:00
fut: Some(ReadBody::new(res.take_payload(), 262_144)),
2019-03-27 03:45:00 +00:00
}
}
/// Change max size of payload. By default max size is 256kB
2019-03-27 03:45:00 +00:00
pub fn limit(mut self, limit: usize) -> Self {
2019-04-02 17:53:44 +00:00
if let Some(ref mut fut) = self.fut {
fut.limit = limit;
}
2019-03-27 03:45:00 +00:00
self
}
fn err(e: PayloadError) -> Self {
MessageBody {
fut: None,
err: Some(e),
length: None,
}
}
}
impl<S> Future for MessageBody<S>
where
S: Stream<Item = Result<Bytes, PayloadError>> + Unpin,
2019-03-27 03:45:00 +00:00
{
type Output = Result<Bytes, PayloadError>;
2019-03-27 03:45:00 +00:00
2019-12-08 06:31:16 +00:00
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
if let Some(err) = this.err.take() {
return Poll::Ready(Err(err));
2019-03-27 03:45:00 +00:00
}
if let Some(len) = this.length.take() {
if len > this.fut.as_ref().unwrap().limit {
return Poll::Ready(Err(PayloadError::Overflow));
2019-03-27 03:45:00 +00:00
}
}
Pin::new(&mut this.fut.as_mut().unwrap()).poll(cx)
2019-03-27 03:45:00 +00:00
}
}
2019-03-27 04:54:57 +00:00
2019-04-01 18:51:18 +00:00
/// Response's payload json parser, it resolves to a deserialized `T` value.
///
/// Returns error:
///
/// * content type is not `application/json`
/// * content length is greater than 64k
pub struct JsonBody<S, U> {
length: Option<usize>,
err: Option<JsonPayloadError>,
2019-04-02 17:53:44 +00:00
fut: Option<ReadBody<S>>,
2021-01-04 00:49:02 +00:00
_phantom: PhantomData<U>,
2019-04-01 18:51:18 +00:00
}
impl<S, U> JsonBody<S, U>
where
S: Stream<Item = Result<Bytes, PayloadError>>,
2019-04-01 18:51:18 +00:00
U: DeserializeOwned,
{
/// Create `JsonBody` for request.
pub fn new(req: &mut ClientResponse<S>) -> Self {
// check content-type
let json = if let Ok(Some(mime)) = req.mime_type() {
mime.subtype() == mime::JSON || mime.suffix() == Some(mime::JSON)
} else {
false
};
if !json {
return JsonBody {
length: None,
fut: None,
err: Some(JsonPayloadError::ContentType),
2021-01-04 00:49:02 +00:00
_phantom: PhantomData,
2019-04-01 18:51:18 +00:00
};
}
let mut len = None;
2021-02-13 15:08:43 +00:00
if let Some(l) = req.headers().get(&header::CONTENT_LENGTH) {
2019-04-01 18:51:18 +00:00
if let Ok(s) = l.to_str() {
if let Ok(l) = s.parse::<usize>() {
len = Some(l)
}
}
}
JsonBody {
length: len,
err: None,
2019-04-02 17:53:44 +00:00
fut: Some(ReadBody::new(req.take_payload(), 65536)),
2021-01-04 00:49:02 +00:00
_phantom: PhantomData,
2019-04-01 18:51:18 +00:00
}
}
/// Change max size of payload. By default max size is 64kB
2019-04-01 18:51:18 +00:00
pub fn limit(mut self, limit: usize) -> Self {
2019-04-02 17:53:44 +00:00
if let Some(ref mut fut) = self.fut {
fut.limit = limit;
}
2019-04-01 18:51:18 +00:00
self
}
}
impl<T, U> Unpin for JsonBody<T, U>
where
T: Stream<Item = Result<Bytes, PayloadError>> + Unpin,
U: DeserializeOwned,
{
}
2019-04-01 18:51:18 +00:00
impl<T, U> Future for JsonBody<T, U>
where
T: Stream<Item = Result<Bytes, PayloadError>> + Unpin,
2019-04-04 17:59:34 +00:00
U: DeserializeOwned,
2019-04-01 18:51:18 +00:00
{
type Output = Result<U, JsonPayloadError>;
2019-04-01 18:51:18 +00:00
2019-12-08 06:31:16 +00:00
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
2019-04-01 18:51:18 +00:00
if let Some(err) = self.err.take() {
return Poll::Ready(Err(err));
2019-04-01 18:51:18 +00:00
}
if let Some(len) = self.length.take() {
2019-04-02 17:53:44 +00:00
if len > self.fut.as_ref().unwrap().limit {
2021-02-11 23:03:17 +00:00
return Poll::Ready(Err(JsonPayloadError::Payload(PayloadError::Overflow)));
2019-04-01 18:51:18 +00:00
}
}
let body = ready!(Pin::new(&mut self.get_mut().fut.as_mut().unwrap()).poll(cx))?;
Poll::Ready(serde_json::from_slice::<U>(&body).map_err(JsonPayloadError::from))
2019-04-02 17:53:44 +00:00
}
}
struct ReadBody<S> {
stream: Payload<S>,
buf: BytesMut,
limit: usize,
}
impl<S> ReadBody<S> {
fn new(stream: Payload<S>, limit: usize) -> Self {
Self {
stream,
buf: BytesMut::with_capacity(std::cmp::min(limit, 32768)),
limit,
}
}
}
impl<S> Future for ReadBody<S>
where
S: Stream<Item = Result<Bytes, PayloadError>> + Unpin,
2019-04-02 17:53:44 +00:00
{
type Output = Result<Bytes, PayloadError>;
2019-12-08 06:31:16 +00:00
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
2019-04-02 17:53:44 +00:00
loop {
return match Pin::new(&mut this.stream).poll_next(cx)? {
Poll::Ready(Some(chunk)) => {
if (this.buf.len() + chunk.len()) > this.limit {
Poll::Ready(Err(PayloadError::Overflow))
2019-04-02 17:53:44 +00:00
} else {
this.buf.extend_from_slice(&chunk);
2019-04-02 17:53:44 +00:00
continue;
}
2019-04-01 18:51:18 +00:00
}
2019-12-05 17:35:43 +00:00
Poll::Ready(None) => Poll::Ready(Ok(this.buf.split().freeze())),
Poll::Pending => Poll::Pending,
2019-04-02 17:53:44 +00:00
};
}
2019-04-01 18:51:18 +00:00
}
}
2019-03-27 04:54:57 +00:00
#[cfg(test)]
mod tests {
use super::*;
2019-04-01 18:51:18 +00:00
use serde::{Deserialize, Serialize};
2019-03-27 04:54:57 +00:00
2019-05-12 15:34:51 +00:00
use crate::{http::header, test::TestResponse};
2019-03-27 04:54:57 +00:00
2019-11-26 05:25:50 +00:00
#[actix_rt::test]
async fn test_body() {
let mut req = TestResponse::with_header(header::CONTENT_LENGTH, "xxxx").finish();
match req.body().await.err().unwrap() {
2021-01-04 01:01:35 +00:00
PayloadError::UnknownLength => {}
2019-11-26 05:25:50 +00:00
_ => unreachable!("error"),
}
2019-03-27 04:54:57 +00:00
2021-02-11 23:03:17 +00:00
let mut req = TestResponse::with_header(header::CONTENT_LENGTH, "1000000").finish();
2019-11-26 05:25:50 +00:00
match req.body().await.err().unwrap() {
2021-01-04 01:01:35 +00:00
PayloadError::Overflow => {}
2019-11-26 05:25:50 +00:00
_ => unreachable!("error"),
}
2019-03-27 04:54:57 +00:00
2019-11-26 05:25:50 +00:00
let mut req = TestResponse::default()
.set_payload(Bytes::from_static(b"test"))
.finish();
assert_eq!(req.body().await.ok().unwrap(), Bytes::from_static(b"test"));
let mut req = TestResponse::default()
.set_payload(Bytes::from_static(b"11111111111111"))
.finish();
match req.body().limit(5).await.err().unwrap() {
2021-01-04 01:01:35 +00:00
PayloadError::Overflow => {}
2019-11-26 05:25:50 +00:00
_ => unreachable!("error"),
}
2019-03-27 04:54:57 +00:00
}
2019-04-01 18:51:18 +00:00
#[derive(Serialize, Deserialize, PartialEq, Debug)]
struct MyObject {
name: String,
}
fn json_eq(err: JsonPayloadError, other: JsonPayloadError) -> bool {
match err {
JsonPayloadError::Payload(PayloadError::Overflow) => {
matches!(other, JsonPayloadError::Payload(PayloadError::Overflow))
}
JsonPayloadError::ContentType => {
matches!(other, JsonPayloadError::ContentType)
}
2019-04-01 18:51:18 +00:00
_ => false,
}
}
2019-11-26 05:25:50 +00:00
#[actix_rt::test]
async fn test_json_body() {
let mut req = TestResponse::default().finish();
let json = JsonBody::<_, MyObject>::new(&mut req).await;
assert!(json_eq(json.err().unwrap(), JsonPayloadError::ContentType));
let mut req = TestResponse::default()
.header(
header::CONTENT_TYPE,
header::HeaderValue::from_static("application/text"),
)
.finish();
let json = JsonBody::<_, MyObject>::new(&mut req).await;
assert!(json_eq(json.err().unwrap(), JsonPayloadError::ContentType));
let mut req = TestResponse::default()
.header(
header::CONTENT_TYPE,
header::HeaderValue::from_static("application/json"),
)
.header(
header::CONTENT_LENGTH,
header::HeaderValue::from_static("10000"),
)
.finish();
let json = JsonBody::<_, MyObject>::new(&mut req).limit(100).await;
assert!(json_eq(
json.err().unwrap(),
JsonPayloadError::Payload(PayloadError::Overflow)
));
let mut req = TestResponse::default()
.header(
header::CONTENT_TYPE,
header::HeaderValue::from_static("application/json"),
)
.header(
header::CONTENT_LENGTH,
header::HeaderValue::from_static("16"),
)
.set_payload(Bytes::from_static(b"{\"name\": \"test\"}"))
.finish();
let json = JsonBody::<_, MyObject>::new(&mut req).await;
assert_eq!(
json.ok().unwrap(),
MyObject {
name: "test".to_owned()
}
);
2019-04-01 18:51:18 +00:00
}
2019-03-27 04:54:57 +00:00
}