1use core::iter::FusedIterator;
5use core::mem;
6use core::ops::{Deref, Range};
7
8use zerocopy::byteorder::LittleEndian;
9use zerocopy::{
10 FromBytes, Immutable, IntoBytes, KnownLayout, Ref, SplitByteSlice, SplitByteSliceMut,
11 Unaligned, U32,
12};
13
14use crate::error::{NtHiveError, Result};
15use crate::helpers::byte_subrange;
16use crate::hive::Hive;
17use crate::key_node::{KeyNode, KeyNodeMut};
18use crate::leaf::LeafItemRanges;
19
20#[derive(FromBytes, Immutable, IntoBytes, KnownLayout, Unaligned)]
22#[repr(packed)]
23struct IndexRootItem {
24 subkeys_list_offset: U32<LittleEndian>,
25}
26
27pub(crate) struct IndexRootItemRange(Range<usize>);
29
30impl IndexRootItemRange {
31 pub fn subkeys_list_offset<B>(&self, hive: &Hive<B>) -> u32
32 where
33 B: SplitByteSlice,
34 {
35 let item = Ref::<&[u8], IndexRootItem>::from_bytes(&hive.data[self.0.clone()]).unwrap();
36 item.subkeys_list_offset.get()
37 }
38}
39
40impl Deref for IndexRootItemRange {
41 type Target = Range<usize>;
42
43 fn deref(&self) -> &Self::Target {
44 &self.0
45 }
46}
47
48#[derive(Clone)]
54pub(crate) struct IndexRootItemRanges {
55 items_range: Range<usize>,
56}
57
58impl IndexRootItemRanges {
59 fn new(count: u16, count_field_offset: usize, data_range: Range<usize>) -> Result<Self> {
60 let byte_count = count as usize * mem::size_of::<IndexRootItem>();
61
62 let items_range = byte_subrange(&data_range, byte_count).ok_or_else(|| {
63 NtHiveError::InvalidSizeField {
64 offset: count_field_offset,
65 expected: byte_count,
66 actual: data_range.len(),
67 }
68 })?;
69
70 Ok(Self { items_range })
71 }
72}
73
74impl Iterator for IndexRootItemRanges {
75 type Item = IndexRootItemRange;
76
77 fn next(&mut self) -> Option<Self::Item> {
78 let item_range = byte_subrange(&self.items_range, mem::size_of::<IndexRootItem>())?;
79 self.items_range.start += mem::size_of::<IndexRootItem>();
80
81 Some(IndexRootItemRange(item_range))
82 }
83
84 fn count(self) -> usize {
85 let (size, _) = self.size_hint();
86 size
87 }
88
89 fn last(mut self) -> Option<Self::Item> {
90 let (size, _) = self.size_hint();
91 if size == 0 {
92 return None;
93 }
94
95 self.nth(size - 1)
96 }
97
98 fn nth(&mut self, n: usize) -> Option<Self::Item> {
99 let bytes_to_skip = n.checked_mul(mem::size_of::<IndexRootItem>())?;
101 self.items_range.start = self.items_range.start.checked_add(bytes_to_skip)?;
102 self.next()
103 }
104
105 fn size_hint(&self) -> (usize, Option<usize>) {
106 let size = self.items_range.len() / mem::size_of::<IndexRootItem>();
107 (size, Some(size))
108 }
109}
110
111impl ExactSizeIterator for IndexRootItemRanges {}
112impl FusedIterator for IndexRootItemRanges {}
113
114impl<B: SplitByteSlice> From<IndexRootKeyNodes<'_, B>> for IndexRootItemRanges {
115 fn from(index_root_key_nodes: IndexRootKeyNodes<'_, B>) -> IndexRootItemRanges {
116 index_root_key_nodes.index_root_item_ranges
117 }
118}
119
120#[derive(Clone)]
129pub struct IndexRootKeyNodes<'h, B: SplitByteSlice> {
130 hive: &'h Hive<B>,
131 index_root_item_ranges: IndexRootItemRanges,
132 leaf_item_ranges: Option<LeafItemRanges>,
133}
134
135impl<'h, B> IndexRootKeyNodes<'h, B>
136where
137 B: SplitByteSlice,
138{
139 pub(crate) fn new(
140 hive: &'h Hive<B>,
141 count: u16,
142 count_field_offset: usize,
143 data_range: Range<usize>,
144 ) -> Result<Self> {
145 let index_root_item_ranges =
146 IndexRootItemRanges::new(count, count_field_offset, data_range)?;
147
148 Ok(Self {
149 hive,
150 index_root_item_ranges,
151 leaf_item_ranges: None,
152 })
153 }
154}
155
156impl<'h, B> Iterator for IndexRootKeyNodes<'h, B>
157where
158 B: SplitByteSlice,
159{
160 type Item = Result<KeyNode<'h, B>>;
161
162 fn next(&mut self) -> Option<Self::Item> {
163 loop {
164 if let Some(leaf_item_ranges) = self.leaf_item_ranges.as_mut() {
165 if let Some(leaf_item_range) = leaf_item_ranges.next() {
166 let key_node =
167 iter_try!(KeyNode::from_leaf_item_range(self.hive, leaf_item_range));
168 return Some(Ok(key_node));
169 }
170 }
171
172 let index_root_item_range = self.index_root_item_ranges.next()?;
175 let leaf_item_ranges = iter_try!(LeafItemRanges::from_index_root_item_range(
176 self.hive,
177 index_root_item_range
178 ));
179 self.leaf_item_ranges = Some(leaf_item_ranges);
180 }
181 }
182}
183
184impl<B> FusedIterator for IndexRootKeyNodes<'_, B> where B: SplitByteSlice {}
185
186pub(crate) struct IndexRootKeyNodesMut<'h, B: SplitByteSliceMut> {
195 hive: &'h mut Hive<B>,
196 index_root_item_ranges: IndexRootItemRanges,
197 leaf_item_ranges: Option<LeafItemRanges>,
198}
199
200impl<'h, B> IndexRootKeyNodesMut<'h, B>
201where
202 B: SplitByteSliceMut,
203{
204 pub(crate) fn new(
205 hive: &'h mut Hive<B>,
206 count: u16,
207 count_field_offset: usize,
208 data_range: Range<usize>,
209 ) -> Result<Self> {
210 let index_root_item_ranges =
211 IndexRootItemRanges::new(count, count_field_offset, data_range)?;
212
213 Ok(Self {
214 hive,
215 index_root_item_ranges,
216 leaf_item_ranges: None,
217 })
218 }
219
220 pub(crate) fn next<'a>(&'a mut self) -> Option<Result<KeyNodeMut<'a, B>>>
221 where
222 'h: 'a,
223 {
224 loop {
225 if let Some(leaf_item_ranges) = self.leaf_item_ranges.as_mut() {
226 if let Some(leaf_item_range) = leaf_item_ranges.next() {
227 let key_node =
228 iter_try!(KeyNodeMut::from_leaf_item_range(self.hive, leaf_item_range));
229 return Some(Ok(key_node));
230 }
231 }
232
233 let index_root_item_range = self.index_root_item_ranges.next()?;
236 let leaf_item_ranges = iter_try!(LeafItemRanges::from_index_root_item_range(
237 self.hive,
238 index_root_item_range
239 ));
240 self.leaf_item_ranges = Some(leaf_item_ranges);
241 }
242 }
243}