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}