runefs/
codec.rs

1//! (De)compression and enciphering/deciphering.
2
3use std::convert::TryFrom;
4#[cfg(feature = "rs3")]
5use std::io::BufReader;
6use std::io::{self, Read, Write};
7
8use bzip2::{read::BzDecoder, write::BzEncoder};
9use flate2::{bufread::GzDecoder, write::GzEncoder};
10#[cfg(feature = "rs3")]
11use lzma_rs::{compress, decompress, lzma_compress_with_options, lzma_decompress_with_options};
12use nom::{
13    combinator::cond,
14    number::complete::{be_i16, be_u32, be_u8},
15};
16
17use crate::{error::CompressionUnsupported, xtea};
18
19use std::marker::PhantomData;
20
21/// Supported compression types.
22#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug)]
23pub enum Compression {
24    None,
25    Bzip2,
26    Gzip,
27    #[cfg(feature = "rs3")]
28    #[cfg_attr(docsrs, doc(cfg(feature = "rs3")))]
29    Lzma,
30}
31
32/// Marker struct conveying `State` of a [`Buffer`](Buffer).
33pub struct Encoded;
34/// Marker struct conveying `State` of a [`Buffer`](Buffer).
35pub struct Decoded;
36
37/// Primary way to store bytes for encoding and decoding.
38pub struct Buffer<State> {
39    compression: Compression,
40    buffer: Vec<u8>,
41    version: Option<i16>,
42    keys: Option<[u32; 4]>,
43    _state: PhantomData<State>,
44}
45
46impl Buffer<Decoded> {
47    /// Encodes the buffer, consuming self and returning a `Buffer<Encoded>`.
48    ///
49    /// The following process takes place when encoding:
50    /// 1. Compress the buffer with the selected compression format.
51    /// 2. Allocate a new buffer.
52    /// 3. Push the compression type as a byte into the new buffer.
53    /// 4. Push the length (u32) into the buffer of the compressed data from step 1.
54    /// 5. If a compression type was selected (and not `Compression::None`) insert the uncompressed length as u32.
55    /// 6. Extend the buffer with the compressed data.
56    /// 7. Add the `version` as i16 if present.
57    /// 8. Encode complete.
58    ///
59    /// **NOTE: When compressing with gzip the header is removed
60    /// before the compressed data is returned.
61    /// The encoded buffer will not contain the gzip header.**
62    ///
63    /// # Errors
64    ///
65    /// Returns an error if the data couldn't be compressed or is invalid.
66    pub fn encode(self) -> crate::Result<Buffer<Encoded>> {
67        let decompressed_len = self.buffer.len();
68        let mut compressed_data = match self.compression {
69            Compression::None => self.buffer,
70            Compression::Bzip2 => compress_bzip2(&self.buffer)?,
71            Compression::Gzip => compress_gzip(&self.buffer)?,
72            #[cfg(feature = "rs3")]
73            Compression::Lzma => compress_lzma(&self.buffer)?,
74        };
75        if let Some(keys) = &self.keys {
76            xtea::encipher(&mut compressed_data, keys);
77        }
78        let mut buffer = Vec::with_capacity(compressed_data.len() + 11);
79        buffer.write_all(&[self.compression as u8])?;
80        buffer.write_all(&u32::to_be_bytes(compressed_data.len() as u32))?;
81        if self.compression != Compression::None {
82            buffer.write_all(&u32::to_be_bytes(decompressed_len as u32))?;
83        }
84        buffer.extend(compressed_data);
85        if let Some(version) = self.version {
86            buffer.write_all(&i16::to_be_bytes(version))?;
87        }
88
89        Ok(Buffer {
90            compression: self.compression,
91            buffer,
92            version: self.version,
93            keys: self.keys,
94            _state: PhantomData,
95        })
96    }
97}
98
99impl Buffer<Encoded> {
100    /// Decodes the buffer, consuming self and returning a `Buffer<Decoded>`.
101    ///
102    /// The following process takes place when decoding:
103    /// 1. Read the first byte to determine which compression type should be used to decompress.
104    /// 2. Read the length of the rest of the buffer.
105    /// 3. Decompress the remaining bytes.
106    ///
107    /// # Panics
108    /// 
109    /// When data can't be decompressed using LZMA this function panics.
110    /// 
111    /// # Errors
112    ///
113    /// Returns an error if the remaining bytes couldn't be decompressed.
114    pub fn decode(self) -> crate::Result<Buffer<Decoded>> {
115        let (buffer, compression) = be_u8(self.buffer.as_slice())?;
116        let compression = Compression::try_from(compression)?;
117
118        let (buffer, compressed_len) = be_u32(buffer)?;
119        let compressed_len = compressed_len as usize;
120
121        let mut buffer = std::borrow::Cow::from(buffer);
122        if let Some(keys) = self.keys {
123            xtea::decipher(buffer.to_mut(), &keys);
124        }
125
126        let (version, buffer) = match compression {
127            Compression::None => decompress_none(&buffer, compressed_len)?,
128            Compression::Bzip2 => decompress_bzip2(&buffer, compressed_len)?,
129            Compression::Gzip => decompress_gzip(&buffer, compressed_len)?,
130            #[cfg(feature = "rs3")]
131            Compression::Lzma => decompress_lzma(&buffer, compressed_len)?,
132        };
133
134        Ok(Buffer {
135            compression,
136            buffer,
137            version,
138            keys: self.keys,
139            _state: PhantomData,
140        })
141    }
142}
143
144impl<State> Buffer<State> {
145    /// Set the compression format for this buffer returning a new instance of `Self`.
146    pub fn with_compression(mut self, compression: Compression) -> Self {
147        self.compression = compression;
148        self
149    }
150
151    /// Set the version for this buffer returning a new instance of `Self`.
152    pub fn with_version(mut self, version: i16) -> Self {
153        self.version = Some(version);
154        self
155    }
156
157    /// Set the xtea keys for this buffer returning a new instance of `Self`.
158    pub fn with_xtea_keys(mut self, keys: [u32; 4]) -> Self {
159        self.keys = Some(keys);
160        self
161    }
162
163    /// Convert the `Buffer` with its current state into a raw `Vec<u8>`.
164    #[inline]
165    pub fn finalize(self) -> Vec<u8> {
166        self.buffer
167    }
168}
169
170impl<State> Default for Buffer<State> {
171    fn default() -> Self {
172        Self {
173            compression: Compression::None,
174            buffer: Vec::new(),
175            version: None,
176            keys: None,
177            _state: PhantomData,
178        }
179    }
180}
181
182impl<State> std::fmt::Debug for Buffer<State> {
183    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
184        f.debug_struct("Buffer")
185            .field("compression", &self.compression)
186            .field("keys", &self.keys)
187            .field("version", &self.version)
188            .field("buffer", &self.buffer)
189            .finish()
190    }
191}
192
193impl<State> From<&[u8]> for Buffer<State> {
194    fn from(buffer: &[u8]) -> Self {
195        Self {
196            buffer: Vec::from(buffer),
197            ..Self::default()
198        }
199    }
200}
201
202impl<State> From<Vec<u8>> for Buffer<State> {
203    fn from(buffer: Vec<u8>) -> Self {
204        Self {
205            buffer,
206            ..Self::default()
207        }
208    }
209}
210
211impl<State> std::ops::Deref for Buffer<State> {
212    type Target = Vec<u8>;
213
214    #[inline]
215    fn deref(&self) -> &Self::Target {
216        &self.buffer
217    }
218}
219
220impl<State> std::ops::DerefMut for Buffer<State> {
221    #[inline]
222    fn deref_mut(&mut self) -> &mut Self::Target {
223        &mut self.buffer
224    }
225}
226
227impl<State> std::convert::AsRef<[u8]> for Buffer<State> {
228    #[inline]
229    fn as_ref(&self) -> &[u8] {
230        self.buffer.as_slice()
231    }
232}
233
234impl<State> std::io::Write for Buffer<State> {
235    fn write(&mut self, buffer: &[u8]) -> io::Result<usize> {
236        self.buffer.write(buffer)
237    }
238
239    fn flush(&mut self) -> io::Result<()> {
240        self.buffer.flush()
241    }
242}
243
244fn compress_bzip2(data: &[u8]) -> io::Result<Vec<u8>> {
245    let mut compressor = BzEncoder::new(Vec::with_capacity(data.len()), bzip2::Compression::fast());
246    compressor.write_all(data)?;
247    let mut compressed_data = compressor.finish()?;
248    compressed_data.drain(..4);
249
250    Ok(compressed_data)
251}
252
253fn compress_gzip(data: &[u8]) -> io::Result<Vec<u8>> {
254    let mut compressor =
255        GzEncoder::new(Vec::with_capacity(data.len()), flate2::Compression::best());
256    compressor.write_all(data)?;
257    let compressed_data: Vec<u8> = compressor.finish()?;
258
259    Ok(compressed_data)
260}
261
262#[cfg(feature = "rs3")]
263fn compress_lzma(data: &[u8]) -> io::Result<Vec<u8>> {
264    let mut input = std::io::BufReader::new(data);
265    let mut output = Vec::with_capacity(data.len());
266    let options = compress::Options {
267        unpacked_size: compress::UnpackedSize::SkipWritingToHeader,
268    };
269
270    lzma_compress_with_options(&mut input, &mut output, &options)?;
271
272    Ok(output)
273}
274
275fn decompress_none(buffer: &[u8], len: usize) -> crate::Result<(Option<i16>, Vec<u8>)> {
276    let (buffer, data) = nom::bytes::complete::take(len)(buffer)?;
277    let (_, version) = cond(buffer.len() >= 2, be_i16)(buffer)?;
278
279    Ok((version, data.to_vec()))
280}
281
282fn decompress_bzip2(buffer: &[u8], len: usize) -> crate::Result<(Option<i16>, Vec<u8>)> {
283    let (buffer, decompressed_len) = be_u32(buffer)?;
284    let (buffer, data) = nom::bytes::complete::take(len)(buffer)?;
285    let (_, version) = cond(buffer.len() >= 2, be_i16)(buffer)?;
286
287    let mut compressed_data = data.to_vec();
288    compressed_data[4..len].copy_from_slice(&data[..len - 4]);
289    compressed_data[..4].copy_from_slice(b"BZh1");
290
291    let mut decompressor = BzDecoder::new(compressed_data.as_slice());
292    let mut decompressed_data = vec![0; decompressed_len as usize];
293    decompressor.read_exact(&mut decompressed_data)?;
294
295    Ok((version, decompressed_data))
296}
297
298fn decompress_gzip(buffer: &[u8], len: usize) -> crate::Result<(Option<i16>, Vec<u8>)> {
299    let (buffer, decompressed_len) = be_u32(buffer)?;
300    let (buffer, data) = nom::bytes::complete::take(len)(buffer)?;
301    let (_, version) = cond(buffer.len() >= 2, be_i16)(buffer)?;
302
303    let mut decompressor = GzDecoder::new(data);
304    let mut decompressed_data = vec![0; decompressed_len as usize];
305    decompressor.read_exact(&mut decompressed_data)?;
306
307    Ok((version, decompressed_data))
308}
309
310#[cfg(feature = "rs3")]
311fn decompress_lzma(buffer: &[u8], len: usize) -> crate::Result<(Option<i16>, Vec<u8>)> {
312    let (buffer, decompressed_len) = be_u32(buffer)?;
313    let (buffer, data) = nom::bytes::complete::take(len)(buffer)?;
314    let (_, version) = cond(buffer.len() >= 2, be_i16)(buffer)?;
315
316    let mut decompressed_data = Vec::with_capacity(decompressed_len as usize);
317    let mut wrapper = BufReader::new(data);
318    let options = decompress::Options {
319        unpacked_size: decompress::UnpackedSize::UseProvided(Some(decompressed_len as u64)),
320        ..decompress::Options::default()
321    };
322
323    lzma_decompress_with_options(&mut wrapper, &mut decompressed_data, &options).unwrap();
324
325    Ok((version, decompressed_data))
326}
327
328impl Default for Compression {
329    #[inline]
330    fn default() -> Self {
331        Self::None
332    }
333}
334
335impl From<Compression> for u8 {
336    fn from(compression: Compression) -> Self {
337        match compression {
338            Compression::None => 0,
339            Compression::Bzip2 => 1,
340            Compression::Gzip => 2,
341            #[cfg(feature = "rs3")]
342            Compression::Lzma => 3,
343        }
344    }
345}
346
347impl std::convert::TryFrom<u8> for Compression {
348    type Error = CompressionUnsupported;
349
350    fn try_from(compression: u8) -> Result<Self, Self::Error> {
351        match compression {
352            0 => Ok(Self::None),
353            1 => Ok(Self::Bzip2),
354            2 => Ok(Self::Gzip),
355            #[cfg(feature = "rs3")]
356            3 => Ok(Self::Lzma),
357            _ => Err(CompressionUnsupported(compression)),
358        }
359    }
360}
runefs/codec.rs

runefs/
codec.rs
