split compression code, return read bytes

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<Packet> 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,

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<Packet> 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<Packet> 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 {}

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;
 

src/compression.rs

@@ -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<Payload, CompressionError> {
-    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<Payload, CompressionError> {
-    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
-            })
-        }
-    }
-}

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<Payload, CompressionError> {
+        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<u8>, 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)
+            }
+        }
+    }
+}

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<u8>, usize), CompressionError> {
+        let mut output: Vec<u8> = 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<Payload, CompressionError> {
+        lzma::compress(payload, 6).map_err(|e| {
+            error!("failed to compress data: {e}");
+            CompressionError::CompressionFailed
+        })
+    }
+}

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<Vec<u8>, CompressionError>;
+    fn decompress(
+        payload: &[u8],
+        expected_size_hint: usize,
+    ) -> Result<(Vec<u8>, usize), CompressionError>;
+}
+
+pub(crate) fn decompress(
+    kind: CompressionCode,
+    #[allow(unused, reason = "depends on features")] payload: &[u8],
+    expected_size_hint: usize,
+) -> Result<(Vec<u8>, 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<Vec<u8>, 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),
+    }
+}

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<Vec<u8>, CompressionError> {
+        Ok(payload.to_vec())
+    }
+
+    fn decompress(
+        payload: &[u8],
+        expected_size_hint: usize,
+    ) -> Result<(Vec<u8>, usize), CompressionError> {
+        Ok((
+            payload[..expected_size_hint.min(payload.len())].to_vec(),
+            expected_size_hint,
+        ))
+    }
+}

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<Payload, CompressionError> {
+        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<u8>, 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)
+            }
+        }
+    }
+}

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<Payload, CompressionError> {
+        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<u8>, 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))
+    }
+}

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;

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>());
             u16::from_be_bytes([slice[0], slice[1]])

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,
+                }
+            })
         );
     }
 }