diff --git a/src/encoding/frame_compressor.rs b/src/encoding/frame_compressor.rs index 29cdfad..519bd52 100644 --- a/src/encoding/frame_compressor.rs +++ b/src/encoding/frame_compressor.rs @@ -76,6 +76,9 @@ impl FrameCompressor { /// /// This will repeatedly call [Read::read] on the source to fill up blocks until the source returns 0 on the read call. /// Also [Write::write_all] will be called on the drain after each block has been encoded. + /// + /// To avoid endlessly encoding from a potentially endless source (like a network socket) you can use the + /// [Read::take] function pub fn compress(&mut self) { self.match_generator.reset(self.compression_level); let source = self.uncompressed_data.as_mut().unwrap(); diff --git a/src/io_nostd.rs b/src/io_nostd.rs index b0f989b..5eb1437 100644 --- a/src/io_nostd.rs +++ b/src/io_nostd.rs @@ -127,6 +127,10 @@ pub trait Read { } Ok(()) } + + fn take(self, limit: u64) -> Take where Self: Sized { + Take { inner: self, limit } + } } impl Read for &[u8] { @@ -154,6 +158,46 @@ where } } +pub struct Take { + inner: R, + limit: u64, +} + +impl Take { + pub fn limit(&self) -> u64 { + self.limit + } + + pub fn set_limit(&mut self, limit: u64) { + self.limit = limit; + } + + pub fn get_ref(&self) -> &R { + &self.inner + } + + pub fn get_mut(&mut self) -> &mut R { + &mut self.inner + } + + pub fn into_inner(self) -> R { + self.inner + } +} + +impl Read for Take { + fn read(&mut self, buf: &mut [u8]) -> Result { + if self.limit == 0 { + return Ok(0); + } + + let at_most = (self.limit as usize).min(buf.len()); + let bytes = self.inner.read(&mut buf[..at_most])?; + self.limit -= bytes as u64; + Ok(bytes) + } +} + pub trait Write { fn write(&mut self, buf: &[u8]) -> Result; fn flush(&mut self) -> Result<(), Error>;