mpq_rs/core/
archive.rs

1use std::io::{Read, Seek};
2
3use crate::core::error::*;
4use crate::core::seeker::*;
5use crate::core::table::*;
6use crate::core::util::*;
7
8#[derive(Debug)]
9/// Implementation of a MoPaQ archive viewer.
10///
11/// Refer to top-level documentation to see which features are supported.
12///
13/// Will work on any reader that implements `Read + Seek`.
14pub struct Archive<R: Read + Seek> {
15    seeker: Seeker<R>,
16    hash_table: FileHashTable,
17    block_table: FileBlockTable,
18}
19
20impl<R: Read + Seek> Archive<R> {
21    /// Try to open an MPQ archive from the specified `reader`.
22    ///
23    /// Immediately, this will perform the following:
24    ///
25    /// 1. Locate an MPQ header.
26    /// 2. Locate and read the Hash Table.
27    /// 3. Locate and read the Block Table.
28    ///
29    /// If any of these steps fail, the archive is deemed corrupted and
30    /// an appropriate error is returned.
31    ///
32    /// No other operations will be performed.
33    pub fn open(reader: R) -> Result<Archive<R>, MpqError> {
34        let mut seeker = Seeker::new(reader)?;
35
36        let hash_table = FileHashTable::from_seeker(&mut seeker)?;
37        let block_table = FileBlockTable::from_seeker(&mut seeker)?;
38
39        Ok(Archive {
40            seeker,
41            hash_table,
42            block_table,
43        })
44    }
45
46    /// Read a file's contents.
47    ///
48    /// Notably, the filename resolution algorithm
49    /// is case, and will treat backslashes (`\`) and forward slashes (`/`)
50    /// as different characters.
51    ///
52    /// Does not support single-unit files or uncompressed files.
53    pub fn read_file(&mut self, name: &str) -> Result<Vec<u8>, MpqError> {
54        // find the hash entry and use it to find the block entry
55        let hash_entry = self
56            .hash_table
57            .find_entry(name)
58            .ok_or(MpqError::FileNotFound)?;
59        let block_entry = self
60            .block_table
61            .get(hash_entry.block_index as usize)
62            .ok_or(MpqError::FileNotFound)?;
63
64        // calculate the file key
65        let encryption_key = if block_entry.is_encrypted() {
66            Some(calculate_file_key(
67                name,
68                block_entry.file_pos as u32,
69                block_entry.uncompressed_size as u32,
70                block_entry.is_key_adjusted(),
71            ))
72        } else {
73            None
74        };
75
76        // read the sector offsets
77        let sector_offsets = SectorOffsets::from_reader(
78            &mut self.seeker,
79            block_entry,
80            encryption_key.map(|k| k - 1),
81        )?;
82
83        // read out all the sectors
84        let sector_range = sector_offsets.all();
85        let raw_data = self.seeker.read(
86            block_entry.file_pos + u64::from(sector_range.0),
87            u64::from(sector_range.1),
88        )?;
89
90        let mut result = Vec::with_capacity(block_entry.uncompressed_size as usize);
91
92        let sector_size = self.seeker.info().sector_size;
93        let sector_count = sector_offsets.count();
94        let first_sector_offset = sector_offsets.one(0).unwrap().0;
95        for i in 0..sector_count {
96            let sector_offset = sector_offsets.one(i).unwrap();
97            let slice_start = (sector_offset.0 - first_sector_offset) as usize;
98            let slice_end = slice_start + sector_offset.1 as usize;
99
100            // if this is the last sector, then its size will be less than
101            // one archive sector size, so account for that
102            let uncompressed_size = if (i + 1) == sector_count {
103                let size = block_entry.uncompressed_size % sector_size;
104
105                if size == 0 {
106                    sector_size
107                } else {
108                    size
109                }
110            } else {
111                sector_size
112            };
113
114            // decode the block and append it to the final result buffer
115            let decoded_sector = decode_mpq_block(
116                &raw_data[slice_start..slice_end],
117                uncompressed_size,
118                encryption_key.map(|k| k + i as u32),
119            )?;
120
121            result.extend(decoded_sector.iter());
122        }
123
124        Ok(result)
125    }
126
127    /// If the archive contains a `(listfile)`, this will method
128    /// parse it and return a `Vec` containing all known filenames.
129    pub fn files(&mut self) -> Option<Vec<String>> {
130        let listfile = self.read_file("(listfile)").ok()?;
131
132        let mut list = Vec::new();
133        let mut line_start = 0;
134        for i in 0..listfile.len() {
135            let byte = listfile[i];
136
137            if byte == b'\r' || byte == b'\n' {
138                if i - line_start > 0 {
139                    let line = &listfile[line_start..i];
140                    let line = std::str::from_utf8(line);
141
142                    if let Ok(line) = line {
143                        list.push(line.to_string());
144                    }
145                }
146
147                line_start = i + 1;
148            }
149        }
150
151        Some(list)
152    }
153
154    // Returns the start of the archive in the reader, which is the MPQ header,
155    // relative to the beginning of the reader.
156    pub fn start(&self) -> u64 {
157        self.seeker.info().header_offset
158    }
159
160    // Returns the end of the archive in the reader, relative to the beginning of the reader.
161    pub fn end(&self) -> u64 {
162        self.seeker.info().header_offset + self.seeker.info().archive_size
163    }
164
165    // Returns the size of the archive as specified in the MPQ header.
166    pub fn size(&self) -> u64 {
167        self.seeker.info().archive_size
168    }
169
170    // Returns a mutable reference to the underlying reader.
171    pub fn reader(&mut self) -> &mut R {
172        self.seeker.reader()
173    }
174}