diff --git a/src/byte_counter.rs b/src/byte_counter.rs new file mode 100644 index 0000000..547523b --- /dev/null +++ b/src/byte_counter.rs @@ -0,0 +1,75 @@ +// Copyright 2021, The Tari Project +// +// Redistribution and use in source and binary forms, with or without modification, are permitted provided that the +// following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following +// disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the +// following disclaimer in the documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote +// products derived from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE +// USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +use std::io; + +#[derive(Debug, Clone, Default)] +pub struct ByteCounter { + count: usize, +} + +impl ByteCounter { + pub fn new() -> Self { + Default::default() + } + + pub fn get(&self) -> usize { + self.count + } +} + +impl io::Write for ByteCounter { + fn write(&mut self, buf: &[u8]) -> io::Result { + let len = buf.len(); + self.count += len; + Ok(len) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} + +#[cfg(test)] +mod test { + use std::io::Write; + + use super::*; + + #[test] + fn write_test() { + let mut byte_counter = ByteCounter::new(); + let buf = [0u8, 1u8, 2u8, 3u8]; + let new_count = byte_counter.write(&buf).unwrap(); + assert_eq!(byte_counter.get(), new_count); + assert_eq!(byte_counter.get(), buf.len()); + } + + #[test] + fn flush_test() { + let mut byte_counter = ByteCounter::new(); + let buf = [0u8, 1u8, 2u8, 3u8]; + let _count_bytes = byte_counter.write(&buf).unwrap(); + // test passes if the following method does not return an error + byte_counter.flush().unwrap(); + } +} diff --git a/src/lib.rs b/src/lib.rs index 69ed574..c6409dc 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -23,12 +23,14 @@ //! A set of useful and commonly used utilities that are used in several places in the Tari project. pub mod bit; pub mod byte_array; +pub mod byte_counter; pub mod convert; pub mod encoding; pub mod epoch_time; pub mod fixed_set; pub mod hash; pub mod hex; +pub mod limited_reader; #[macro_use] pub mod locks; pub mod hidden; diff --git a/src/limited_reader.rs b/src/limited_reader.rs new file mode 100644 index 0000000..80dfd36 --- /dev/null +++ b/src/limited_reader.rs @@ -0,0 +1,74 @@ +// Copyright 2022, The Tari Project +// +// Redistribution and use in source and binary forms, with or without modification, are permitted provided that the +// following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following +// disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the +// following disclaimer in the documentation and/or other materials provided with the distribution. +// +// 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote +// products derived from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +// INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE +// USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +use std::{io, io::Read}; + +pub struct LimitedBytesReader { + byte_limit: usize, + num_read: usize, + inner: R, +} + +impl LimitedBytesReader { + pub fn new(byte_limit: usize, reader: R) -> Self { + Self { + byte_limit, + num_read: 0, + inner: reader, + } + } +} +impl Read for LimitedBytesReader { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + let read = self.inner.read(buf)?; + self.num_read += read; + if self.num_read > self.byte_limit { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("Read more bytes than the maximum ({})", self.byte_limit), + )); + } + Ok(read) + } +} + +#[cfg(test)] +mod test { + use std::io::Read; + + use super::*; + + #[test] + fn read_test() { + // read should work fine in the case of a buffer whose length is within byte_limit + let inner: &[u8] = &[0u8, 1u8, 2u8, 3u8, 4u8]; + let mut reader = LimitedBytesReader::new(3, inner); + let mut buf = [0u8; 3]; + let output = reader.read(&mut buf).unwrap(); + assert_eq!(output, buf.len()); + + // in case of buffer with length strictly bigger than reader byte_limit, the code should throw an error + let mut new_buf = [0u8; 4]; + let output = reader.read(&mut new_buf); + assert!(output.is_err()); + } +}