#[cfg(feature = "compression")]
use super::compression::{compress, CompressionEncoding, SingleMessageCompressionOverride};
use super::{EncodeBuf, Encoder, HEADER_SIZE};
use crate::{Code, Status};
use bytes::{BufMut, Bytes, BytesMut};
use futures_core::{Stream, TryStream};
use futures_util::{ready, StreamExt, TryStreamExt};
use http::HeaderMap;
use http_body::Body;
use pin_project::pin_project;
use std::{
pin::Pin,
task::{Context, Poll},
};
pub(super) const BUFFER_SIZE: usize = 8 * 1024;
pub(crate) fn encode_server<T, U>(
encoder: T,
source: U,
#[cfg(feature = "compression")] compression_encoding: Option<CompressionEncoding>,
#[cfg(feature = "compression")] compression_override: SingleMessageCompressionOverride,
) -> EncodeBody<impl Stream<Item = Result<Bytes, Status>>>
where
T: Encoder<Error = Status>,
U: Stream<Item = Result<T::Item, Status>>,
{
let stream = encode(
encoder,
source,
#[cfg(feature = "compression")]
compression_encoding,
#[cfg(feature = "compression")]
compression_override,
)
.into_stream();
EncodeBody::new_server(stream)
}
pub(crate) fn encode_client<T, U>(
encoder: T,
source: U,
#[cfg(feature = "compression")] compression_encoding: Option<CompressionEncoding>,
) -> EncodeBody<impl Stream<Item = Result<Bytes, Status>>>
where
T: Encoder<Error = Status>,
U: Stream<Item = T::Item>,
{
let stream = encode(
encoder,
source.map(Ok),
#[cfg(feature = "compression")]
compression_encoding,
#[cfg(feature = "compression")]
SingleMessageCompressionOverride::default(),
)
.into_stream();
EncodeBody::new_client(stream)
}
fn encode<T, U>(
mut encoder: T,
source: U,
#[cfg(feature = "compression")] compression_encoding: Option<CompressionEncoding>,
#[cfg(feature = "compression")] compression_override: SingleMessageCompressionOverride,
) -> impl TryStream<Ok = Bytes, Error = Status>
where
T: Encoder<Error = Status>,
U: Stream<Item = Result<T::Item, Status>>,
{
async_stream::stream! {
let mut buf = BytesMut::with_capacity(BUFFER_SIZE);
#[cfg(feature = "compression")]
let (compression_enabled_for_stream, mut uncompression_buf) = match compression_encoding {
Some(CompressionEncoding::Gzip) => (true, BytesMut::with_capacity(BUFFER_SIZE)),
None => (false, BytesMut::new()),
};
#[cfg(feature = "compression")]
let compress_item = compression_enabled_for_stream && compression_override == SingleMessageCompressionOverride::Inherit;
#[cfg(not(feature = "compression"))]
let compress_item = false;
futures_util::pin_mut!(source);
loop {
match source.next().await {
Some(Ok(item)) => {
buf.reserve(HEADER_SIZE);
unsafe {
buf.advance_mut(HEADER_SIZE);
}
if compress_item {
#[cfg(feature = "compression")]
{
uncompression_buf.clear();
encoder.encode(item, &mut EncodeBuf::new(&mut uncompression_buf))
.map_err(|err| Status::internal(format!("Error encoding: {}", err)))?;
let uncompressed_len = uncompression_buf.len();
compress(
compression_encoding.unwrap(),
&mut uncompression_buf,
&mut buf,
uncompressed_len,
).map_err(|err| Status::internal(format!("Error compressing: {}", err)))?;
}
#[cfg(not(feature = "compression"))]
unreachable!("compression disabled, should not take this branch");
} else {
encoder.encode(item, &mut EncodeBuf::new(&mut buf))
.map_err(|err| Status::internal(format!("Error encoding: {}", err)))?;
}
let len = buf.len() - HEADER_SIZE;
assert!(len <= std::u32::MAX as usize);
{
let mut buf = &mut buf[..HEADER_SIZE];
buf.put_u8(compress_item as u8);
buf.put_u32(len as u32);
}
yield Ok(buf.split_to(len + HEADER_SIZE).freeze());
},
Some(Err(status)) => yield Err(status),
None => break,
}
}
}
}
#[derive(Debug)]
enum Role {
Client,
Server,
}
#[pin_project]
#[derive(Debug)]
pub(crate) struct EncodeBody<S> {
#[pin]
inner: S,
error: Option<Status>,
role: Role,
is_end_stream: bool,
}
impl<S> EncodeBody<S>
where
S: Stream<Item = Result<Bytes, Status>>,
{
pub(crate) fn new_client(inner: S) -> Self {
Self {
inner,
error: None,
role: Role::Client,
is_end_stream: false,
}
}
pub(crate) fn new_server(inner: S) -> Self {
Self {
inner,
error: None,
role: Role::Server,
is_end_stream: false,
}
}
}
impl<S> Body for EncodeBody<S>
where
S: Stream<Item = Result<Bytes, Status>>,
{
type Data = Bytes;
type Error = Status;
fn is_end_stream(&self) -> bool {
self.is_end_stream
}
fn poll_data(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Self::Data, Self::Error>>> {
let mut self_proj = self.project();
match ready!(self_proj.inner.try_poll_next_unpin(cx)) {
Some(Ok(d)) => Some(Ok(d)).into(),
Some(Err(status)) => match self_proj.role {
Role::Client => Some(Err(status)).into(),
Role::Server => {
*self_proj.error = Some(status);
None.into()
}
},
None => None.into(),
}
}
fn poll_trailers(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Result<Option<HeaderMap>, Status>> {
match self.role {
Role::Client => Poll::Ready(Ok(None)),
Role::Server => {
let self_proj = self.project();
if *self_proj.is_end_stream {
return Poll::Ready(Ok(None));
}
let status = if let Some(status) = self_proj.error.take() {
*self_proj.is_end_stream = true;
status
} else {
Status::new(Code::Ok, "")
};
Poll::Ready(Ok(Some(status.to_header_map()?)))
}
}
}
}