imperator_save/
file.rs

1use crate::{
2    flavor::ImperatorFlavor, Encoding, ImperatorError, ImperatorErrorKind, ImperatorMelter,
3    SaveHeader,
4};
5use jomini::{
6    binary::{BinaryDeserializerBuilder, FailedResolveStrategy, TokenResolver},
7    text::ObjectReader,
8    BinaryDeserializer, BinaryTape, TextDeserializer, TextTape, Utf8Encoding,
9};
10use serde::Deserialize;
11use std::io::Cursor;
12use zip::result::ZipError;
13
14enum FileKind<'a> {
15    Text(&'a [u8]),
16    Binary(&'a [u8]),
17    Zip {
18        archive: ImperatorZipFiles<'a>,
19        metadata: &'a [u8],
20        gamestate: VerifiedIndex,
21        is_text: bool,
22    },
23}
24
25/// Entrypoint for parsing Imperator saves
26///
27/// Only consumes enough data to determine encoding of the file
28pub struct ImperatorFile<'a> {
29    header: SaveHeader,
30    kind: FileKind<'a>,
31}
32
33impl<'a> ImperatorFile<'a> {
34    /// Creates a Imperator file from a slice of data
35    pub fn from_slice(data: &[u8]) -> Result<ImperatorFile, ImperatorError> {
36        let header = SaveHeader::from_slice(data)?;
37        let data = &data[header.header_len()..];
38
39        let reader = Cursor::new(data);
40        match zip::ZipArchive::new(reader) {
41            Ok(mut zip) => {
42                let metadata = &data[..zip.offset() as usize];
43                let files = ImperatorZipFiles::new(&mut zip, data);
44                let gamestate_idx = files
45                    .gamestate_index()
46                    .ok_or(ImperatorErrorKind::ZipMissingEntry)?;
47
48                let is_text = !header.kind().is_binary();
49                Ok(ImperatorFile {
50                    header,
51                    kind: FileKind::Zip {
52                        archive: files,
53                        gamestate: gamestate_idx,
54                        metadata,
55                        is_text,
56                    },
57                })
58            }
59            Err(ZipError::InvalidArchive(_)) => {
60                if header.kind().is_binary() {
61                    Ok(ImperatorFile {
62                        header,
63                        kind: FileKind::Binary(data),
64                    })
65                } else {
66                    Ok(ImperatorFile {
67                        header,
68                        kind: FileKind::Text(data),
69                    })
70                }
71            }
72            Err(e) => Err(ImperatorErrorKind::ZipArchive(e).into()),
73        }
74    }
75
76    /// Returns the detected decoding of the file
77    pub fn encoding(&self) -> Encoding {
78        match &self.kind {
79            FileKind::Text(_) => Encoding::Text,
80            FileKind::Binary(_) => Encoding::Binary,
81            FileKind::Zip { is_text, .. } if *is_text => Encoding::TextZip,
82            FileKind::Zip { .. } => Encoding::BinaryZip,
83        }
84    }
85
86    /// Returns the size of the file
87    ///
88    /// The size includes the inflated size of the zip
89    pub fn size(&self) -> usize {
90        match &self.kind {
91            FileKind::Text(x) | FileKind::Binary(x) => x.len(),
92            FileKind::Zip { gamestate, .. } => gamestate.size,
93        }
94    }
95
96    pub fn parse_metadata(&self) -> Result<ImperatorParsedFile<'a>, ImperatorError> {
97        match &self.kind {
98            FileKind::Text(x) => {
99                // The metadata section should be way smaller than the total
100                // length so if the total data isn't significantly bigger (2x or
101                // more), assume that the header doesn't accurately represent
102                // the metadata length. Like maybe someone accidentally
103                // converted the line endings from unix to dos.
104                let len = self.header.metadata_len() as usize;
105                let data = if len * 2 > x.len() {
106                    x
107                } else {
108                    &x[..len.min(x.len())]
109                };
110
111                let text = ImperatorText::from_raw(data)?;
112                Ok(ImperatorParsedFile {
113                    kind: ImperatorParsedFileKind::Text(text),
114                })
115            }
116            FileKind::Binary(x) => {
117                let metadata = x.get(..self.header.metadata_len() as usize).unwrap_or(x);
118                let binary = ImperatorBinary::from_raw(metadata, self.header.clone())?;
119                Ok(ImperatorParsedFile {
120                    kind: ImperatorParsedFileKind::Binary(binary),
121                })
122            }
123            FileKind::Zip {
124                metadata, is_text, ..
125            } if *is_text => {
126                let text = ImperatorText::from_raw(metadata)?;
127                Ok(ImperatorParsedFile {
128                    kind: ImperatorParsedFileKind::Text(text),
129                })
130            }
131            FileKind::Zip { metadata, .. } => {
132                let binary = ImperatorBinary::from_raw(metadata, self.header.clone())?;
133                Ok(ImperatorParsedFile {
134                    kind: ImperatorParsedFileKind::Binary(binary),
135                })
136            }
137        }
138    }
139
140    /// Parses the entire file
141    ///
142    /// If the file is a zip, the zip contents will be inflated into the zip
143    /// sink before being parsed
144    pub fn parse(
145        &self,
146        zip_sink: &'a mut Vec<u8>,
147    ) -> Result<ImperatorParsedFile<'a>, ImperatorError> {
148        match &self.kind {
149            FileKind::Text(x) => {
150                let text = ImperatorText::from_raw(x)?;
151                Ok(ImperatorParsedFile {
152                    kind: ImperatorParsedFileKind::Text(text),
153                })
154            }
155            FileKind::Binary(x) => {
156                let binary = ImperatorBinary::from_raw(x, self.header.clone())?;
157                Ok(ImperatorParsedFile {
158                    kind: ImperatorParsedFileKind::Binary(binary),
159                })
160            }
161            FileKind::Zip {
162                archive,
163                gamestate,
164                is_text,
165                ..
166            } => {
167                let zip = archive.retrieve_file(*gamestate);
168                zip.read_to_end(zip_sink)?;
169
170                if *is_text {
171                    let text = ImperatorText::from_raw(zip_sink)?;
172                    Ok(ImperatorParsedFile {
173                        kind: ImperatorParsedFileKind::Text(text),
174                    })
175                } else {
176                    let binary = ImperatorBinary::from_raw(zip_sink, self.header.clone())?;
177                    Ok(ImperatorParsedFile {
178                        kind: ImperatorParsedFileKind::Binary(binary),
179                    })
180                }
181            }
182        }
183    }
184}
185
186/// Contains the parsed Imperator file
187pub enum ImperatorParsedFileKind<'a> {
188    /// The Imperator file as text
189    Text(ImperatorText<'a>),
190
191    /// The Imperator file as binary
192    Binary(ImperatorBinary<'a>),
193}
194
195/// An Imperator file that has been parsed
196pub struct ImperatorParsedFile<'a> {
197    kind: ImperatorParsedFileKind<'a>,
198}
199
200impl<'a> ImperatorParsedFile<'a> {
201    /// Returns the file as text
202    pub fn as_text(&self) -> Option<&ImperatorText> {
203        match &self.kind {
204            ImperatorParsedFileKind::Text(x) => Some(x),
205            _ => None,
206        }
207    }
208
209    /// Returns the file as binary
210    pub fn as_binary(&self) -> Option<&ImperatorBinary> {
211        match &self.kind {
212            ImperatorParsedFileKind::Binary(x) => Some(x),
213            _ => None,
214        }
215    }
216
217    /// Returns the kind of file (binary or text)
218    pub fn kind(&self) -> &ImperatorParsedFileKind {
219        &self.kind
220    }
221
222    /// Prepares the file for deserialization into a custom structure
223    pub fn deserializer(&self) -> ImperatorDeserializer {
224        match &self.kind {
225            ImperatorParsedFileKind::Text(x) => ImperatorDeserializer {
226                kind: ImperatorDeserializerKind::Text(x),
227            },
228            ImperatorParsedFileKind::Binary(x) => ImperatorDeserializer {
229                kind: ImperatorDeserializerKind::Binary(x.deserializer()),
230            },
231        }
232    }
233}
234
235#[derive(Debug, Clone, Copy)]
236struct VerifiedIndex {
237    data_start: usize,
238    data_end: usize,
239    size: usize,
240}
241
242#[derive(Debug, Clone)]
243struct ImperatorZipFiles<'a> {
244    archive: &'a [u8],
245    gamestate_index: Option<VerifiedIndex>,
246}
247
248impl<'a> ImperatorZipFiles<'a> {
249    pub fn new(archive: &mut zip::ZipArchive<Cursor<&'a [u8]>>, data: &'a [u8]) -> Self {
250        let mut gamestate_index = None;
251
252        for index in 0..archive.len() {
253            if let Ok(file) = archive.by_index_raw(index) {
254                let size = file.size() as usize;
255                let data_start = file.data_start() as usize;
256                let data_end = data_start + file.compressed_size() as usize;
257
258                if file.name() == "gamestate" {
259                    gamestate_index = Some(VerifiedIndex {
260                        data_start,
261                        data_end,
262                        size,
263                    })
264                }
265            }
266        }
267
268        Self {
269            archive: data,
270            gamestate_index,
271        }
272    }
273
274    pub fn retrieve_file(&self, index: VerifiedIndex) -> ImperatorZipFile {
275        let raw = &self.archive[index.data_start..index.data_end];
276        ImperatorZipFile {
277            raw,
278            size: index.size,
279        }
280    }
281
282    pub fn gamestate_index(&self) -> Option<VerifiedIndex> {
283        self.gamestate_index
284    }
285}
286
287struct ImperatorZipFile<'a> {
288    raw: &'a [u8],
289    size: usize,
290}
291
292impl<'a> ImperatorZipFile<'a> {
293    pub fn read_to_end(&self, buf: &mut Vec<u8>) -> Result<(), ImperatorError> {
294        let start_len = buf.len();
295        buf.resize(start_len + self.size(), 0);
296        let body = &mut buf[start_len..];
297        crate::deflate::inflate_exact(self.raw, body).map_err(ImperatorErrorKind::from)?;
298        Ok(())
299    }
300
301    pub fn size(&self) -> usize {
302        self.size
303    }
304}
305
306/// A parsed Imperator text document
307pub struct ImperatorText<'a> {
308    tape: TextTape<'a>,
309}
310
311impl<'a> ImperatorText<'a> {
312    pub fn from_slice(data: &'a [u8]) -> Result<Self, ImperatorError> {
313        let header = SaveHeader::from_slice(data)?;
314        Self::from_raw(&data[..header.header_len()])
315    }
316
317    pub(crate) fn from_raw(data: &'a [u8]) -> Result<Self, ImperatorError> {
318        let tape = TextTape::from_slice(data).map_err(ImperatorErrorKind::Parse)?;
319        Ok(ImperatorText { tape })
320    }
321
322    pub fn reader(&self) -> ObjectReader<Utf8Encoding> {
323        self.tape.utf8_reader()
324    }
325
326    pub fn deserialize<T>(&self) -> Result<T, ImperatorError>
327    where
328        T: Deserialize<'a>,
329    {
330        let result = TextDeserializer::from_utf8_tape(&self.tape)
331            .map_err(ImperatorErrorKind::Deserialize)?;
332        Ok(result)
333    }
334}
335
336/// A parsed Imperator binary document
337pub struct ImperatorBinary<'a> {
338    tape: BinaryTape<'a>,
339    header: SaveHeader,
340}
341
342impl<'a> ImperatorBinary<'a> {
343    pub fn from_slice(data: &'a [u8]) -> Result<Self, ImperatorError> {
344        let header = SaveHeader::from_slice(data)?;
345        Self::from_raw(&data[..header.header_len()], header)
346    }
347
348    pub(crate) fn from_raw(data: &'a [u8], header: SaveHeader) -> Result<Self, ImperatorError> {
349        let tape = BinaryTape::from_slice(data).map_err(ImperatorErrorKind::Parse)?;
350        Ok(ImperatorBinary { tape, header })
351    }
352
353    pub fn deserializer<'b>(&'b self) -> ImperatorBinaryDeserializer<'a, 'b> {
354        ImperatorBinaryDeserializer {
355            builder: BinaryDeserializer::builder_flavor(ImperatorFlavor),
356            tape: &self.tape,
357        }
358    }
359
360    pub fn melter<'b>(&'b self) -> ImperatorMelter<'a, 'b> {
361        ImperatorMelter::new(&self.tape, &self.header)
362    }
363}
364
365enum ImperatorDeserializerKind<'a, 'b> {
366    Text(&'b ImperatorText<'a>),
367    Binary(ImperatorBinaryDeserializer<'a, 'b>),
368}
369
370/// A deserializer for custom structures
371pub struct ImperatorDeserializer<'a, 'b> {
372    kind: ImperatorDeserializerKind<'a, 'b>,
373}
374
375impl<'a, 'b> ImperatorDeserializer<'a, 'b> {
376    pub fn on_failed_resolve(&mut self, strategy: FailedResolveStrategy) -> &mut Self {
377        if let ImperatorDeserializerKind::Binary(x) = &mut self.kind {
378            x.on_failed_resolve(strategy);
379        }
380        self
381    }
382
383    pub fn build<T, R>(&self, resolver: &'a R) -> Result<T, ImperatorError>
384    where
385        R: TokenResolver,
386        T: Deserialize<'a>,
387    {
388        match &self.kind {
389            ImperatorDeserializerKind::Text(x) => x.deserialize(),
390            ImperatorDeserializerKind::Binary(x) => x.build(resolver),
391        }
392    }
393}
394
395/// Deserializes binary data into custom structures
396pub struct ImperatorBinaryDeserializer<'a, 'b> {
397    builder: BinaryDeserializerBuilder<ImperatorFlavor>,
398    tape: &'b BinaryTape<'a>,
399}
400
401impl<'a, 'b> ImperatorBinaryDeserializer<'a, 'b> {
402    pub fn on_failed_resolve(&mut self, strategy: FailedResolveStrategy) -> &mut Self {
403        self.builder.on_failed_resolve(strategy);
404        self
405    }
406
407    pub fn build<T, R>(&self, resolver: &'a R) -> Result<T, ImperatorError>
408    where
409        R: TokenResolver,
410        T: Deserialize<'a>,
411    {
412        let result = self
413            .builder
414            .from_tape(self.tape, resolver)
415            .map_err(|e| match e.kind() {
416                jomini::ErrorKind::Deserialize(e2) => match e2.kind() {
417                    &jomini::DeserializeErrorKind::UnknownToken { token_id } => {
418                        ImperatorErrorKind::UnknownToken { token_id }
419                    }
420                    _ => ImperatorErrorKind::Deserialize(e),
421                },
422                _ => ImperatorErrorKind::Deserialize(e),
423            })?;
424        Ok(result)
425    }
426}