1use core::cmp::Ordering;
5use core::mem;
6use core::ops::Range;
7use core::ptr;
8
9use bitflags::bitflags;
10use zerocopy::byteorder::LittleEndian;
11use zerocopy::{
12 FromBytes, Immutable, IntoBytes, KnownLayout, Ref, SplitByteSlice, SplitByteSliceMut,
13 Unaligned, U16, U32, U64,
14};
15
16use crate::error::{NtHiveError, Result};
17use crate::helpers::byte_subrange;
18use crate::hive::Hive;
19use crate::index_root::IndexRootItemRanges;
20use crate::key_value::KeyValue;
21use crate::key_values_list::KeyValues;
22use crate::leaf::{LeafItemRange, LeafItemRanges};
23use crate::string::NtHiveNameString;
24use crate::subkeys_list::{SubKeyNodes, SubKeyNodesMut};
25
26bitflags! {
27 struct KeyNodeFlags: u16 {
28 const KEY_IS_VOLATILE = 0x0001;
30 const KEY_HIVE_EXIT = 0x0002;
32 const KEY_HIVE_ENTRY = 0x0004;
34 const KEY_NO_DELETE = 0x0008;
36 const KEY_SYM_LINK = 0x0010;
38 const KEY_COMP_NAME = 0x0020;
40 const KEY_PREDEF_HANDLE = 0x0040;
42 const KEY_VIRT_MIRRORED = 0x0080;
44 const KEY_VIRT_TARGET = 0x0100;
46 const KEY_VIRTUAL_STORE = 0x0200;
48 }
49}
50
51#[allow(dead_code)]
53#[derive(FromBytes, Immutable, IntoBytes, KnownLayout, Unaligned)]
54#[repr(packed)]
55struct KeyNodeHeader {
56 signature: [u8; 2],
57 flags: U16<LittleEndian>,
58 timestamp: U64<LittleEndian>,
59 spare: U32<LittleEndian>,
60 parent: U32<LittleEndian>,
61 subkey_count: U32<LittleEndian>,
62 volatile_subkey_count: U32<LittleEndian>,
63 subkeys_list_offset: U32<LittleEndian>,
64 volatile_subkeys_list_offset: U32<LittleEndian>,
65 key_values_count: U32<LittleEndian>,
66 key_values_list_offset: U32<LittleEndian>,
67 key_security_offset: U32<LittleEndian>,
68 class_name_offset: U32<LittleEndian>,
69 max_subkey_name: U32<LittleEndian>,
70 max_subkey_class_name: U32<LittleEndian>,
71 max_value_name: U32<LittleEndian>,
72 max_value_data: U32<LittleEndian>,
73 work_var: U32<LittleEndian>,
74 key_name_length: U16<LittleEndian>,
75 class_name_length: U16<LittleEndian>,
76}
77
78#[derive(Clone, Eq, PartialEq)]
80struct KeyNodeItemRange {
81 header_range: Range<usize>,
82 data_range: Range<usize>,
83}
84
85impl KeyNodeItemRange {
86 fn from_cell_range<B>(hive: &Hive<B>, cell_range: Range<usize>) -> Result<Self>
87 where
88 B: SplitByteSlice,
89 {
90 let header_range =
91 byte_subrange(&cell_range, mem::size_of::<KeyNodeHeader>()).ok_or_else(|| {
92 NtHiveError::InvalidHeaderSize {
93 offset: hive.offset_of_data_offset(cell_range.start),
94 expected: mem::size_of::<KeyNodeHeader>(),
95 actual: cell_range.len(),
96 }
97 })?;
98 let data_range = header_range.end..cell_range.end;
99
100 let key_node_item_range = Self {
101 header_range,
102 data_range,
103 };
104 key_node_item_range.validate_signature(hive)?;
105
106 Ok(key_node_item_range)
107 }
108
109 fn from_leaf_item_range<B>(hive: &Hive<B>, leaf_item_range: LeafItemRange) -> Result<Self>
110 where
111 B: SplitByteSlice,
112 {
113 let key_node_offset = leaf_item_range.key_node_offset(hive);
114 let cell_range = hive.cell_range_from_data_offset(key_node_offset)?;
115 let key_node = Self::from_cell_range(hive, cell_range)?;
116 Ok(key_node)
117 }
118
119 fn binary_search_subkey_in_index_root<B>(
120 &self,
121 hive: &Hive<B>,
122 name: &str,
123 index_root_item_ranges: IndexRootItemRanges,
124 ) -> Option<Result<Self>>
125 where
126 B: SplitByteSlice,
127 {
128 assert!(index_root_item_ranges.len() <= u16::MAX as usize);
131 let mut left = 0i32;
132 let mut right = index_root_item_ranges.len() as i32 - 1;
133
134 while left <= right {
135 let mid = (left + right) / 2;
138
139 let index_root_item_range = index_root_item_ranges.clone().nth(mid as usize).unwrap();
140 let leaf_item_ranges = iter_try!(LeafItemRanges::from_index_root_item_range(
141 hive,
142 index_root_item_range
143 ));
144
145 let leaf_item_range = leaf_item_ranges.clone().next().unwrap();
147 let key_node_item_range = iter_try!(Self::from_leaf_item_range(hive, leaf_item_range));
148 let key_node_name = iter_try!(key_node_item_range.name(hive));
149
150 match key_node_name.partial_cmp(name).unwrap() {
151 Ordering::Equal => return Some(Ok(key_node_item_range)),
152 Ordering::Less => (),
153 Ordering::Greater => {
154 right = mid - 1;
158 continue;
159 }
160 }
161
162 let leaf_item_range = leaf_item_ranges.clone().last().unwrap();
164 let key_node_item_range = iter_try!(Self::from_leaf_item_range(hive, leaf_item_range));
165 let key_node_name = iter_try!(key_node_item_range.name(hive));
166
167 match key_node_name.partial_cmp(name).unwrap() {
168 Ordering::Equal => return Some(Ok(key_node_item_range)),
169 Ordering::Less => {
170 left = mid + 1;
174 continue;
175 }
176 Ordering::Greater => (),
177 }
178
179 return self.binary_search_subkey_in_leaf(hive, name, leaf_item_ranges);
181 }
182
183 None
184 }
185
186 fn binary_search_subkey_in_leaf<B>(
187 &self,
188 hive: &Hive<B>,
189 name: &str,
190 leaf_item_ranges: LeafItemRanges,
191 ) -> Option<Result<Self>>
192 where
193 B: SplitByteSlice,
194 {
195 assert!(leaf_item_ranges.len() <= u16::MAX as usize);
198 let mut left = 0i32;
199 let mut right = leaf_item_ranges.len() as i32 - 1;
200
201 while left <= right {
202 let mid = (left + right) / 2;
204
205 let leaf_item_range = leaf_item_ranges.clone().nth(mid as usize).unwrap();
206 let key_node_item_range = iter_try!(Self::from_leaf_item_range(hive, leaf_item_range));
207 let key_node_name = iter_try!(key_node_item_range.name(hive));
208
209 match key_node_name.partial_cmp(name).unwrap() {
211 Ordering::Equal => return Some(Ok(key_node_item_range)),
212 Ordering::Less => left = mid + 1,
213 Ordering::Greater => right = mid - 1,
214 }
215 }
216
217 None
218 }
219
220 fn class_name<'h, B>(&self, hive: &'h Hive<B>) -> Option<Result<NtHiveNameString<'h>>>
221 where
222 B: SplitByteSlice,
223 {
224 let header = self.header(hive);
225 let class_name_offset = header.class_name_offset.get();
226 if class_name_offset == u32::MAX {
227 return None;
229 }
230
231 let class_name_length = header.class_name_length.get() as usize;
232 let class_name_offset_range =
233 iter_try!(hive.cell_range_from_data_offset(class_name_offset));
234
235 let class_name_range = iter_try!(byte_subrange(
236 &class_name_offset_range,
237 class_name_length
238 )
239 .ok_or_else(|| NtHiveError::InvalidSizeField {
240 offset: hive.offset_of_field(&header.class_name_length),
241 expected: class_name_length,
242 actual: class_name_offset_range.len(),
243 }));
244 let class_name_bytes = &hive.data[class_name_range];
245
246 Some(Ok(NtHiveNameString::Utf16LE(class_name_bytes)))
247 }
248
249 fn header<'h, B>(&self, hive: &'h Hive<B>) -> Ref<&'h [u8], KeyNodeHeader>
250 where
251 B: SplitByteSlice,
252 {
253 Ref::from_bytes(&hive.data[self.header_range.clone()]).unwrap()
254 }
255
256 fn header_mut<'h, B>(&self, hive: &'h mut Hive<B>) -> Ref<&'h mut [u8], KeyNodeHeader>
257 where
258 B: SplitByteSliceMut,
259 {
260 Ref::from_bytes(&mut hive.data[self.header_range.clone()]).unwrap()
261 }
262
263 fn name<'h, B>(&self, hive: &'h Hive<B>) -> Result<NtHiveNameString<'h>>
264 where
265 B: SplitByteSlice,
266 {
267 let header = self.header(hive);
268 let flags = KeyNodeFlags::from_bits_truncate(header.flags.get());
269 let key_name_length = header.key_name_length.get() as usize;
270
271 let key_name_range = byte_subrange(&self.data_range, key_name_length).ok_or_else(|| {
272 NtHiveError::InvalidSizeField {
273 offset: hive.offset_of_field(&header.key_name_length),
274 expected: key_name_length,
275 actual: self.data_range.len(),
276 }
277 })?;
278 let key_name_bytes = &hive.data[key_name_range];
279
280 if flags.contains(KeyNodeFlags::KEY_COMP_NAME) {
281 Ok(NtHiveNameString::Latin1(key_name_bytes))
282 } else {
283 Ok(NtHiveNameString::Utf16LE(key_name_bytes))
284 }
285 }
286
287 fn subkey<B>(&self, hive: &Hive<B>, name: &str) -> Option<Result<Self>>
288 where
289 B: SplitByteSlice,
290 {
291 let cell_range = iter_try!(self.subkeys_cell_range(hive)?);
292 let subkeys = iter_try!(SubKeyNodes::new(hive, cell_range));
293
294 match subkeys {
295 SubKeyNodes::IndexRoot(iter) => {
296 let index_root_item_ranges = IndexRootItemRanges::from(iter);
297 self.binary_search_subkey_in_index_root(hive, name, index_root_item_ranges)
298 }
299 SubKeyNodes::Leaf(iter) => {
300 let leaf_item_ranges = LeafItemRanges::from(iter);
301 self.binary_search_subkey_in_leaf(hive, name, leaf_item_ranges)
302 }
303 }
304 }
305
306 fn subkeys_cell_range<B>(&self, hive: &Hive<B>) -> Option<Result<Range<usize>>>
307 where
308 B: SplitByteSlice,
309 {
310 let header = self.header(hive);
311 let subkeys_list_offset = header.subkeys_list_offset.get();
312 if subkeys_list_offset == u32::MAX {
313 return None;
315 }
316
317 let cell_range = iter_try!(hive.cell_range_from_data_offset(subkeys_list_offset));
318 Some(Ok(cell_range))
319 }
320
321 fn subpath<B>(&self, hive: &Hive<B>, path: &str) -> Option<Result<Self>>
322 where
323 B: SplitByteSlice,
324 {
325 let mut key_node_item_range = self.clone();
326
327 for component in path.split('\\') {
328 if !component.is_empty() {
330 key_node_item_range = iter_try!(key_node_item_range.subkey(hive, component)?);
331 }
332 }
333
334 Some(Ok(key_node_item_range))
335 }
336
337 fn validate_signature<B>(&self, hive: &Hive<B>) -> Result<()>
338 where
339 B: SplitByteSlice,
340 {
341 let header = self.header(hive);
342 let signature = &header.signature;
343 let expected_signature = b"nk";
344
345 if signature == expected_signature {
346 Ok(())
347 } else {
348 Err(NtHiveError::InvalidTwoByteSignature {
349 offset: hive.offset_of_field(signature),
350 expected: expected_signature,
351 actual: *signature,
352 })
353 }
354 }
355
356 fn value<'h, B>(&self, hive: &'h Hive<B>, name: &str) -> Option<Result<KeyValue<'h, B>>>
357 where
358 B: SplitByteSlice,
359 {
360 let mut values = iter_try!(self.values(hive)?);
361
362 values.find(|key_value| {
364 let key_value = match key_value {
365 Ok(key_value) => key_value,
366 Err(_) => return true,
367 };
368 let key_value_name = match key_value.name() {
369 Ok(name) => name,
370 Err(_) => return true,
371 };
372
373 key_value_name == name
374 })
375 }
376
377 fn values<'h, B>(&self, hive: &'h Hive<B>) -> Option<Result<KeyValues<'h, B>>>
378 where
379 B: SplitByteSlice,
380 {
381 let header = self.header(hive);
382 let key_values_list_offset = header.key_values_list_offset.get();
383 if key_values_list_offset == u32::MAX {
384 return None;
386 }
387
388 let cell_range = iter_try!(hive.cell_range_from_data_offset(key_values_list_offset));
389 let count = header.key_values_count.get();
390 let count_field_offset = hive.offset_of_field(&header.key_values_count);
391
392 Some(KeyValues::new(hive, count, count_field_offset, cell_range))
393 }
394}
395
396#[derive(Clone)]
403pub struct KeyNode<'h, B: SplitByteSlice> {
404 hive: &'h Hive<B>,
405 item_range: KeyNodeItemRange,
406}
407
408impl<'h, B> KeyNode<'h, B>
409where
410 B: SplitByteSlice,
411{
412 pub(crate) fn from_cell_range(hive: &'h Hive<B>, cell_range: Range<usize>) -> Result<Self> {
413 let item_range = KeyNodeItemRange::from_cell_range(hive, cell_range)?;
414 Ok(Self { hive, item_range })
415 }
416
417 pub(crate) fn from_leaf_item_range(
418 hive: &'h Hive<B>,
419 leaf_item_range: LeafItemRange,
420 ) -> Result<Self> {
421 let item_range = KeyNodeItemRange::from_leaf_item_range(hive, leaf_item_range)?;
422 Ok(Self { hive, item_range })
423 }
424
425 pub fn class_name(&self) -> Option<Result<NtHiveNameString>> {
427 self.item_range.class_name(self.hive)
428 }
429
430 pub fn name(&self) -> Result<NtHiveNameString> {
432 self.item_range.name(self.hive)
433 }
434
435 pub fn subkey(&self, name: &str) -> Option<Result<KeyNode<'h, B>>> {
437 let item_range = iter_try!(self.item_range.subkey(self.hive, name)?);
438
439 Some(Ok(KeyNode {
440 hive: self.hive,
441 item_range,
442 }))
443 }
444
445 pub fn subkeys(&self) -> Option<Result<SubKeyNodes<'h, B>>> {
447 let cell_range = iter_try!(self.item_range.subkeys_cell_range(self.hive)?);
448 Some(SubKeyNodes::new(self.hive, cell_range))
449 }
450
451 pub fn subpath(&self, path: &str) -> Option<Result<KeyNode<'h, B>>> {
455 let item_range = iter_try!(self.item_range.subpath(self.hive, path)?);
456
457 Some(Ok(KeyNode {
458 hive: self.hive,
459 item_range,
460 }))
461 }
462
463 pub fn value(&self, name: &str) -> Option<Result<KeyValue<'h, B>>> {
465 self.item_range.value(self.hive, name)
466 }
467
468 pub fn values(&self) -> Option<Result<KeyValues<'h, B>>> {
470 self.item_range.values(self.hive)
471 }
472}
473
474impl<B> PartialEq for KeyNode<'_, B>
475where
476 B: SplitByteSlice,
477{
478 fn eq(&self, other: &Self) -> bool {
479 ptr::eq(self.hive, other.hive) && self.item_range == other.item_range
480 }
481}
482
483impl<B> Eq for KeyNode<'_, B> where B: SplitByteSlice {}
484
485pub(crate) struct KeyNodeMut<'h, B: SplitByteSliceMut> {
486 hive: &'h mut Hive<B>,
487 item_range: KeyNodeItemRange,
488}
489
490impl<'h, B> KeyNodeMut<'h, B>
491where
492 B: SplitByteSliceMut,
493{
494 pub(crate) fn from_cell_range(hive: &'h mut Hive<B>, cell_range: Range<usize>) -> Result<Self> {
495 let item_range = KeyNodeItemRange::from_cell_range(hive, cell_range)?;
496 Ok(Self { hive, item_range })
497 }
498
499 pub(crate) fn from_leaf_item_range(
500 hive: &'h mut Hive<B>,
501 leaf_item_range: LeafItemRange,
502 ) -> Result<Self> {
503 let item_range = KeyNodeItemRange::from_leaf_item_range(hive, leaf_item_range)?;
504 Ok(Self { hive, item_range })
505 }
506
507 pub(crate) fn clear_volatile_subkeys(&mut self) -> Result<()> {
508 let mut header = self.item_range.header_mut(self.hive);
509 header.volatile_subkey_count.set(0);
510
511 if let Some(subkeys) = self.subkeys_mut() {
512 let mut subkeys = subkeys?;
513 while let Some(subkey) = subkeys.next() {
514 subkey?.clear_volatile_subkeys()?;
515 }
516 }
517
518 Ok(())
519 }
520
521 pub(crate) fn subkeys_mut(&mut self) -> Option<Result<SubKeyNodesMut<B>>> {
522 let cell_range = iter_try!(self.item_range.subkeys_cell_range(self.hive)?);
523 Some(SubKeyNodesMut::new(self.hive, cell_range))
524 }
525}
526
527#[cfg(test)]
528mod tests {
529 use crate::*;
530
531 #[test]
532 fn test_character_encoding() {
533 let testhive = crate::helpers::tests::testhive_vec();
534 let hive = Hive::new(testhive.as_ref()).unwrap();
535 let root_key_node = hive.root_key_node().unwrap();
536 let key_node = root_key_node
537 .subkey("character-encoding-test")
538 .unwrap()
539 .unwrap();
540
541 let subkey = key_node.subkey("äöü").unwrap().unwrap();
543 assert!(matches!(
544 subkey.name().unwrap(),
545 NtHiveNameString::Latin1(&[0xe4, 0xf6, 0xfc])
546 ));
547
548 let subkey1 = key_node.subkey("A").unwrap().unwrap();
552 let subkey2 = key_node.subkey("a").unwrap().unwrap();
553 assert!(subkey1 == subkey2);
554
555 let subkey1 = key_node.subkey("𐐐").unwrap().unwrap();
558 let subkey2 = key_node.subkey("𐐸").unwrap().unwrap();
559 assert!(subkey1 != subkey2);
560 }
561
562 #[test]
563 fn test_subkey() {
564 let testhive = crate::helpers::tests::testhive_vec();
566 let hive = Hive::new(testhive.as_ref()).unwrap();
567 let root_key_node = hive.root_key_node().unwrap();
568 let key_node = root_key_node.subkey("subkey-test").unwrap().unwrap();
569
570 for i in 0..512 {
571 let subkey_name = format!("key{i}");
572 assert!(
573 matches!(key_node.subkey(&subkey_name), Some(Ok(_))),
574 "Could not find subkey \"{subkey_name}\""
575 );
576 }
577 }
578
579 #[test]
580 fn test_subkeys() {
581 let mut key_names = Vec::with_capacity(512);
584 for i in 0..512 {
585 key_names.push(format!("key{}", i));
586 }
587
588 key_names.sort_unstable();
589
590 let testhive = crate::helpers::tests::testhive_vec();
592 let hive = Hive::new(testhive.as_ref()).unwrap();
593 let root_key_node = hive.root_key_node().unwrap();
594 let key_node = root_key_node.subkey("subkey-test").unwrap().unwrap();
595
596 let subkeys = key_node.subkeys().unwrap().unwrap();
597
598 for (subkey, expected_key_name) in subkeys.zip(key_names.iter()) {
599 let subkey = subkey.unwrap();
600 assert_eq!(subkey.name().unwrap(), expected_key_name.as_str());
601 }
602 }
603
604 #[test]
605 fn test_subpath() {
606 let testhive = crate::helpers::tests::testhive_vec();
607 let hive = Hive::new(testhive.as_ref()).unwrap();
608 let root_key_node = hive.root_key_node().unwrap();
609 let key_node = root_key_node.subkey("subpath-test").unwrap().unwrap();
610
611 assert!(matches!(key_node.subpath("no-subkeys"), Some(Ok(_))));
612 assert!(matches!(key_node.subpath("\\no-subkeys"), Some(Ok(_))));
613 assert!(matches!(key_node.subpath("no-subkeys\\"), Some(Ok(_))));
614 assert!(matches!(key_node.subpath("\\no-subkeys\\"), Some(Ok(_))));
615 assert!(key_node.subpath("no-subkeys\\non-existing").is_none());
616
617 assert!(matches!(
618 key_node.subpath("with-single-level-subkey"),
619 Some(Ok(_))
620 ));
621 assert!(matches!(
622 key_node.subpath("with-single-level-subkey\\subkey"),
623 Some(Ok(_))
624 ));
625 assert!(matches!(
626 key_node.subpath("with-single-level-subkey\\\\subkey"),
627 Some(Ok(_))
628 ));
629 assert!(matches!(
630 key_node.subpath("with-single-level-subkey\\\\subkey\\"),
631 Some(Ok(_))
632 ));
633 assert!(key_node
634 .subpath("with-single-level-subkey\\subkey\\non-existing-too")
635 .is_none());
636
637 assert!(matches!(
638 key_node.subpath("with-two-levels-of-subkeys\\subkey1\\subkey2"),
639 Some(Ok(_))
640 ));
641 assert!(matches!(
642 key_node.subpath("with-two-levels-of-subkeys\\subkey1\\\\subkey2"),
643 Some(Ok(_))
644 ));
645
646 assert!(key_node.subpath("non-existing").is_none());
647 assert!(key_node.subpath("non-existing\\sub").is_none());
648 }
649}