pk2/raw/
block_manager.rs

1use std::collections::{HashMap, HashSet};
2use std::io;
3use std::path::{Component, Path};
4
5use crate::blowfish::Blowfish;
6use crate::constants::{PK2_FILE_BLOCK_ENTRY_COUNT, PK2_ROOT_BLOCK, PK2_ROOT_BLOCK_VIRTUAL};
7use crate::error::{ChainLookupError, ChainLookupResult, OpenResult};
8use crate::raw::block_chain::{PackBlock, PackBlockChain};
9use crate::raw::entry::{DirectoryEntry, PackEntry};
10use crate::raw::{BlockOffset, ChainIndex};
11
12/// Simple BlockManager backed by a hashmap.
13pub struct BlockManager {
14    chains: HashMap<ChainIndex, PackBlockChain, NoHashHasherBuilder>,
15}
16
17impl BlockManager {
18    /// Parses the complete index of a pk2 file
19    pub fn new<F: io::Read + io::Seek>(bf: Option<&Blowfish>, mut stream: F) -> OpenResult<Self> {
20        let mut chains = HashMap::with_capacity_and_hasher(32, NoHashHasherBuilder);
21        // used to prevent an infinite loop that can be caused by specific files
22        let mut visited_block_set = HashSet::with_capacity_and_hasher(32, NoHashHasherBuilder);
23        let mut offsets = vec![PK2_ROOT_BLOCK];
24        while let Some(offset) = offsets.pop() {
25            if chains.contains_key(&offset) {
26                // skip offsets that are being pointed to multiple times
27                continue;
28            }
29            let block_chain =
30                Self::read_chain_from_stream_at(&mut visited_block_set, bf, &mut stream, offset)?;
31            visited_block_set.clear();
32
33            // put all folder offsets of this chain into the stack to parse them next
34            offsets.extend(
35                block_chain
36                    .entries()
37                    .filter_map(PackEntry::as_directory)
38                    .filter(|d| d.is_normal_link())
39                    .map(DirectoryEntry::children_position),
40            );
41            chains.insert(offset, block_chain);
42        }
43        let mut this = BlockManager { chains };
44        this.insert_virtual_root();
45        Ok(this)
46    }
47
48    fn insert_virtual_root(&mut self) {
49        // dummy entry to give root a proper name
50        let mut virtual_root = PackBlockChain::from_blocks(vec![(
51            PK2_ROOT_BLOCK_VIRTUAL.into(),
52            PackBlock::default(),
53        )]);
54        virtual_root[0] = PackEntry::new_directory("/", PK2_ROOT_BLOCK, None);
55        self.chains.insert(virtual_root.chain_index(), virtual_root);
56    }
57
58    /// Reads a [`PackBlockChain`] from the given file at the specified offset.
59    fn read_chain_from_stream_at<F: io::Read + io::Seek>(
60        visited_block_set: &mut HashSet<BlockOffset, NoHashHasherBuilder>,
61        bf: Option<&Blowfish>,
62        stream: &mut F,
63        offset: ChainIndex,
64    ) -> OpenResult<PackBlockChain> {
65        let mut blocks = Vec::new();
66        let mut offset = offset.into();
67
68        while visited_block_set.insert(offset) {
69            let block = crate::io::read_block_at(bf, &mut *stream, offset)?;
70            let nc = block.entries().last().and_then(PackEntry::next_block);
71            blocks.push((offset, block));
72            match nc {
73                Some(nc) => offset = BlockOffset(nc.get()),
74                None => break,
75            }
76        }
77        Ok(PackBlockChain::from_blocks(blocks))
78    }
79
80    pub fn get(&self, chain: ChainIndex) -> Option<&PackBlockChain> {
81        self.chains.get(&chain)
82    }
83
84    pub fn get_mut(&mut self, chain: ChainIndex) -> Option<&mut PackBlockChain> {
85        assert_ne!(chain, PK2_ROOT_BLOCK_VIRTUAL);
86        self.chains.get_mut(&chain)
87    }
88
89    pub fn insert(&mut self, chain: ChainIndex, block: PackBlockChain) {
90        self.chains.insert(chain, block);
91    }
92
93    pub fn resolve_path_to_parent<'path>(
94        &self,
95        current_chain: ChainIndex,
96        path: &'path Path,
97    ) -> ChainLookupResult<(ChainIndex, &'path str)> {
98        let mut components = path.components();
99
100        if let Some(c) = components.next_back() {
101            let parent_index =
102                self.resolve_path_to_block_chain_index_at(current_chain, components.as_path())?;
103            let name = c.as_os_str().to_str().ok_or(ChainLookupError::InvalidPath)?;
104            Ok((parent_index, name))
105        } else {
106            Err(ChainLookupError::InvalidPath)
107        }
108    }
109
110    /// Resolves a path from the specified chain to a parent chain and the entry
111    /// Returns Ok(None) if the path is empty, otherwise (blockchain,
112    /// entry_index, entry)
113    pub fn resolve_path_to_entry_and_parent(
114        &self,
115        current_chain: ChainIndex,
116        path: &Path,
117    ) -> ChainLookupResult<(ChainIndex, usize, &PackEntry)> {
118        self.resolve_path_to_parent(current_chain, path).and_then(|(parent_index, name)| {
119            self.chains
120                .get(&parent_index)
121                .ok_or(ChainLookupError::InvalidChainIndex)?
122                .entries()
123                .enumerate()
124                .find(|(_, entry)| entry.name_eq_ignore_ascii_case(name))
125                .ok_or(ChainLookupError::NotFound)
126                .map(|(idx, entry)| (parent_index, idx, entry))
127        })
128    }
129
130    pub fn resolve_path_to_entry_and_parent_mut(
131        &mut self,
132        current_chain: ChainIndex,
133        path: &Path,
134    ) -> ChainLookupResult<(ChainIndex, usize, &mut PackEntry)> {
135        self.resolve_path_to_parent(current_chain, path).and_then(move |(parent_index, name)| {
136            self.chains
137                .get_mut(&parent_index)
138                .ok_or(ChainLookupError::InvalidChainIndex)?
139                .entries_mut()
140                .enumerate()
141                .find(|(_, entry)| entry.name_eq_ignore_ascii_case(name))
142                .ok_or(ChainLookupError::NotFound)
143                .map(|(idx, entry)| (parent_index, idx, entry))
144        })
145    }
146
147    /// Resolves a path to a [`PackBlockChain`] index starting from the given
148    /// blockchain returning the index of the last blockchain.
149    pub fn resolve_path_to_block_chain_index_at(
150        &self,
151        current_chain: ChainIndex,
152        path: &Path,
153    ) -> ChainLookupResult<ChainIndex> {
154        path.components().try_fold(current_chain, |idx, component| {
155            let comp = component.as_os_str().to_str().ok_or(ChainLookupError::InvalidPath)?;
156            self.chains
157                .get(&idx)
158                .ok_or(ChainLookupError::InvalidChainIndex)?
159                .find_block_chain_index_of(comp)
160        })
161    }
162
163    /// Traverses the path until it hits a non-existent component and returns
164    /// the rest of the path as a peekable as well as the chain index of the
165    /// last valid part.
166    /// A return value of Ok(None) means the entire path has been searched
167    pub fn validate_dir_path_until<'p>(
168        &self,
169        mut chain: ChainIndex,
170        path: &'p Path,
171    ) -> ChainLookupResult<Option<(ChainIndex, std::iter::Peekable<std::path::Components<'p>>)>>
172    {
173        let mut components = path.components().peekable();
174        while let Some(component) = components.peek() {
175            let name = component.as_os_str().to_str().ok_or(ChainLookupError::InvalidPath)?;
176            match self
177                .chains
178                .get(&chain)
179                .ok_or(ChainLookupError::InvalidChainIndex)?
180                .find_block_chain_index_of(name)
181            {
182                Ok(i) => chain = i,
183                // lies outside of the archive
184                Err(ChainLookupError::NotFound) if component == &Component::ParentDir => {
185                    return Err(ChainLookupError::InvalidPath)
186                }
187                // found a non-existent part, we are done here
188                Err(ChainLookupError::NotFound) => break,
189                Err(ChainLookupError::ExpectedDirectory) => {
190                    return if components.count() == 1 {
191                        // found a file name at the end of the path
192                        // this means the path has been fully searched
193                        Ok(None)
194                    } else {
195                        Err(ChainLookupError::ExpectedDirectory)
196                    };
197                }
198                Err(_) => unreachable!(),
199            }
200            let _ = components.next();
201        }
202        if components.clone().count() == 0 {
203            Ok(None)
204        } else {
205            Ok(Some((chain, components)))
206        }
207    }
208
209    pub fn sort(&mut self) {
210        let scratch = &mut Vec::with_capacity(4 * PK2_FILE_BLOCK_ENTRY_COUNT);
211        for chain in self.chains.values_mut() {
212            chain.sort(scratch);
213            scratch.clear();
214        }
215    }
216}
217
218#[derive(Default)]
219struct NoHashHasherBuilder;
220impl std::hash::BuildHasher for NoHashHasherBuilder {
221    type Hasher = NoHashHasher;
222    #[inline(always)]
223    fn build_hasher(&self) -> Self::Hasher {
224        NoHashHasher(0)
225    }
226}
227
228struct NoHashHasher(u64);
229impl std::hash::Hasher for NoHashHasher {
230    #[inline(always)]
231    fn finish(&self) -> u64 {
232        self.0
233    }
234
235    fn write(&mut self, _: &[u8]) {
236        panic!("ChainIndex has been hashed wrong. This is a bug!");
237    }
238
239    #[inline(always)]
240    fn write_u64(&mut self, chain: u64) {
241        self.0 = chain;
242    }
243}
pk2/raw/block_manager.rs

pk2/raw/
block_manager.rs
