Skip to main content

rpfm_lib/compression/
mod.rs

1//---------------------------------------------------------------------------//
2// Copyright (c) 2017-2026 Ismael Gutiérrez González. All rights reserved.
3//
4// This file is part of the Rusted PackFile Manager (RPFM) project,
5// which can be found here: https://github.com/Frodo45127/rpfm.
6//
7// This file is licensed under the MIT license, which can be found here:
8// https://github.com/Frodo45127/rpfm/blob/master/LICENSE.
9//---------------------------------------------------------------------------//
10
11//! This module contains the code to compress/decompress data for Total War games.
12//!
13//! The traits [`Compressible`] and [`Decompressible`] within this module contain functions to compress and decompress
14//! data from/to CA's different supported compression formats. Implementations of these two traits for &[[`u8`]] are provided within this module.
15//!
16//! # Supported Formats
17//!
18//! See [`CompressionFormat`] for details on the supported compression formats (LZMA1, Lz4, Zstd, Lz4Block) and their file structure.
19//!
20//! # Important Notes
21//!
22//! * Due to an game bug, compressing tables tends to cause crashes when starting for some people. This bug seems to have been fixed in WH3, but all other games before WH3
23//!   may still suffer from it, so unless manually forced to, this lib will not compress tables in those games. Tables will only be compressed in WH3 and newer games.
24//!
25//! * Compressed files are **only supported on PFH5 and newer Packs** (Since Total War: Warhammer 2), and on **PFH3 Packs from Shogun 2 Complete Edition**
26//!   (only with [`CompressionFormat::Lz4Block`]).
27
28use lz4_flex::frame::{FrameDecoder, FrameEncoder};
29use lzma_rs::{lzma_compress, lzma_decompress};
30use serde_derive::{Serialize, Deserialize};
31
32use std::fmt::Display;
33use std::io::{Cursor, Read, Seek, Write};
34
35use crate::binary::{ReadBytes, WriteBytes};
36use crate::error::{RLibError, Result};
37
38#[cfg(test)]
39mod test;
40
41// LZMA Alone doesn't have a defined magic number, but it always starts with one of these, depending on the compression level.
42const MAGIC_NUMBERS_LZMA: [u32; 9] = [
43    0x0100005D,
44    0x1000005D,
45    0x0800005D,
46    0x1000005D,
47    0x2000005D,
48    0x4000005D,
49    0x8000005D,
50    0x0000005D,
51    0x0400005D,
52];
53const MAGIC_NUMBER_LZ4: u32 = 0x184D2204;
54const MAGIC_NUMBER_ZSTD: u32 = 0xfd2fb528;
55
56//---------------------------------------------------------------------------//
57//                                  Traits
58//---------------------------------------------------------------------------//
59
60/// Internal trait to implement compression over a data type.
61///
62/// See also [`Decompressible`] for the reverse operation.
63pub trait Compressible {
64
65    /// This function compress the data of a file, returning the compressed data.
66    ///
67    /// # Arguments
68    ///
69    /// * `format` - The [`CompressionFormat`] to use for compression.
70    ///
71    /// # Returns
72    ///
73    /// A [`Vec<u8>`] containing the compressed data, or an error if compression failed.
74    fn compress(&self, format: CompressionFormat) -> Result<Vec<u8>>;
75}
76
77/// Internal trait to implement decompression over a data type.
78///
79/// See also [`Compressible`] for the reverse operation.
80pub trait Decompressible {
81
82    /// This function decompress the provided data, returning the decompressed data, or an error if the decompression failed.
83    ///
84    /// Compression format is auto-detected using each format's magic numbers. See [`CompressionFormat`] for details
85    /// on the supported formats.
86    ///
87    /// # Returns
88    ///
89    /// A [`Vec<u8>`] containing the decompressed data, or an error if decompression failed.
90    fn decompress(&self) -> Result<Vec<u8>>;
91
92    /// This function decompress the provided [`CompressionFormat::Lz4Block`] data, returning the decompressed data, or an error if the decompression failed.
93    ///
94    /// This format cannot be auto-detected, as it has neither magic number nor size header, so it needs its own function.
95    ///
96    /// # Arguments
97    ///
98    /// * `uncompressed_size` - The size of the data once decompressed. Decompression fails if the result doesn't match it.
99    ///
100    /// # Returns
101    ///
102    /// A [`Vec<u8>`] containing the decompressed data, or an error if decompression failed.
103    fn decompress_lz4_block(&self, uncompressed_size: usize) -> Result<Vec<u8>>;
104}
105
106/// Compression formats supported by TW Games.
107///
108/// Not all games support all formats. Check their game info to know what formats each game support.
109#[derive(Debug, Copy, Clone, Default, PartialEq, Serialize, Deserialize)]
110pub enum CompressionFormat {
111
112    /// Dummy variant to disable compression.
113    #[default]None,
114
115    /// Legacy format. Supported by all PFH5 games (all Post-WH2 games).
116    ///
117    /// Specifically, Total War games use the Non-Streamed LZMA1 format with the following custom header:
118    ///
119    /// | Bytes | Type     | Data                                                                                |
120    /// | ----- | -------- | ----------------------------------------------------------------------------------- |
121    /// |  4    | [`u32`]  | Uncompressed size (as u32, max at 4GB).                                             |
122    /// |  1    | [`u8`]   | LZMA model properties (lc, lp, pb) in encoded form... I think. Usually it's `0x5D`. |
123    /// |  4    | [`u32`]  | Dictionary size (as u32)... I think. It's usually `[0x00, 0x00, 0x40, 0x00]`.       |
124    ///
125    /// For reference, a normal Non-Streamed LZMA1 header (from the original spec) contains:
126    ///
127    /// | Bytes | Type     | Data                                                        |
128    /// | ----- | -------- | ----------------------------------------------------------- |
129    /// |  1    | [`u8`]   | LZMA model properties (lc, lp, pb) in encoded form.         |
130    /// |  4    | [`u32`]  | Dictionary size (32-bit unsigned integer, little-endian).   |
131    /// |  8    | [`u64`]  | Uncompressed size (64-bit unsigned integer, little-endian). |
132    ///
133    /// This means one has to move the uncompressed size to the correct place in order for a compressed file to be readable,
134    /// and one has to remove the uncompressed size and prepend it to the file in order for the game to read the compressed file.
135    Lzma1,
136
137    /// New format introduced in WH3 6.2.
138    ///
139    /// This is a standard Lz4 implementation, with the following tweaks:
140    ///
141    /// | Bytes | Type      | Data                                          |
142    /// | ----- | --------- | --------------------------------------------- |
143    /// |  4    | [`u32`]   | Uncompressed size (as u32, max at 4GB).       |
144    /// |  *    | &[[`u8`]] | Lz4 data, starting with the Lz4 Magic Number. |
145    Lz4,
146
147    /// New format introduced in WH3 6.2.
148    ///
149    /// This is a standard Zstd implementation, with the following tweaks:
150    ///
151    /// | Bytes | Type      | Data                                            |
152    /// | ----- | --------- | ----------------------------------------------- |
153    /// |  4    | [`u32`]   | Uncompressed size (as u32, max at 4GB).         |
154    /// |  *    | &[[`u8`]] | Zstd data, starting with the Zstd Magic Number. |
155    ///
156    /// By default the Zstd compression is done with the checksum and content size flags enabled.
157    Zstd,
158
159    /// Format used by Shogun 2 Complete Edition PFH3 Packs.
160    ///
161    /// This is a raw Lz4 block, without frame, magic number nor size header. The uncompressed size is stored in the Pack's index instead,
162    /// which means this format cannot be auto-detected. Use [`Decompressible::decompress_lz4_block`] to decompress it.
163    Lz4Block,
164}
165
166//---------------------------------------------------------------------------//
167//                              Implementations
168//---------------------------------------------------------------------------//
169
170impl Compressible for [u8] {
171    fn compress(&self, format: CompressionFormat) -> Result<Vec<u8>> {
172        match format {
173            CompressionFormat::None => Ok(self.to_vec()),
174            CompressionFormat::Lzma1 => {
175                let mut dst = vec![];
176                dst.write_i32(self.len() as i32)?;
177
178                let mut compressed_data = vec![];
179                let mut src = Cursor::new(self);
180                lzma_compress(&mut src, &mut compressed_data).unwrap();
181
182                if compressed_data.len() < 13 {
183                    return Err(RLibError::DataCannotBeCompressed);
184                }
185
186                dst.extend_from_slice(&compressed_data[..5]);
187                dst.extend_from_slice(&compressed_data[13..]);
188
189                Ok(dst)
190            },
191            CompressionFormat::Lz4 => {
192                let mut dst = vec![];
193                dst.write_u32(self.len() as u32)?;
194
195                let mut encoder = FrameEncoder::new(&mut dst);
196                encoder.write_all(self)?;
197                encoder.finish()?;
198
199                Ok(dst)
200            },
201            CompressionFormat::Zstd => {
202                let mut dst = vec![];
203                dst.write_u32(self.len() as u32)?;
204
205                let mut encoder = zstd::Encoder::new(&mut dst, 3)?;
206                encoder.include_checksum(true)?;
207                encoder.include_contentsize(true)?;
208                encoder.set_pledged_src_size(Some(self.len() as u64))?;
209
210                let mut src = Cursor::new(self.to_vec());
211                std::io::copy(&mut src, &mut encoder)?;
212                encoder.finish()?;
213                Ok(dst)
214            },
215            CompressionFormat::Lz4Block => Ok(lz4_flex::block::compress(self)),
216        }
217    }
218}
219
220impl Decompressible for &[u8] {
221    fn decompress(&self) -> Result<Vec<u8>> {
222        if self.is_empty() {
223            return Ok(vec![]);
224        }
225
226        // We use the magic numbers to know in what format are the files compressed.
227        let mut src = Cursor::new(self);
228        let u_size = src.read_u32()?;
229        let magic_number = src.read_u32()?;
230
231        let format = if magic_number == MAGIC_NUMBER_ZSTD {
232            CompressionFormat::Zstd
233        } else if magic_number == MAGIC_NUMBER_LZ4 {
234            CompressionFormat::Lz4
235        } else if MAGIC_NUMBERS_LZMA.contains(&magic_number) {
236            CompressionFormat::Lzma1
237        }
238
239        // Special case files marked as compressed but not being compressed. This allows fixing them so they're readable again.
240        else {
241            CompressionFormat::None
242        };
243
244        // Fix the starting position of the file before processing it.
245        src.seek_relative(-4)?;
246
247        match format {
248            CompressionFormat::None => Ok(self.to_vec()),
249            CompressionFormat::Lzma1 => {
250
251                // LZMA1 headers have 13 bytes, but we only have 9 due to using a u32 size.
252                if self.len() < 9 {
253                    return Err(RLibError::DataCannotBeDecompressed);
254                }
255
256                // Unlike other formats, in this one we need to inject the uncompressed size in the file header. Otherwise it won't be a valid lzma file.
257                let mut fixed_data: Vec<u8> = Vec::with_capacity(self.len() + 4);
258                fixed_data.extend_from_slice(&src.read_slice(5, false)?);
259                fixed_data.write_u64(u_size as u64)?;
260                src.read_to_end(&mut fixed_data)?;
261
262                // Vanilla compressed files are LZMA Alone (or legacy) level 3 compressed files, reproducible by compressing them
263                // with default settings with 7-Zip. This should do the trick to get them decoded.
264                let mut dst = Vec::with_capacity(u_size as usize);
265                let mut reader = Cursor::new(fixed_data);
266                let result = lzma_decompress(&mut reader, &mut dst);
267
268                // Ok, history lesson. That method breaks sometimes due to difference in program's behavior when reading LZMA1 files with uncompressed size set.
269                // If that fails, we try passing a unknown size (u64::MAX) instead. This usually deals with the errors.
270                if result.is_err() {
271                    src.set_position(4);
272
273                    let mut fixed_data = Vec::with_capacity(self.len() + 4);
274                    fixed_data.extend_from_slice(&src.read_slice(5, false)?);
275                    fixed_data.write_u64(u64::MAX)?;
276                    src.read_to_end(&mut fixed_data)?;
277
278                    let mut dst = Vec::with_capacity(u_size as usize);
279                    let mut reader = Cursor::new(fixed_data);
280                    lzma_decompress(&mut reader, &mut dst)?;
281
282                    Ok(dst)
283                } else {
284                    Ok(dst)
285                }
286            },
287            CompressionFormat::Lz4 => {
288                let mut dst = Vec::with_capacity(u_size as usize);
289                let mut reader = FrameDecoder::new(src);
290                std::io::copy(&mut reader, &mut dst)?;
291                Ok(dst)
292            },
293            CompressionFormat::Zstd => {
294                let mut dst = Vec::with_capacity(u_size as usize);
295                zstd::stream::copy_decode(src, &mut dst)?;
296                Ok(dst)
297            },
298
299            // Not detectable by magic number, so this is never reached.
300            CompressionFormat::Lz4Block => Ok(self.to_vec()),
301        }
302    }
303
304    fn decompress_lz4_block(&self, uncompressed_size: usize) -> Result<Vec<u8>> {
305        if uncompressed_size == 0 {
306            return Ok(vec![]);
307        }
308
309        let dst = lz4_flex::block::decompress(self, uncompressed_size)?;
310        if dst.len() != uncompressed_size {
311            return Err(RLibError::DataCannotBeDecompressed);
312        }
313
314        Ok(dst)
315    }
316}
317
318impl From<&str> for CompressionFormat {
319    fn from(value: &str) -> Self {
320        match value {
321            "Lzma1" => Self::Lzma1,
322            "Lz4" => Self::Lz4,
323            "Zstd" => Self::Zstd,
324            "Lz4Block" => Self::Lz4Block,
325            _ => Self::None,
326        }
327    }
328}
329
330impl Display for CompressionFormat {
331    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
332        match self {
333            Self::Lzma1 => write!(f, "Lzma1"),
334            Self::Lz4 => write!(f, "Lz4"),
335            Self::Zstd => write!(f, "Zstd"),
336            Self::Lz4Block => write!(f, "Lz4Block"),
337            Self::None => write!(f, "None"),
338        }
339    }
340}