tet_libp2p_deflate/
lib.rs

1// Copyright 2019 Parity Technologies (UK) Ltd.
2//
3// Permission is hereby granted, free of charge, to any person obtaining a
4// copy of this software and associated documentation files (the "Software"),
5// to deal in the Software without restriction, including without limitation
6// the rights to use, copy, modify, merge, publish, distribute, sublicense,
7// and/or sell copies of the Software, and to permit persons to whom the
8// Software is furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
14// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
18// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
19// DEALINGS IN THE SOFTWARE.
20
21use futures::{prelude::*, ready};
22use tet_libp2p_core::{InboundUpgrade, OutboundUpgrade, UpgradeInfo};
23use std::{io, iter, pin::Pin, task::Context, task::Poll};
24
25#[derive(Debug, Copy, Clone)]
26pub struct DeflateConfig {
27    compression: flate2::Compression,
28}
29
30impl Default for DeflateConfig {
31    fn default() -> Self {
32        DeflateConfig {
33            compression: flate2::Compression::fast(),
34        }
35    }
36}
37
38impl UpgradeInfo for DeflateConfig {
39    type Info = &'static [u8];
40    type InfoIter = iter::Once<Self::Info>;
41
42    fn protocol_info(&self) -> Self::InfoIter {
43        iter::once(b"/deflate/1.0.0")
44    }
45}
46
47impl<C> InboundUpgrade<C> for DeflateConfig
48where
49    C: AsyncRead + AsyncWrite,
50{
51    type Output = DeflateOutput<C>;
52    type Error = io::Error;
53    type Future = future::Ready<Result<Self::Output, Self::Error>>;
54
55    fn upgrade_inbound(self, r: C, _: Self::Info) -> Self::Future {
56        future::ok(DeflateOutput::new(r, self.compression))
57    }
58}
59
60impl<C> OutboundUpgrade<C> for DeflateConfig
61where
62    C: AsyncRead + AsyncWrite,
63{
64    type Output = DeflateOutput<C>;
65    type Error = io::Error;
66    type Future = future::Ready<Result<Self::Output, Self::Error>>;
67
68    fn upgrade_outbound(self, w: C, _: Self::Info) -> Self::Future {
69        future::ok(DeflateOutput::new(w, self.compression))
70    }
71}
72
73/// Decodes and encodes traffic using DEFLATE.
74#[derive(Debug)]
75pub struct DeflateOutput<S> {
76    /// Inner stream where we read compressed data from and write compressed data to.
77    inner: S,
78    /// Internal object used to hold the state of the compression.
79    compress: flate2::Compress,
80    /// Internal object used to hold the state of the decompression.
81    decompress: flate2::Decompress,
82    /// Temporary buffer between `compress` and `inner`. Stores compressed bytes that need to be
83    /// sent out once `inner` is ready to accept more.
84    write_out: Vec<u8>,
85    /// Temporary buffer between `decompress` and `inner`. Stores compressed bytes that need to be
86    /// given to `decompress`.
87    read_interm: Vec<u8>,
88    /// When we read from `inner` and `Ok(0)` is returned, we set this to `true` so that we don't
89    /// read from it again.
90    inner_read_eof: bool,
91}
92
93impl<S> DeflateOutput<S> {
94    fn new(inner: S, compression: flate2::Compression) -> Self {
95        DeflateOutput {
96            inner,
97            compress: flate2::Compress::new(compression, false),
98            decompress: flate2::Decompress::new(false),
99            write_out: Vec::with_capacity(256),
100            read_interm: Vec::with_capacity(256),
101            inner_read_eof: false,
102        }
103    }
104
105    /// Tries to write the content of `self.write_out` to `self.inner`.
106    /// Returns `Ready(Ok(()))` if `self.write_out` is empty.
107    fn flush_write_out(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>>
108        where S: AsyncWrite + Unpin
109    {
110        loop {
111            if self.write_out.is_empty() {
112                return Poll::Ready(Ok(()))
113            }
114
115            match AsyncWrite::poll_write(Pin::new(&mut self.inner), cx, &self.write_out) {
116                Poll::Ready(Ok(0)) => return Poll::Ready(Err(io::ErrorKind::WriteZero.into())),
117                Poll::Ready(Ok(n)) => self.write_out = self.write_out.split_off(n),
118                Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
119                Poll::Pending => return Poll::Pending,
120            };
121        }
122    }
123}
124
125impl<S> AsyncRead for DeflateOutput<S>
126    where S: AsyncRead + Unpin
127{
128    fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<Result<usize, io::Error>> {
129        // We use a `this` variable because the compiler doesn't allow multiple mutable borrows
130        // across a `Deref`.
131        let this = &mut *self;
132
133        loop {
134            // Read from `self.inner` into `self.read_interm` if necessary.
135            if this.read_interm.is_empty() && !this.inner_read_eof {
136                unsafe {
137                    this.read_interm.reserve(256);
138                    this.read_interm.set_len(this.read_interm.capacity());
139                }
140
141                match AsyncRead::poll_read(Pin::new(&mut this.inner), cx, &mut this.read_interm) {
142                    Poll::Ready(Ok(0)) => {
143                        this.inner_read_eof = true;
144                        this.read_interm.clear();
145                    }
146                    Poll::Ready(Ok(n)) => {
147                        this.read_interm.truncate(n)
148                    },
149                    Poll::Ready(Err(err)) => {
150                        this.read_interm.clear();
151                        return Poll::Ready(Err(err))
152                    },
153                    Poll::Pending => {
154                        this.read_interm.clear();
155                        return Poll::Pending
156                    },
157                }
158            }
159            debug_assert!(!this.read_interm.is_empty() || this.inner_read_eof);
160
161            let before_out = this.decompress.total_out();
162            let before_in = this.decompress.total_in();
163            let ret = this.decompress.decompress(&this.read_interm, buf, if this.inner_read_eof { flate2::FlushDecompress::Finish } else { flate2::FlushDecompress::None })?;
164
165            // Remove from `self.read_interm` the bytes consumed by the decompressor.
166            let consumed = (this.decompress.total_in() - before_in) as usize;
167            this.read_interm = this.read_interm.split_off(consumed);
168
169            let read = (this.decompress.total_out() - before_out) as usize;
170            if read != 0 || ret == flate2::Status::StreamEnd {
171                return Poll::Ready(Ok(read))
172            }
173        }
174    }
175}
176
177impl<S> AsyncWrite for DeflateOutput<S>
178    where S: AsyncWrite + Unpin
179{
180    fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8])
181        -> Poll<Result<usize, io::Error>>
182    {
183        // We use a `this` variable because the compiler doesn't allow multiple mutable borrows
184        // across a `Deref`.
185        let this = &mut *self;
186
187        // We don't want to accumulate too much data in `self.write_out`, so we only proceed if it
188        // is empty.
189        ready!(this.flush_write_out(cx))?;
190
191        // We special-case this, otherwise an empty buffer would make the loop below infinite.
192        if buf.is_empty() {
193            return Poll::Ready(Ok(0));
194        }
195
196        // Unfortunately, the compressor might be in a "flushing mode", not accepting any input
197        // data. We don't want to return `Ok(0)` in that situation, as that would be wrong.
198        // Instead, we invoke the compressor in a loop until it accepts some of our data.
199        loop {
200            let before_in = this.compress.total_in();
201            this.write_out.reserve(256);  // compress_vec uses the Vec's capacity
202            let ret = this.compress.compress_vec(buf, &mut this.write_out, flate2::FlushCompress::None)?;
203            let written = (this.compress.total_in() - before_in) as usize;
204
205            if written != 0 || ret == flate2::Status::StreamEnd {
206                return Poll::Ready(Ok(written));
207            }
208        }
209    }
210
211    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
212        // We use a `this` variable because the compiler doesn't allow multiple mutable borrows
213        // across a `Deref`.
214        let this = &mut *self;
215
216        ready!(this.flush_write_out(cx))?;
217        this.compress.compress_vec(&[], &mut this.write_out, flate2::FlushCompress::Sync)?;
218
219        loop {
220            ready!(this.flush_write_out(cx))?;
221
222            debug_assert!(this.write_out.is_empty());
223            // We ask the compressor to flush everything into `self.write_out`.
224            this.write_out.reserve(256);  // compress_vec uses the Vec's capacity
225            this.compress.compress_vec(&[], &mut this.write_out, flate2::FlushCompress::None)?;
226            if this.write_out.is_empty() {
227                break;
228            }
229        }
230
231        AsyncWrite::poll_flush(Pin::new(&mut this.inner), cx)
232    }
233
234    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
235        // We use a `this` variable because the compiler doesn't allow multiple mutable borrows
236        // across a `Deref`.
237        let this = &mut *self;
238
239        loop {
240            ready!(this.flush_write_out(cx))?;
241
242            // We ask the compressor to flush everything into `self.write_out`.
243            debug_assert!(this.write_out.is_empty());
244            this.write_out.reserve(256);  // compress_vec uses the Vec's capacity
245            this.compress.compress_vec(&[], &mut this.write_out, flate2::FlushCompress::Finish)?;
246            if this.write_out.is_empty() {
247                break;
248            }
249        }
250
251        AsyncWrite::poll_close(Pin::new(&mut this.inner), cx)
252    }
253}
tet_libp2p_deflate/lib.rs

tet_libp2p_deflate/
lib.rs
