Skip to main content

lz_fear/framed/
decompress.rs

1use byteorder::{LE, ReadBytesExt};
2use std::hash::Hasher;
3use std::io::{self, Read, BufRead, ErrorKind};
4use std::cmp;
5use std::convert::TryInto;
6use twox_hash::XxHash32;
7use thiserror::Error;
8use culpa::{throw, throws};
9
10use super::{MAGIC, INCOMPRESSIBLE, WINDOW_SIZE};
11use super::header::{self, Flags, BlockDescriptor};
12use crate::raw;
13
14
15/// Errors when decompressing an LZ4 frame.
16#[derive(Error, Debug)]
17pub enum DecompressionError {
18    #[error("error reading from the input you gave me")]
19    InputError(#[from] io::Error),
20    #[error("the raw LZ4 decompression failed (data corruption?)")]
21    CodecError(#[from] raw::DecodeError),
22    #[error("invalid header")]
23    HeaderParseError(#[from] header::ParseError),
24    #[error("wrong magic number in file header: {0:08x}")]
25    WrongMagic(u32),
26    #[error("the header checksum was invalid")]
27    HeaderChecksumFail,
28    #[error("a block checksum was invalid")]
29    BlockChecksumFail,
30    #[error("the frame checksum was invalid")]
31    FrameChecksumFail,
32    #[error("stream contains a compressed block with a size so large we can't even compute it (let alone fit the block in memory...)")]
33    BlockLengthOverflow,
34    #[error("a block decompressed to more data than allowed")]
35    BlockSizeOverflow,
36}
37type Error = DecompressionError; // do it this way for better docs
38
39impl From<Error> for io::Error {
40    fn from(e: Error) -> io::Error {
41        io::Error::new(ErrorKind::Other, e)
42    }
43}
44
45/// Wrapper around `LZ4FrameReader` that implements `Read` and `BufRead`.
46pub struct LZ4FrameIoReader<'a, R: Read> {
47    frame_reader: LZ4FrameReader<R>,
48    bytes_taken: usize,
49    buffer: Vec<u8>,
50    dictionary: &'a [u8],
51}
52impl<R: Read> Read for LZ4FrameIoReader<'_, R> {
53    #[throws(io::Error)]
54    fn read(&mut self, buf: &mut [u8]) -> usize {
55        let mybuf = self.fill_buf()?;
56        let bytes_to_take = cmp::min(mybuf.len(), buf.len());
57        buf[..bytes_to_take].copy_from_slice(&mybuf[..bytes_to_take]);
58        self.consume(bytes_to_take);
59        bytes_to_take
60    }
61}
62impl<R: Read> BufRead for LZ4FrameIoReader<'_, R> {
63    #[throws(io::Error)]
64    fn fill_buf(&mut self) -> &[u8] {
65        if self.bytes_taken == self.buffer.len() {
66            self.buffer.clear();
67            self.frame_reader.decode_block(&mut self.buffer, self.dictionary)?;
68            self.bytes_taken = 0;
69        }
70        &self.buffer[self.bytes_taken..]
71    }
72
73    fn consume(&mut self, amt: usize) {
74        self.bytes_taken += amt;
75        assert!(self.bytes_taken <= self.buffer.len(), "You consumed more bytes than I even gave you!");
76    }
77}
78
79/// Read an LZ4-compressed frame.
80///
81/// This reader reads the blocks inside a frame one by one.
82pub struct LZ4FrameReader<R: Read> {
83    reader: R,
84    flags: Flags,
85    block_maxsize: usize,
86    read_buf: Vec<u8>,
87    content_size: Option<u64>,
88    dictionary_id: Option<u32>,
89    content_hasher: Option<XxHash32>,
90    carryover_window: Option<Vec<u8>>,
91    finished: bool,
92}
93
94impl<R: Read> LZ4FrameReader<R> {
95    /// Create a new LZ4FrameReader over an underlying reader and parse the header.
96    ///
97    /// A typical LZ4 file consists of exactly one frame.
98    /// This reader will stop reading at the end of the frame.
99    /// If you want to read any data following this frame, you should probably
100    /// pass in your reader by reference, rather than by value.
101    #[throws]
102    pub fn new(mut reader: R) -> Self {
103        let magic = reader.read_u32::<LE>()?;
104        if magic != MAGIC {
105            throw!(Error::WrongMagic(magic));
106        }
107
108        let flags_byte = reader.read_u8()?;
109        let flags = Flags::parse(flags_byte)?;
110        let bd = BlockDescriptor::parse(reader.read_u8()?)?;
111
112        let mut hasher = XxHash32::with_seed(0);
113        hasher.write_u8(flags_byte);
114        hasher.write_u8(bd.0);
115
116        let content_size = if flags.content_size() {
117            let i = reader.read_u64::<LE>()?;
118            hasher.write_u64(i);
119            Some(i)
120        } else {
121            None
122        };
123
124        let dictionary_id = if flags.dictionary_id() {
125            let i = reader.read_u32::<LE>()?;
126            hasher.write_u32(i);
127            Some(i)
128        } else {
129            None
130        };
131
132        let header_checksum_desired = reader.read_u8()?;
133        let header_checksum_actual = (hasher.finish() >> 8) as u8;
134        if header_checksum_desired != header_checksum_actual {
135            throw!(Error::HeaderChecksumFail);
136        }
137
138        let content_hasher = if flags.content_checksum() {
139            Some(XxHash32::with_seed(0))
140        } else {
141            None
142        };
143
144        let carryover_window = if flags.independent_blocks() {
145            None
146        } else {
147            Some(Vec::with_capacity(WINDOW_SIZE))
148        };
149
150        LZ4FrameReader {
151            reader,
152            flags,
153            block_maxsize: bd.block_maxsize()?,
154            content_size,
155            dictionary_id,
156            content_hasher,
157            carryover_window,
158            finished: false,
159            read_buf: Vec::new()
160        }
161    }
162
163    /// Returns the maximum number of bytes a block can decompress to (as specified by the file header).
164    ///
165    /// In general, all blocks in a frame except for the final one will have exactly this size.
166    /// (Although this is not strictly enforced and may be violated by hand-crafted inputs)
167    pub fn block_size(&self) -> usize { self.block_maxsize }
168    /// Returns the number of bytes that this entire frame is supposed to decompress to.
169    /// This value is read directly from the file header and may be incorrect for malicious inputs.
170    pub fn frame_size(&self) -> Option<u64> { self.content_size }
171    /// Return an identifier for the dictionary that was used to compress this frame.
172    ///
173    /// Dictionary identifiers are always application-specific. Note that the lz4 command line utility never
174    /// specifies a dictionary id, even if a dictionary was used.
175    pub fn dictionary_id(&self) -> Option<u32> { self.dictionary_id }
176
177    /// Convert this `LZ4FrameReader` into something that implements `std::io::BufRead`.
178    ///
179    /// Note that `io::copy` has a small performance issue: https://github.com/rust-lang/rust/issues/49921
180    pub fn into_read_with_dictionary(self, dictionary: &[u8]) -> LZ4FrameIoReader<R> {
181        LZ4FrameIoReader {
182            buffer: Vec::with_capacity(self.block_size()),
183            bytes_taken: 0,
184            frame_reader: self,
185            dictionary,
186        }
187    }
188
189    /// Convenience wrapper in case you don't want to specify a dictionary.
190    pub fn into_read(self) -> LZ4FrameIoReader<'static, R> {
191        self.into_read_with_dictionary(&[])
192    }
193
194    /// Decode a single block.
195    ///
196    /// The `output` buffer must be empty upon calling this method.
197    #[throws]
198    pub fn decode_block(&mut self, output: &mut Vec<u8>, dictionary: &[u8]) {
199        assert!(output.is_empty(), "You must pass an empty buffer to this interface.");
200        
201        if self.finished { return; }
202
203        let reader = &mut self.reader;
204
205        let block_length = reader.read_u32::<LE>()?;
206        if block_length == 0 {
207            if let Some(hasher) = self.content_hasher.take() {
208                let checksum = reader.read_u32::<LE>()?;
209                if hasher.finish() != checksum.into() {
210                    throw!(Error::FrameChecksumFail);
211                }
212            }
213            self.finished = true;
214            return;
215        }
216
217        let is_compressed = block_length & INCOMPRESSIBLE == 0;
218        let block_length = block_length & !INCOMPRESSIBLE;
219
220        if block_length > self.block_maxsize as u32 {
221            throw!(Error::BlockSizeOverflow);
222        }
223
224        let buf = &mut self.read_buf;
225        buf.resize(block_length.try_into().or(Err(Error::BlockLengthOverflow))?, 0);
226        reader.read_exact(buf.as_mut_slice())?;
227
228        if self.flags.block_checksums() {
229            let checksum = reader.read_u32::<LE>()?;
230            let mut hasher = XxHash32::with_seed(0);
231            hasher.write(buf);
232            if hasher.finish() != checksum.into() {
233                throw!(Error::BlockChecksumFail);
234            }
235        }
236
237        // set up the prefix properly
238        let dec_prefix = if let Some(window) = self.carryover_window.as_mut() {
239            if window.is_empty() {
240                window.extend_from_slice(dictionary);
241            }
242            window
243        } else {
244            dictionary
245        };
246        // decompress or copy, depending on whether this block is compressed
247        if is_compressed {
248            raw::decompress_raw(buf, dec_prefix, output, self.block_maxsize)?;
249        } else {
250            output.extend_from_slice(buf);
251        }
252        // finally, push data back into the window as needed
253        if let Some(window) = self.carryover_window.as_mut() {
254            let outlen = output.len();
255            if outlen < WINDOW_SIZE {
256                let available_bytes = window.len() + outlen;
257                if let Some(surplus_bytes) = available_bytes.checked_sub(WINDOW_SIZE) {
258                    // remove as many bytes from front as we are replacing
259                    window.drain(..surplus_bytes);
260                }
261                window.extend_from_slice(output);
262            } else {
263                // TODO: optimize this case to avoid the copy
264                window.clear();
265                window.extend_from_slice(&output[outlen - WINDOW_SIZE..]);
266            }
267
268            assert!(window.len() <= WINDOW_SIZE);
269        }
270
271
272        if output.len() > self.block_maxsize {
273            throw!(Error::BlockSizeOverflow);
274        }
275
276        if let Some(hasher) = self.content_hasher.as_mut() {
277            hasher.write(output);
278        }
279    }
280}
281
282/// Convenience wrapper around `LZ4FrameReader` that reads everything into a vector and returns it.
283#[throws]
284pub fn decompress_frame<R: Read>(reader: R) -> Vec<u8> {
285    let mut plaintext = Vec::new();
286    LZ4FrameReader::new(reader)?.into_read().read_to_end(&mut plaintext)?;
287    plaintext
288}
289