use futures::{prelude::*, ready};
use libp2p_core::{InboundUpgrade, OutboundUpgrade, UpgradeInfo};
use std::{io, iter, pin::Pin, task::Context, task::Poll};
#[derive(Debug, Copy, Clone)]
pub struct DeflateConfig {
compression: flate2::Compression,
}
impl Default for DeflateConfig {
fn default() -> Self {
DeflateConfig {
compression: flate2::Compression::fast(),
}
}
}
impl UpgradeInfo for DeflateConfig {
type Info = &'static [u8];
type InfoIter = iter::Once<Self::Info>;
fn protocol_info(&self) -> Self::InfoIter {
iter::once(b"/deflate/1.0.0")
}
}
impl<C> InboundUpgrade<C> for DeflateConfig
where
C: AsyncRead + AsyncWrite,
{
type Output = DeflateOutput<C>;
type Error = io::Error;
type Future = future::Ready<Result<Self::Output, Self::Error>>;
fn upgrade_inbound(self, r: C, _: Self::Info) -> Self::Future {
future::ok(DeflateOutput::new(r, self.compression))
}
}
impl<C> OutboundUpgrade<C> for DeflateConfig
where
C: AsyncRead + AsyncWrite,
{
type Output = DeflateOutput<C>;
type Error = io::Error;
type Future = future::Ready<Result<Self::Output, Self::Error>>;
fn upgrade_outbound(self, w: C, _: Self::Info) -> Self::Future {
future::ok(DeflateOutput::new(w, self.compression))
}
}
#[derive(Debug)]
pub struct DeflateOutput<S> {
inner: S,
compress: flate2::Compress,
decompress: flate2::Decompress,
write_out: Vec<u8>,
read_interm: Vec<u8>,
inner_read_eof: bool,
}
impl<S> DeflateOutput<S> {
fn new(inner: S, compression: flate2::Compression) -> Self {
DeflateOutput {
inner,
compress: flate2::Compress::new(compression, false),
decompress: flate2::Decompress::new(false),
write_out: Vec::with_capacity(256),
read_interm: Vec::with_capacity(256),
inner_read_eof: false,
}
}
fn flush_write_out(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>>
where S: AsyncWrite + Unpin
{
loop {
if self.write_out.is_empty() {
return Poll::Ready(Ok(()))
}
match AsyncWrite::poll_write(Pin::new(&mut self.inner), cx, &self.write_out) {
Poll::Ready(Ok(0)) => return Poll::Ready(Err(io::ErrorKind::WriteZero.into())),
Poll::Ready(Ok(n)) => self.write_out = self.write_out.split_off(n),
Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
Poll::Pending => return Poll::Pending,
};
}
}
}
impl<S> AsyncRead for DeflateOutput<S>
where S: AsyncRead + Unpin
{
fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<Result<usize, io::Error>> {
let this = &mut *self;
loop {
if this.read_interm.is_empty() && !this.inner_read_eof {
unsafe {
this.read_interm.reserve(256);
this.read_interm.set_len(this.read_interm.capacity());
}
match AsyncRead::poll_read(Pin::new(&mut this.inner), cx, &mut this.read_interm) {
Poll::Ready(Ok(0)) => {
this.inner_read_eof = true;
this.read_interm.clear();
}
Poll::Ready(Ok(n)) => {
this.read_interm.truncate(n)
},
Poll::Ready(Err(err)) => {
this.read_interm.clear();
return Poll::Ready(Err(err))
},
Poll::Pending => {
this.read_interm.clear();
return Poll::Pending
},
}
}
debug_assert!(!this.read_interm.is_empty() || this.inner_read_eof);
let before_out = this.decompress.total_out();
let before_in = this.decompress.total_in();
let ret = this.decompress.decompress(&this.read_interm, buf, if this.inner_read_eof { flate2::FlushDecompress::Finish } else { flate2::FlushDecompress::None })?;
let consumed = (this.decompress.total_in() - before_in) as usize;
this.read_interm = this.read_interm.split_off(consumed);
let read = (this.decompress.total_out() - before_out) as usize;
if read != 0 || ret == flate2::Status::StreamEnd {
return Poll::Ready(Ok(read))
}
}
}
}
impl<S> AsyncWrite for DeflateOutput<S>
where S: AsyncWrite + Unpin
{
fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8])
-> Poll<Result<usize, io::Error>>
{
let this = &mut *self;
ready!(this.flush_write_out(cx))?;
if buf.is_empty() {
return Poll::Ready(Ok(0));
}
loop {
let before_in = this.compress.total_in();
this.write_out.reserve(256); let ret = this.compress.compress_vec(buf, &mut this.write_out, flate2::FlushCompress::None)?;
let written = (this.compress.total_in() - before_in) as usize;
if written != 0 || ret == flate2::Status::StreamEnd {
return Poll::Ready(Ok(written));
}
}
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
let this = &mut *self;
ready!(this.flush_write_out(cx))?;
this.compress.compress_vec(&[], &mut this.write_out, flate2::FlushCompress::Sync)?;
loop {
ready!(this.flush_write_out(cx))?;
debug_assert!(this.write_out.is_empty());
this.write_out.reserve(256); this.compress.compress_vec(&[], &mut this.write_out, flate2::FlushCompress::None)?;
if this.write_out.is_empty() {
break;
}
}
AsyncWrite::poll_flush(Pin::new(&mut this.inner), cx)
}
fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
let this = &mut *self;
loop {
ready!(this.flush_write_out(cx))?;
debug_assert!(this.write_out.is_empty());
this.write_out.reserve(256); this.compress.compress_vec(&[], &mut this.write_out, flate2::FlushCompress::Finish)?;
if this.write_out.is_empty() {
break;
}
}
AsyncWrite::poll_close(Pin::new(&mut this.inner), cx)
}
}