1
0
mirror of https://github.com/actix/actix-extras.git synced 2024-11-28 09:42:40 +01:00
actix-extras/src/body.rs

445 lines
11 KiB
Rust
Raw Normal View History

use std::marker::PhantomData;
use std::{fmt, mem};
2017-10-24 08:25:32 +02:00
2018-11-14 07:53:30 +01:00
use bytes::{Bytes, BytesMut};
use futures::{Async, Poll, Stream};
2018-12-06 23:32:52 +01:00
use crate::error::{Error, PayloadError};
/// Type represent streaming payload
pub type PayloadStream = Box<dyn Stream<Item = Bytes, Error = PayloadError>>;
2017-10-24 08:25:32 +02:00
2018-11-19 02:52:56 +01:00
#[derive(Debug, PartialEq, Copy, Clone)]
/// Different type of body
pub enum BodyLength {
2018-11-14 07:53:30 +01:00
None,
Empty,
2018-11-14 07:53:30 +01:00
Sized(usize),
Sized64(u64),
2018-11-18 05:21:28 +01:00
Chunked,
Stream,
2018-11-14 07:53:30 +01:00
}
2019-02-06 20:44:15 +01:00
impl BodyLength {
pub fn is_eof(&self) -> bool {
match self {
BodyLength::None
| BodyLength::Empty
| BodyLength::Sized(0)
| BodyLength::Sized64(0) => true,
_ => false,
}
}
}
2018-11-14 07:53:30 +01:00
/// Type that provides this trait can be streamed to a peer.
pub trait MessageBody {
fn length(&self) -> BodyLength;
2018-11-14 07:53:30 +01:00
fn poll_next(&mut self) -> Poll<Option<Bytes>, Error>;
}
impl MessageBody for () {
fn length(&self) -> BodyLength {
BodyLength::Empty
2018-11-14 07:53:30 +01:00
}
fn poll_next(&mut self) -> Poll<Option<Bytes>, Error> {
Ok(Async::Ready(None))
}
}
2018-11-21 16:49:24 +01:00
pub enum ResponseBody<B> {
Body(B),
Other(Body),
}
2019-02-06 20:44:15 +01:00
impl ResponseBody<Body> {
pub fn into_body<B>(self) -> ResponseBody<B> {
match self {
ResponseBody::Body(b) => ResponseBody::Other(b),
ResponseBody::Other(b) => ResponseBody::Other(b),
}
}
}
2018-11-21 16:49:24 +01:00
impl<B: MessageBody> ResponseBody<B> {
pub fn as_ref(&self) -> Option<&B> {
if let ResponseBody::Body(ref b) = self {
Some(b)
} else {
None
}
}
}
impl<B: MessageBody> MessageBody for ResponseBody<B> {
fn length(&self) -> BodyLength {
match self {
ResponseBody::Body(ref body) => body.length(),
ResponseBody::Other(ref body) => body.length(),
}
}
fn poll_next(&mut self) -> Poll<Option<Bytes>, Error> {
match self {
ResponseBody::Body(ref mut body) => body.poll_next(),
ResponseBody::Other(ref mut body) => body.poll_next(),
}
}
}
/// Represents various types of http message body.
pub enum Body {
/// Empty response. `Content-Length` header is not set.
None,
/// Zero sized response body. `Content-Length` header is set to `0`.
Empty,
/// Specific response body.
2017-10-24 08:25:32 +02:00
Bytes(Bytes),
/// Generic message body.
Message(Box<dyn MessageBody>),
2017-10-24 08:25:32 +02:00
}
impl Body {
/// Create body from slice (copy)
pub fn from_slice(s: &[u8]) -> Body {
Body::Bytes(Bytes::from(s))
2017-10-24 08:25:32 +02:00
}
2018-01-14 23:40:39 +01:00
/// Create body from generic message body.
pub fn from_message<B: MessageBody + 'static>(body: B) -> Body {
Body::Message(Box::new(body))
2018-01-14 23:40:39 +01:00
}
}
impl MessageBody for Body {
fn length(&self) -> BodyLength {
match self {
Body::None => BodyLength::None,
Body::Empty => BodyLength::Empty,
Body::Bytes(ref bin) => BodyLength::Sized(bin.len()),
Body::Message(ref body) => body.length(),
2018-01-14 23:40:39 +01:00
}
}
2017-10-24 08:25:32 +02:00
fn poll_next(&mut self) -> Poll<Option<Bytes>, Error> {
2018-01-11 00:28:33 +01:00
match self {
Body::None => Ok(Async::Ready(None)),
Body::Empty => Ok(Async::Ready(None)),
Body::Bytes(ref mut bin) => {
2018-11-19 02:52:56 +01:00
let len = bin.len();
if len == 0 {
Ok(Async::Ready(None))
} else {
2018-11-19 02:52:56 +01:00
Ok(Async::Ready(Some(bin.split_to(len))))
}
}
Body::Message(ref mut body) => body.poll_next(),
2018-01-11 00:28:33 +01:00
}
}
}
impl PartialEq for Body {
fn eq(&self, other: &Body) -> bool {
match *self {
Body::None => match *other {
Body::None => true,
_ => false,
},
Body::Empty => match *other {
Body::Empty => true,
_ => false,
},
Body::Bytes(ref b) => match *other {
Body::Bytes(ref b2) => b == b2,
_ => false,
},
Body::Message(_) => false,
}
2017-10-24 08:25:32 +02:00
}
}
impl fmt::Debug for Body {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match *self {
Body::None => write!(f, "Body::None"),
Body::Empty => write!(f, "Body::Zero"),
Body::Bytes(ref b) => write!(f, "Body::Bytes({:?})", b),
Body::Message(_) => write!(f, "Body::Message(_)"),
}
2017-10-30 06:50:21 +01:00
}
}
impl From<&'static str> for Body {
fn from(s: &'static str) -> Body {
Body::Bytes(Bytes::from_static(s.as_ref()))
2017-10-24 08:39:01 +02:00
}
}
impl From<&'static [u8]> for Body {
fn from(s: &'static [u8]) -> Body {
2019-01-29 19:14:00 +01:00
Body::Bytes(Bytes::from_static(s))
2017-10-24 08:39:01 +02:00
}
}
impl From<Vec<u8>> for Body {
fn from(vec: Vec<u8>) -> Body {
Body::Bytes(Bytes::from(vec))
2017-10-24 08:49:27 +02:00
}
}
impl From<String> for Body {
fn from(s: String) -> Body {
s.into_bytes().into()
2017-10-24 08:49:27 +02:00
}
}
impl<'a> From<&'a String> for Body {
fn from(s: &'a String) -> Body {
Body::Bytes(Bytes::from(AsRef::<[u8]>::as_ref(&s)))
}
}
impl From<Bytes> for Body {
fn from(s: Bytes) -> Body {
Body::Bytes(s)
}
}
impl From<BytesMut> for Body {
fn from(s: BytesMut) -> Body {
Body::Bytes(s.freeze())
2017-10-24 08:25:32 +02:00
}
}
2017-10-24 09:09:52 +02:00
2018-11-14 07:53:30 +01:00
impl MessageBody for Bytes {
fn length(&self) -> BodyLength {
BodyLength::Sized(self.len())
2018-11-14 07:53:30 +01:00
}
fn poll_next(&mut self) -> Poll<Option<Bytes>, Error> {
if self.is_empty() {
Ok(Async::Ready(None))
} else {
Ok(Async::Ready(Some(mem::replace(self, Bytes::new()))))
}
}
}
2018-11-18 05:21:28 +01:00
impl MessageBody for BytesMut {
fn length(&self) -> BodyLength {
BodyLength::Sized(self.len())
}
fn poll_next(&mut self) -> Poll<Option<Bytes>, Error> {
if self.is_empty() {
Ok(Async::Ready(None))
} else {
Ok(Async::Ready(Some(
mem::replace(self, BytesMut::new()).freeze(),
)))
}
}
}
2018-11-14 07:53:30 +01:00
impl MessageBody for &'static str {
fn length(&self) -> BodyLength {
BodyLength::Sized(self.len())
2018-11-14 07:53:30 +01:00
}
fn poll_next(&mut self) -> Poll<Option<Bytes>, Error> {
if self.is_empty() {
Ok(Async::Ready(None))
} else {
Ok(Async::Ready(Some(Bytes::from_static(
mem::replace(self, "").as_ref(),
))))
}
}
}
impl MessageBody for &'static [u8] {
fn length(&self) -> BodyLength {
BodyLength::Sized(self.len())
2018-11-14 07:53:30 +01:00
}
fn poll_next(&mut self) -> Poll<Option<Bytes>, Error> {
if self.is_empty() {
Ok(Async::Ready(None))
} else {
Ok(Async::Ready(Some(Bytes::from_static(mem::replace(
self, b"",
)))))
}
}
}
impl MessageBody for Vec<u8> {
fn length(&self) -> BodyLength {
BodyLength::Sized(self.len())
2018-11-14 07:53:30 +01:00
}
fn poll_next(&mut self) -> Poll<Option<Bytes>, Error> {
if self.is_empty() {
Ok(Async::Ready(None))
} else {
Ok(Async::Ready(Some(Bytes::from(mem::replace(
self,
Vec::new(),
)))))
}
}
}
impl MessageBody for String {
fn length(&self) -> BodyLength {
BodyLength::Sized(self.len())
2018-11-14 07:53:30 +01:00
}
fn poll_next(&mut self) -> Poll<Option<Bytes>, Error> {
if self.is_empty() {
Ok(Async::Ready(None))
} else {
Ok(Async::Ready(Some(Bytes::from(
mem::replace(self, String::new()).into_bytes(),
))))
}
}
}
/// Type represent streaming body.
/// Response does not contain `content-length` header and appropriate transfer encoding is used.
pub struct BodyStream<S, E> {
2018-11-14 07:53:30 +01:00
stream: S,
_t: PhantomData<E>,
2018-11-14 07:53:30 +01:00
}
impl<S, E> BodyStream<S, E>
2018-11-14 07:53:30 +01:00
where
S: Stream<Item = Bytes, Error = E>,
E: Into<Error>,
2018-11-14 07:53:30 +01:00
{
pub fn new(stream: S) -> Self {
BodyStream {
stream,
_t: PhantomData,
}
2018-11-14 07:53:30 +01:00
}
}
impl<S, E> MessageBody for BodyStream<S, E>
2018-11-14 07:53:30 +01:00
where
S: Stream<Item = Bytes, Error = E>,
E: Into<Error>,
2018-11-14 07:53:30 +01:00
{
fn length(&self) -> BodyLength {
2018-11-18 05:21:28 +01:00
BodyLength::Chunked
2018-11-14 07:53:30 +01:00
}
fn poll_next(&mut self) -> Poll<Option<Bytes>, Error> {
self.stream.poll().map_err(|e| e.into())
}
}
/// Type represent streaming body. This body implementation should be used
/// if total size of stream is known. Data get sent as is without using transfer encoding.
pub struct SizedStream<S> {
size: usize,
stream: S,
}
impl<S> SizedStream<S>
where
S: Stream<Item = Bytes, Error = Error>,
{
pub fn new(size: usize, stream: S) -> Self {
SizedStream { size, stream }
}
}
impl<S> MessageBody for SizedStream<S>
where
S: Stream<Item = Bytes, Error = Error>,
{
fn length(&self) -> BodyLength {
BodyLength::Sized(self.size)
}
2018-11-14 07:53:30 +01:00
fn poll_next(&mut self) -> Poll<Option<Bytes>, Error> {
self.stream.poll()
}
}
2017-10-24 09:09:52 +02:00
#[cfg(test)]
mod tests {
use super::*;
impl Body {
pub(crate) fn get_ref(&self) -> &[u8] {
match *self {
Body::Bytes(ref bin) => &bin,
_ => panic!(),
}
}
2017-11-20 05:32:37 +01:00
}
2018-11-21 16:49:24 +01:00
impl ResponseBody<Body> {
pub(crate) fn get_ref(&self) -> &[u8] {
match *self {
ResponseBody::Body(ref b) => b.get_ref(),
ResponseBody::Other(ref b) => b.get_ref(),
}
}
}
2017-10-24 09:09:52 +02:00
#[test]
fn test_static_str() {
assert_eq!(Body::from("").length(), BodyLength::Sized(0));
assert_eq!(Body::from("test").length(), BodyLength::Sized(4));
assert_eq!(Body::from("test").get_ref(), b"test");
2017-10-24 09:09:52 +02:00
}
#[test]
fn test_static_bytes() {
assert_eq!(Body::from(b"test".as_ref()).length(), BodyLength::Sized(4));
assert_eq!(Body::from(b"test".as_ref()).get_ref(), b"test");
assert_eq!(
Body::from_slice(b"test".as_ref()).length(),
BodyLength::Sized(4)
);
assert_eq!(Body::from_slice(b"test".as_ref()).get_ref(), b"test");
2017-10-24 09:09:52 +02:00
}
#[test]
fn test_vec() {
assert_eq!(Body::from(Vec::from("test")).length(), BodyLength::Sized(4));
assert_eq!(Body::from(Vec::from("test")).get_ref(), b"test");
2017-10-24 09:09:52 +02:00
}
#[test]
fn test_bytes() {
assert_eq!(
Body::from(Bytes::from("test")).length(),
BodyLength::Sized(4)
);
assert_eq!(Body::from(Bytes::from("test")).get_ref(), b"test");
2017-10-24 09:09:52 +02:00
}
2017-11-20 05:32:37 +01:00
#[test]
fn test_string() {
let b = "test".to_owned();
assert_eq!(Body::from(b.clone()).length(), BodyLength::Sized(4));
assert_eq!(Body::from(b.clone()).get_ref(), b"test");
assert_eq!(Body::from(&b).length(), BodyLength::Sized(4));
assert_eq!(Body::from(&b).get_ref(), b"test");
}
2017-11-20 05:32:37 +01:00
#[test]
fn test_bytes_mut() {
2018-04-14 01:02:01 +02:00
let b = BytesMut::from("test");
assert_eq!(Body::from(b.clone()).length(), BodyLength::Sized(4));
assert_eq!(Body::from(b).get_ref(), b"test");
2018-01-11 00:28:33 +01:00
}
2017-10-24 09:09:52 +02:00
}