diff --git a/src/commands/bitmap.rs b/src/commands/bitmap.rs index 94c73b5..158e369 100644 --- a/src/commands/bitmap.rs +++ b/src/commands/bitmap.rs @@ -1,7 +1,7 @@ use crate::{ command_code::{CommandCode, InvalidCommandCodeError}, commands::errors::{TryFromPacketError, TryIntoPacketError}, - compression::{compress, decompress, CompressionError}, + compression::{compress, decompress}, Bitmap, CompressionCode, DataRef, Grid, Header, Origin, Packet, Pixels, TypedCommand, TILE_SIZE, }; @@ -51,7 +51,6 @@ impl TryFrom<&BitmapCommand> for Packet { let data_ref = value.bitmap.data_ref(); let payload = match compress(value.compression, data_ref) { Ok(payload) => payload, - Err(CompressionError::NoCompression) => data_ref.to_vec(), Err(_) => return Err(TryIntoPacketError::CompressionFailed), }; @@ -102,11 +101,11 @@ impl TryFrom for BitmapCommand { actual: 0, expected, })?; - let payload = match decompress(compression, &payload) { - Ok(payload) => payload, - Err(CompressionError::NoCompression) => payload, - Err(_) => return Err(TryFromPacketError::DecompressionFailed), - }; + let (payload, read_payload_bytes) = + match decompress(compression, &payload, expected) { + Ok(payload) => payload, + Err(_) => return Err(TryFromPacketError::DecompressionFailed), + }; let bitmap = Bitmap::load( tile_w as usize * TILE_SIZE, pixel_h as usize, diff --git a/src/commands/bitvec.rs b/src/commands/bitvec.rs index c3f86e4..d24076f 100644 --- a/src/commands/bitvec.rs +++ b/src/commands/bitvec.rs @@ -1,7 +1,7 @@ use crate::{ command_code::{CommandCode, InvalidCommandCodeError}, commands::errors::TryFromPacketError, - compression::{compress, decompress, CompressionError}, + compression::{compress, decompress}, CompressionCode, DisplayBitVec, Header, Offset, Packet, TryIntoPacketError, TypedCommand, }; @@ -68,7 +68,6 @@ impl TryFrom<&BitVecCommand> for Packet { let length = data_ref.len().try_into()?; let payload = match compress(value.compression, data_ref) { Ok(payload) => payload, - Err(CompressionError::NoCompression) => data_ref.to_vec(), Err(_) => return Err(TryIntoPacketError::CompressionFailed), }; Ok(Packet { @@ -100,6 +99,7 @@ impl TryFrom for BitVecCommand { }, payload, } = packet; + let expected_len = expected_len as usize; let command_code = CommandCode::try_from(command_code)?; let operation = match command_code { CommandCode::BitmapLinear => BinaryOperation::Overwrite, @@ -117,17 +117,17 @@ impl TryFrom for BitVecCommand { let compression = CompressionCode::try_from(sub)?; let payload = payload.ok_or(TryFromPacketError::UnexpectedPayloadSize { - expected: expected_len as usize, + expected: expected_len, actual: 0, })?; - let payload = match decompress(compression, &payload) { - Ok(payload) => payload, - Err(CompressionError::NoCompression) => payload.clone(), - Err(_) => return Err(TryFromPacketError::DecompressionFailed), - }; - if payload.len() != expected_len as usize { + let (payload, read_payload_bytes) = + match decompress(compression, &payload, expected_len) { + Ok(payload) => payload, + Err(_) => return Err(TryFromPacketError::DecompressionFailed), + }; + if payload.len() != expected_len { return Err(TryFromPacketError::UnexpectedPayloadSize { - expected: expected_len as usize, + expected: expected_len, actual: payload.len(), }); } @@ -162,7 +162,7 @@ mod tests { use super::*; use crate::{ commands, commands::tests::TestImplementsCommand, - compression_code::InvalidCompressionCodeError, PIXEL_WIDTH, + InvalidCompressionCodeError, PIXEL_WIDTH, }; impl TestImplementsCommand for BitVecCommand {} diff --git a/src/commands/errors.rs b/src/commands/errors.rs index bdb9b62..40dbbbc 100644 --- a/src/commands/errors.rs +++ b/src/commands/errors.rs @@ -1,6 +1,6 @@ use crate::{ - command_code::InvalidCommandCodeError, - compression_code::InvalidCompressionCodeError, LoadBitmapError, + command_code::InvalidCommandCodeError, InvalidCompressionCodeError, + LoadBitmapError, }; use std::num::TryFromIntError; diff --git a/src/compression.rs b/src/compression.rs deleted file mode 100644 index a8b8c9a..0000000 --- a/src/compression.rs +++ /dev/null @@ -1,177 +0,0 @@ -#[cfg(feature = "compression_bzip2")] -use bzip2::read::{BzDecoder, BzEncoder}; -#[cfg(feature = "compression_zlib")] -use flate2::{FlushCompress, FlushDecompress, Status}; -#[allow(unused, reason = "used depending on enabled features")] -use log::error; -#[allow(unused, reason = "used depending on enabled features")] -use std::io::{Read, Write}; -#[cfg(feature = "compression_zstd")] -use zstd::{Decoder as ZstdDecoder, Encoder as ZstdEncoder}; - -use crate::{CompressionCode, Payload}; - -#[derive(thiserror::Error, Debug, PartialEq)] -pub(crate) enum CompressionError { - #[error("Could not compress or decompress as no compression is used.")] - NoCompression, - #[error("Could not initialize compression library")] - #[allow(unused, reason = "depends on features")] - LibraryError, - #[error("Compression/decompression operation failed")] - #[allow(unused, reason = "depends on features")] - CompressionFailed, -} - -pub(crate) fn decompress( - kind: CompressionCode, - #[allow(unused, reason = "depends on features")] payload: &[u8], -) -> Result { - match kind { - CompressionCode::Uncompressed => Err(CompressionError::NoCompression), - #[cfg(feature = "compression_zlib")] - CompressionCode::Zlib => { - let mut decompress = flate2::Decompress::new(true); - let mut buffer = [0u8; 10000]; - - match decompress.decompress( - payload, - &mut buffer, - FlushDecompress::Finish, - ) { - Ok(Status::Ok) => { - error!("input not big enough"); - Err(CompressionError::CompressionFailed) - } - Ok(Status::BufError) => { - error!("output buffer is too small"); - Err(CompressionError::CompressionFailed) - } - Ok(Status::StreamEnd) => - { - #[allow( - clippy::cast_possible_truncation, - reason = "can never be larger than the fixed buffer size" - )] - Ok(buffer[..decompress.total_out() as usize].to_owned()) - } - Err(e) => { - error!("failed to decompress data: {e}"); - Err(CompressionError::CompressionFailed) - } - } - } - #[cfg(feature = "compression_bzip2")] - CompressionCode::Bzip2 => { - let mut decoder = BzDecoder::new(payload); - let mut decompressed = vec![]; - match decoder.read_to_end(&mut decompressed) { - Ok(_) => Ok(decompressed), - Err(e) => { - error!("failed to decompress data: {e}"); - Err(CompressionError::CompressionFailed) - } - } - } - #[cfg(feature = "compression_lzma")] - CompressionCode::Lzma => lzma::decompress(payload).map_err(|e| { - error!("failed to decompress data: {e}"); - CompressionError::CompressionFailed - }), - #[cfg(feature = "compression_zstd")] - CompressionCode::Zstd => { - let mut decoder = match ZstdDecoder::new(payload) { - Ok(value) => value, - Err(e) => { - error!("failed to create zstd decoder: {e}"); - return Err(CompressionError::LibraryError); - } - }; - let mut decompressed = vec![]; - match decoder.read_to_end(&mut decompressed) { - Err(e) => { - error!("failed to decompress data: {e}"); - Err(CompressionError::CompressionFailed) - } - Ok(_) => Ok(decompressed), - } - } - } -} - -pub(crate) fn compress( - kind: CompressionCode, - #[allow(unused, reason = "depends on features")] payload: &[u8], -) -> Result { - match kind { - CompressionCode::Uncompressed => Err(CompressionError::NoCompression), - #[cfg(feature = "compression_zlib")] - CompressionCode::Zlib => { - let mut compress = - flate2::Compress::new(flate2::Compression::fast(), true); - let mut buffer = [0u8; 10000]; - - match compress.compress(payload, &mut buffer, FlushCompress::Finish) - { - Ok(Status::Ok) => { - error!("output buffer not big enough"); - Err(CompressionError::CompressionFailed) - } - Ok(Status::BufError) => { - error!("Could not compress with buffer error"); - Err(CompressionError::CompressionFailed) - } - Ok(Status::StreamEnd) => - { - #[allow( - clippy::cast_possible_truncation, - reason = "can never be larger than the fixed buffer size" - )] - Ok(buffer[..compress.total_out() as usize].to_owned()) - } - Err(e) => { - error!("failed to compress data: {e}"); - Err(CompressionError::CompressionFailed) - } - } - } - #[cfg(feature = "compression_bzip2")] - CompressionCode::Bzip2 => { - let mut encoder = - BzEncoder::new(payload, bzip2::Compression::fast()); - let mut compressed = vec![]; - match encoder.read_to_end(&mut compressed) { - Ok(_) => Ok(compressed), - Err(e) => { - error!("failed to compress data: {e}"); - Err(CompressionError::CompressionFailed) - } - } - } - #[cfg(feature = "compression_lzma")] - CompressionCode::Lzma => lzma::compress(payload, 6).map_err(|e| { - error!("failed to compress data: {e}"); - CompressionError::CompressionFailed - }), - #[cfg(feature = "compression_zstd")] - CompressionCode::Zstd => { - let buf = Vec::with_capacity(payload.len()); - let mut encoder = - ZstdEncoder::new(buf, zstd::DEFAULT_COMPRESSION_LEVEL) - .map_err(|e| { - error!("failed to create zstd encoder: {e}"); - CompressionError::LibraryError - })?; - - if let Err(e) = encoder.write_all(payload) { - error!("failed to compress data: {e}"); - return Err(CompressionError::CompressionFailed); - } - - encoder.finish().map_err(|e| { - error!("failed to finish compression: {e}"); - CompressionError::CompressionFailed - }) - } - } -} diff --git a/src/compression/bzip2.rs b/src/compression/bzip2.rs new file mode 100644 index 0000000..0f9b93d --- /dev/null +++ b/src/compression/bzip2.rs @@ -0,0 +1,43 @@ +use crate::{ + compression::{CompressionAlgo, CompressionError}, + CompressionCode, Payload, +}; +use bzip2::{ + read::{BzDecoder, BzEncoder}, + Compression, +}; +use log::error; +use std::io::Read; + +pub struct Bzip2; + +impl CompressionAlgo for Bzip2 { + const CODE: CompressionCode = CompressionCode::Bzip2; + + fn compress(payload: &[u8]) -> Result { + let mut encoder = BzEncoder::new(payload, Compression::fast()); + let mut compressed = vec![]; + match encoder.read_to_end(&mut compressed) { + Ok(_) => Ok(compressed), + Err(e) => { + error!("failed to compress data: {e}"); + Err(CompressionError::CompressionFailed) + } + } + } + + fn decompress( + payload: &[u8], + expected_size_hint: usize, + ) -> Result<(Vec, usize), CompressionError> { + let mut decoder = BzDecoder::new(payload); + let mut decompressed = Vec::with_capacity(expected_size_hint); + match decoder.read_to_end(&mut decompressed) { + Ok(_) => Ok((decompressed, decoder.total_in() as usize)), + Err(e) => { + error!("failed to decompress data: {e}"); + Err(CompressionError::CompressionFailed) + } + } + } +} diff --git a/src/compression_code.rs b/src/compression/compression_code.rs similarity index 100% rename from src/compression_code.rs rename to src/compression/compression_code.rs diff --git a/src/compression/lzma.rs b/src/compression/lzma.rs new file mode 100644 index 0000000..8587095 --- /dev/null +++ b/src/compression/lzma.rs @@ -0,0 +1,31 @@ +use crate::{ + compression::{CompressionAlgo, CompressionError}, + CompressionCode, Payload, +}; +use log::error; +use lzma::LzmaReader; +use std::io::Read; + +pub struct Lzma; + +impl CompressionAlgo for Lzma { + const CODE: CompressionCode = CompressionCode::Lzma; + fn decompress( + payload: &[u8], + _: usize, + ) -> Result<(Vec, usize), CompressionError> { + let mut output: Vec = Vec::new(); + let mut reader = LzmaReader::new_decompressor(payload) + .map_err(|_| CompressionError::LibraryError)?; + let read = reader + .read_to_end(&mut output) + .map_err(|_| CompressionError::CompressionFailed)?; + Ok((output, read)) + } + fn compress(payload: &[u8]) -> Result { + lzma::compress(payload, 6).map_err(|e| { + error!("failed to compress data: {e}"); + CompressionError::CompressionFailed + }) + } +} diff --git a/src/compression/mod.rs b/src/compression/mod.rs new file mode 100644 index 0000000..e0ff964 --- /dev/null +++ b/src/compression/mod.rs @@ -0,0 +1,80 @@ +#[cfg(feature = "compression_bzip2")] +mod bzip2; +mod compression_code; +#[cfg(feature = "compression_lzma")] +mod lzma; +mod uncompressed; +#[cfg(feature = "compression_zlib")] +mod zlib; +#[cfg(feature = "compression_zstd")] +mod zstd; + +pub use compression_code::{CompressionCode, InvalidCompressionCodeError}; + +#[derive(thiserror::Error, Debug, PartialEq)] +pub(crate) enum CompressionError { + #[error("Could not initialize compression library")] + #[allow(unused, reason = "depends on features")] + LibraryError, + #[error("Compression/decompression operation failed")] + #[allow(unused, reason = "depends on features")] + CompressionFailed, +} + +// TODO: while were at it, why not decompress into an existing buffer? + +pub(crate) trait CompressionAlgo { + const CODE: CompressionCode; + fn compress(payload: &[u8]) -> Result, CompressionError>; + fn decompress( + payload: &[u8], + expected_size_hint: usize, + ) -> Result<(Vec, usize), CompressionError>; +} + +pub(crate) fn decompress( + kind: CompressionCode, + #[allow(unused, reason = "depends on features")] payload: &[u8], + expected_size_hint: usize, +) -> Result<(Vec, usize), CompressionError> { + match kind { + CompressionCode::Uncompressed => { + uncompressed::Uncompressed::decompress(payload, expected_size_hint) + } + #[cfg(feature = "compression_zlib")] + CompressionCode::Zlib => { + zlib::Zlib::decompress(payload, expected_size_hint) + } + #[cfg(feature = "compression_bzip2")] + CompressionCode::Bzip2 => { + bzip2::Bzip2::decompress(payload, expected_size_hint) + } + #[cfg(feature = "compression_lzma")] + CompressionCode::Lzma => { + lzma::Lzma::decompress(payload, expected_size_hint) + } + #[cfg(feature = "compression_zstd")] + CompressionCode::Zstd => { + zstd::Zstd::decompress(payload, expected_size_hint) + } + } +} + +pub(crate) fn compress( + kind: CompressionCode, + #[allow(unused, reason = "depends on features")] payload: &[u8], +) -> Result, CompressionError> { + match kind { + CompressionCode::Uncompressed => { + uncompressed::Uncompressed::compress(payload) + } + #[cfg(feature = "compression_zlib")] + CompressionCode::Zlib => zlib::Zlib::compress(payload), + #[cfg(feature = "compression_bzip2")] + CompressionCode::Bzip2 => bzip2::Bzip2::compress(payload), + #[cfg(feature = "compression_lzma")] + CompressionCode::Lzma => lzma::Lzma::compress(payload), + #[cfg(feature = "compression_zstd")] + CompressionCode::Zstd => zstd::Zstd::compress(payload), + } +} diff --git a/src/compression/uncompressed.rs b/src/compression/uncompressed.rs new file mode 100644 index 0000000..dd6cef2 --- /dev/null +++ b/src/compression/uncompressed.rs @@ -0,0 +1,24 @@ +use crate::{ + compression::{CompressionAlgo, CompressionError}, + CompressionCode, +}; + +pub struct Uncompressed; + +impl CompressionAlgo for Uncompressed { + const CODE: CompressionCode = CompressionCode::Uncompressed; + + fn compress(payload: &[u8]) -> Result, CompressionError> { + Ok(payload.to_vec()) + } + + fn decompress( + payload: &[u8], + expected_size_hint: usize, + ) -> Result<(Vec, usize), CompressionError> { + Ok(( + payload[..expected_size_hint.min(payload.len())].to_vec(), + expected_size_hint, + )) + } +} diff --git a/src/compression/zlib.rs b/src/compression/zlib.rs new file mode 100644 index 0000000..d1114bd --- /dev/null +++ b/src/compression/zlib.rs @@ -0,0 +1,76 @@ +use crate::{ + compression::{CompressionAlgo, CompressionError}, + CompressionCode, Payload, +}; +use flate2::{FlushCompress, FlushDecompress, Status}; +use log::error; + +pub struct Zlib; + +impl CompressionAlgo for Zlib { + const CODE: CompressionCode = CompressionCode::Zlib; + + fn compress(payload: &[u8]) -> Result { + let mut compress = + flate2::Compress::new(flate2::Compression::fast(), true); + let mut buffer = [0u8; 10000]; + + match compress.compress(payload, &mut buffer, FlushCompress::Finish) { + Ok(Status::Ok) => { + error!("output buffer not big enough"); + Err(CompressionError::CompressionFailed) + } + Ok(Status::BufError) => { + error!("Could not compress with buffer error"); + Err(CompressionError::CompressionFailed) + } + Ok(Status::StreamEnd) => + { + #[allow( + clippy::cast_possible_truncation, + reason = "can never be larger than the fixed buffer size" + )] + Ok(buffer[..compress.total_out() as usize].to_owned()) + } + Err(e) => { + error!("failed to compress data: {e}"); + Err(CompressionError::CompressionFailed) + } + } + } + + fn decompress( + payload: &[u8], + expected_size_hint: usize, + ) -> Result<(Vec, usize), CompressionError> { + let mut buffer = Vec::with_capacity(expected_size_hint); + let mut instance = flate2::Decompress::new(true); + match instance.decompress_vec( + payload, + &mut buffer, + FlushDecompress::Finish, + ) { + Ok(Status::Ok) => { + error!("input not big enough"); + Err(CompressionError::CompressionFailed) + } + Ok(Status::BufError) => { + error!("output buffer is too small"); + Err(CompressionError::CompressionFailed) + } + Ok(Status::StreamEnd) => { + #[allow( + clippy::cast_possible_truncation, + reason = "can never be larger than the fixed buffer size" + )] + let result_buf = + buffer[..instance.total_out() as usize].to_owned(); + Ok((result_buf, instance.total_in() as usize)) + } + Err(e) => { + error!("failed to decompress data: {e}"); + Err(CompressionError::CompressionFailed) + } + } + } +} diff --git a/src/compression/zstd.rs b/src/compression/zstd.rs new file mode 100644 index 0000000..4d5df16 --- /dev/null +++ b/src/compression/zstd.rs @@ -0,0 +1,49 @@ +use crate::{ + compression::{CompressionAlgo, CompressionError}, + CompressionCode, Payload, +}; +use log::error; +use std::io::{Read, Write}; +use zstd::{Decoder, Encoder, DEFAULT_COMPRESSION_LEVEL}; + +pub struct Zstd; + +impl CompressionAlgo for Zstd { + const CODE: CompressionCode = CompressionCode::Zstd; + fn compress(payload: &[u8]) -> Result { + let buf = Vec::with_capacity(payload.len()); + let mut encoder = Encoder::new(buf, DEFAULT_COMPRESSION_LEVEL) + .map_err(|e| { + error!("failed to create zstd encoder: {e}"); + CompressionError::LibraryError + })?; + + if let Err(e) = encoder.write_all(payload) { + error!("failed to compress data: {e}"); + return Err(CompressionError::CompressionFailed); + } + + encoder.finish().map_err(|e| { + error!("failed to finish compression: {e}"); + CompressionError::CompressionFailed + }) + } + fn decompress( + payload: &[u8], + expected_size_hint: usize, + ) -> Result<(Vec, usize), CompressionError> { + let mut decoder = match Decoder::new(payload) { + Ok(value) => value, + Err(e) => { + error!("failed to create zstd decoder: {e}"); + return Err(CompressionError::LibraryError); + } + }; + let mut decompressed = Vec::with_capacity(expected_size_hint); + let read_bytes = decoder + .read_to_end(&mut decompressed) + .inspect_err(|e| error!("failed to decompress data: {e}")) + .map_err(|_| CompressionError::CompressionFailed)?; + Ok((decompressed, read_bytes)) + } +} diff --git a/src/lib.rs b/src/lib.rs index 6d35cc4..1524973 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -81,7 +81,7 @@ pub use crate::brightness::Brightness; pub use crate::command_code::CommandCode; pub use crate::commands::*; -pub use crate::compression_code::CompressionCode; +pub use crate::compression::*; pub use crate::connection::*; pub use crate::constants::*; pub use crate::containers::*; @@ -92,7 +92,6 @@ mod brightness; mod command_code; mod commands; mod compression; -mod compression_code; mod connection; mod constants; mod containers; diff --git a/src/packet/header.rs b/src/packet/header.rs index c77575f..f10311d 100644 --- a/src/packet/header.rs +++ b/src/packet/header.rs @@ -23,8 +23,6 @@ pub struct Header { impl Header { pub fn read_from(value: &[u8]) -> Option<(Header, usize)> { - const NULL_U16_SLICE: &[u8] = &[0, 0]; - fn convert(slice: &[u8]) -> u16 { debug_assert_eq!(slice.len(), size_of::()); u16::from_be_bytes([slice[0], slice[1]]) diff --git a/src/packet/packet.rs b/src/packet/packet.rs index 32289ac..f138e1e 100644 --- a/src/packet/packet.rs +++ b/src/packet/packet.rs @@ -189,11 +189,20 @@ mod tests { } #[test] - fn too_small() { + fn small() { let data = vec![0u8; 4]; assert_eq!( Packet::try_from(data.as_slice()), - Err(SliceSmallerThanHeader) + Ok(Packet { + payload: None, + header: Header { + command_code: 0, + a: 0, + b: 0, + c: 0, + d: 0, + } + }) ); } }