2021-02-15 12:20:43 +00:00
|
|
|
use std::{
|
|
|
|
|
mem,
|
|
|
|
|
pin::Pin,
|
|
|
|
|
task::{Context, Poll},
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
use bytes::Bytes;
|
|
|
|
|
use futures_core::Stream;
|
|
|
|
|
use pin_project::pin_project;
|
|
|
|
|
|
|
|
|
|
use crate::error::Error;
|
|
|
|
|
|
|
|
|
|
use super::{Body, BodySize, MessageBody};
|
|
|
|
|
|
|
|
|
|
#[pin_project(project = ResponseBodyProj)]
|
|
|
|
|
pub enum ResponseBody<B> {
|
|
|
|
|
Body(#[pin] B),
|
|
|
|
|
Other(Body),
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
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),
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
impl<B> ResponseBody<B> {
|
|
|
|
|
pub fn take_body(&mut self) -> ResponseBody<B> {
|
|
|
|
|
mem::replace(self, ResponseBody::Other(Body::None))
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
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 size(&self) -> BodySize {
|
|
|
|
|
match self {
|
|
|
|
|
ResponseBody::Body(ref body) => body.size(),
|
|
|
|
|
ResponseBody::Other(ref body) => body.size(),
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
fn poll_next(
|
|
|
|
|
self: Pin<&mut Self>,
|
|
|
|
|
cx: &mut Context<'_>,
|
|
|
|
|
) -> Poll<Option<Result<Bytes, Error>>> {
|
2021-04-13 13:34:22 +01:00
|
|
|
Stream::poll_next(self, cx)
|
2021-02-15 12:20:43 +00:00
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
impl<B: MessageBody> Stream for ResponseBody<B> {
|
|
|
|
|
type Item = Result<Bytes, Error>;
|
|
|
|
|
|
|
|
|
|
fn poll_next(
|
|
|
|
|
self: Pin<&mut Self>,
|
|
|
|
|
cx: &mut Context<'_>,
|
|
|
|
|
) -> Poll<Option<Self::Item>> {
|
|
|
|
|
match self.project() {
|
|
|
|
|
ResponseBodyProj::Body(body) => body.poll_next(cx),
|
|
|
|
|
ResponseBodyProj::Other(body) => Pin::new(body).poll_next(cx),
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
