distant_net/common/transport/framed/codec/
compression.rs

1use std::io::{self, Read};
2
3use flate2::bufread::{
4    DeflateDecoder, DeflateEncoder, GzDecoder, GzEncoder, ZlibDecoder, ZlibEncoder,
5};
6use flate2::Compression;
7use serde::{Deserialize, Serialize};
8
9use super::{Codec, Frame};
10
11/// Represents the level of compression to apply to data
12#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
13pub enum CompressionLevel {
14    /// Use no compression (can potentially inflate data)
15    Zero = 0,
16
17    /// Optimize for the speed of encoding
18    One = 1,
19
20    Two = 2,
21    Three = 3,
22    Four = 4,
23    Five = 5,
24    Six = 6,
25    Seven = 7,
26    Eight = 8,
27
28    /// Optimize for the size of data being encoded
29    Nine = 9,
30}
31
32impl CompressionLevel {
33    /// Applies best compression to reduce size (slowest)
34    pub const BEST: Self = Self::Nine;
35    /// Applies fastest compression
36    pub const FAST: Self = Self::One;
37    /// Applies no compression
38    pub const NONE: Self = Self::Zero;
39}
40
41impl Default for CompressionLevel {
42    /// Standard compression level used in zlib library is 6, which is also used here
43    fn default() -> Self {
44        Self::Six
45    }
46}
47
48/// Represents the type of compression for a [`CompressionCodec`]
49#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
50pub enum CompressionType {
51    Deflate,
52    Gzip,
53    Zlib,
54
55    /// Indicates an unknown compression type for use in handshakes
56    #[serde(other)]
57    Unknown,
58}
59
60impl CompressionType {
61    /// Returns a list of all variants of the type *except* unknown.
62    pub const fn known_variants() -> &'static [CompressionType] {
63        &[
64            CompressionType::Deflate,
65            CompressionType::Gzip,
66            CompressionType::Zlib,
67        ]
68    }
69
70    /// Returns true if type is unknown
71    pub fn is_unknown(&self) -> bool {
72        matches!(self, Self::Unknown)
73    }
74
75    /// Creates a new [`CompressionCodec`] for this type, failing if this type is unknown
76    pub fn new_codec(&self, level: CompressionLevel) -> io::Result<CompressionCodec> {
77        CompressionCodec::from_type_and_level(*self, level)
78    }
79}
80
81/// Represents a codec that applies compression during encoding and decompression during decoding
82/// of a frame's item
83#[derive(Copy, Clone, Debug, PartialEq, Eq)]
84pub enum CompressionCodec {
85    /// Apply DEFLATE compression/decompression using compression `level`
86    Deflate { level: CompressionLevel },
87
88    /// Apply gzip compression/decompression using compression `level`
89    Gzip { level: CompressionLevel },
90
91    /// Apply zlib compression/decompression using compression `level`
92    Zlib { level: CompressionLevel },
93}
94
95impl CompressionCodec {
96    /// Makes a new [`CompressionCodec`] based on the [`CompressionType`] and [`CompressionLevel`],
97    /// returning error if the type is unknown
98    pub fn from_type_and_level(
99        ty: CompressionType,
100        level: CompressionLevel,
101    ) -> io::Result<CompressionCodec> {
102        match ty {
103            CompressionType::Deflate => Ok(Self::Deflate { level }),
104            CompressionType::Gzip => Ok(Self::Gzip { level }),
105            CompressionType::Zlib => Ok(Self::Zlib { level }),
106            CompressionType::Unknown => Err(io::Error::new(
107                io::ErrorKind::InvalidInput,
108                "Unknown compression type",
109            )),
110        }
111    }
112
113    /// Create a new deflate compression codec with the specified `level`
114    pub fn deflate(level: impl Into<CompressionLevel>) -> Self {
115        Self::Deflate {
116            level: level.into(),
117        }
118    }
119
120    /// Create a new gzip compression codec with the specified `level`
121    pub fn gzip(level: impl Into<CompressionLevel>) -> Self {
122        Self::Gzip {
123            level: level.into(),
124        }
125    }
126
127    /// Create a new zlib compression codec with the specified `level`
128    pub fn zlib(level: impl Into<CompressionLevel>) -> Self {
129        Self::Zlib {
130            level: level.into(),
131        }
132    }
133
134    /// Returns the compression level associated with the codec
135    pub fn level(&self) -> CompressionLevel {
136        match self {
137            Self::Deflate { level } => *level,
138            Self::Gzip { level } => *level,
139            Self::Zlib { level } => *level,
140        }
141    }
142
143    /// Returns the compression type associated with the codec
144    pub fn ty(&self) -> CompressionType {
145        match self {
146            Self::Deflate { .. } => CompressionType::Deflate,
147            Self::Gzip { .. } => CompressionType::Gzip,
148            Self::Zlib { .. } => CompressionType::Zlib,
149        }
150    }
151}
152
153impl Codec for CompressionCodec {
154    fn encode<'a>(&mut self, frame: Frame<'a>) -> io::Result<Frame<'a>> {
155        let item = frame.as_item();
156
157        let mut buf = Vec::new();
158        match *self {
159            Self::Deflate { level } => {
160                DeflateEncoder::new(item, Compression::new(level as u32)).read_to_end(&mut buf)?
161            }
162            Self::Gzip { level } => {
163                GzEncoder::new(item, Compression::new(level as u32)).read_to_end(&mut buf)?
164            }
165            Self::Zlib { level } => {
166                ZlibEncoder::new(item, Compression::new(level as u32)).read_to_end(&mut buf)?
167            }
168        };
169
170        Ok(Frame::from(buf))
171    }
172
173    fn decode<'a>(&mut self, frame: Frame<'a>) -> io::Result<Frame<'a>> {
174        let item = frame.as_item();
175
176        let mut buf = Vec::new();
177        match *self {
178            Self::Deflate { .. } => DeflateDecoder::new(item).read_to_end(&mut buf)?,
179            Self::Gzip { .. } => GzDecoder::new(item).read_to_end(&mut buf)?,
180            Self::Zlib { .. } => ZlibDecoder::new(item).read_to_end(&mut buf)?,
181        };
182
183        Ok(Frame::from(buf))
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use test_log::test;
190
191    use super::*;
192
193    #[test]
194    fn encode_should_apply_appropriate_compression_algorithm() {
195        // Encode using DEFLATE and verify that the compression was as expected by decompressing
196        let mut codec = CompressionCodec::deflate(CompressionLevel::BEST);
197        let frame = codec.encode(Frame::new(b"some bytes")).unwrap();
198
199        let mut item = Vec::new();
200        DeflateDecoder::new(frame.as_item())
201            .read_to_end(&mut item)
202            .unwrap();
203        assert_eq!(item, b"some bytes");
204
205        // Encode using gzip and verify that the compression was as expected by decompressing
206        let mut codec = CompressionCodec::gzip(CompressionLevel::BEST);
207        let frame = codec.encode(Frame::new(b"some bytes")).unwrap();
208
209        let mut item = Vec::new();
210        GzDecoder::new(frame.as_item())
211            .read_to_end(&mut item)
212            .unwrap();
213        assert_eq!(item, b"some bytes");
214
215        // Encode using zlib and verify that the compression was as expected by decompressing
216        let mut codec = CompressionCodec::zlib(CompressionLevel::BEST);
217        let frame = codec.encode(Frame::new(b"some bytes")).unwrap();
218
219        let mut item = Vec::new();
220        ZlibDecoder::new(frame.as_item())
221            .read_to_end(&mut item)
222            .unwrap();
223        assert_eq!(item, b"some bytes");
224    }
225
226    #[test]
227    fn decode_should_apply_appropriate_decompression_algorithm() {
228        // Decode using DEFLATE
229        let frame = {
230            let mut item = Vec::new();
231            DeflateEncoder::new(b"some bytes".as_slice(), Compression::best())
232                .read_to_end(&mut item)
233                .unwrap();
234            Frame::from(item)
235        };
236        let mut codec = CompressionCodec::deflate(CompressionLevel::BEST);
237        let frame = codec.decode(frame).unwrap();
238        assert_eq!(frame, b"some bytes");
239
240        // Decode using gzip
241        let frame = {
242            let mut item = Vec::new();
243            GzEncoder::new(b"some bytes".as_slice(), Compression::best())
244                .read_to_end(&mut item)
245                .unwrap();
246            Frame::from(item)
247        };
248        let mut codec = CompressionCodec::gzip(CompressionLevel::BEST);
249        let frame = codec.decode(frame).unwrap();
250        assert_eq!(frame, b"some bytes");
251
252        // Decode using zlib
253        let frame = {
254            let mut item = Vec::new();
255            ZlibEncoder::new(b"some bytes".as_slice(), Compression::best())
256                .read_to_end(&mut item)
257                .unwrap();
258            Frame::from(item)
259        };
260        let mut codec = CompressionCodec::zlib(CompressionLevel::BEST);
261        let frame = codec.decode(frame).unwrap();
262        assert_eq!(frame, b"some bytes");
263    }
264}
distant_net/common/transport/framed/codec/compression.rs

distant_net/common/transport/framed/codec/
compression.rs
