nt_hive2/hivebin/
cell_iterator.rs

1use std::cell::RefCell;
2use std::fmt::Debug;
3use std::io::{ErrorKind, Seek};
4use std::rc::Rc;
5
6use binread::{derive_binread, BinRead, BinReaderExt, BinResult};
7use derive_getters::Getters;
8use thiserror::Error;
9
10use crate::hivebin::HiveBin;
11use crate::subkeys_list::*;
12use crate::*;
13
14pub struct CellIterator<B>
15where
16    B: BinReaderExt,
17{
18    hive: Rc<RefCell<Hive<B, CleanHive>>>,
19    hivebin_size: usize,
20    consumed_bytes: usize,
21}
22
23impl<B> CellIterator<B>
24where
25    B: BinReaderExt,
26{
27    pub fn new(hivebin: &HiveBin<B>, hive: Rc<RefCell<Hive<B, CleanHive>>>) -> Self {
28        Self {
29            hive,
30            hivebin_size: (*hivebin.size()).try_into().unwrap(),
31
32            // we assume that we already consumed the header
33            consumed_bytes: hivebin.header_size().into(),
34        }
35    }
36
37    fn next_cell(&mut self) -> BinResult<Option<CellSelector>> {
38        const CELL_HEADER_SIZE: usize = 4;
39
40        // if there is not enough space in this hivebin, give up
41        if self.consumed_bytes + CELL_HEADER_SIZE >= self.hivebin_size {
42            return Ok(None);
43        }
44
45        let cell_offset = self.hive.borrow_mut().stream_position().unwrap();
46
47        let header: CellHeader = self.hive.borrow_mut().read_le()?;
48
49        let cell_size = header.size();
50        let content: CellContent = self.hive.borrow_mut().read_le()?;
51        self.consumed_bytes += cell_size;
52
53        let cell_selector = CellSelector {
54            offset: Offset(cell_offset.try_into().unwrap()),
55            header,
56            content,
57        };
58        self.hive.borrow_mut().seek(std::io::SeekFrom::Start(
59            cell_offset + u64::try_from(cell_size).unwrap(),
60        ))?;
61
62        Ok(Some(cell_selector))
63    }
64}
65
66impl<B> Iterator for CellIterator<B>
67where
68    B: BinReaderExt,
69{
70    type Item = CellSelector;
71
72    fn next(&mut self) -> Option<Self::Item> {
73        match self.next_cell() {
74            Ok(v) => v,
75            Err(why) => {
76                if let binread::Error::Io(kind) = &why {
77                    if kind.kind() != ErrorKind::UnexpectedEof {
78                        log::warn!("parser error: {}", why);
79                    }
80                } else {
81                    log::warn!("parser error: {}", why);
82                }
83                None
84            }
85        }
86    }
87}
88
89#[derive(BinRead, Getters)]
90#[getter(get = "pub")]
91pub struct CellSelector {
92    offset: Offset,
93    header: CellHeader,
94    content: CellContent,
95}
96
97#[derive_binread]
98#[derive(Debug)]
99pub enum CellContent {
100    #[br(magic = b"nk")]
101    NK(KeyNode),
102    #[br(magic = b"vk")]
103    VK(KeyValue),
104    #[br(magic = b"sk")]
105    SK,
106    #[br(magic = b"db")]
107    DB,
108
109    #[br(magic = b"li")]
110    LI {
111        #[br(temp)]
112        count: u16,
113
114        #[br(count=count)]
115        items: Vec<IndexLeafItem>,
116    },
117    #[br(magic = b"lf")]
118    LF {
119        #[br(temp)]
120        count: u16,
121
122        #[br(count=count)]
123        items: Vec<FastLeafItem>,
124    },
125
126    #[br(magic = b"lh")]
127    LH {
128        #[br(temp)]
129        count: u16,
130
131        #[br(count=count)]
132        items: Vec<HashLeafItem>,
133    },
134    #[br(magic = b"ri")]
135    RI {
136        #[br(temp)]
137        count: u16,
138
139        #[br(count=count)]
140        items: Vec<IndexRootListElement>,
141    },
142
143    #[allow(clippy::upper_case_acronyms)]
144    UNKNOWN,
145}
146
147#[derive(Error, Debug)]
148pub enum CellLookAheadConversionError {
149    #[error(
150        "tried to extract some type from this cell, which is not actually stored in this cell."
151    )]
152    DifferentCellTypeExpected,
153}
154
155impl CellContent {
156    pub fn is_nk(&self) -> bool {
157        matches!(self, Self::NK(_))
158    }
159}
160
161impl TryInto<KeyNode> for CellSelector {
162    type Error = CellLookAheadConversionError;
163
164    fn try_into(self) -> Result<KeyNode, Self::Error> {
165        match self.content {
166            CellContent::NK(nk) => Ok(nk),
167            _ => Err(CellLookAheadConversionError::DifferentCellTypeExpected),
168        }
169    }
170}