1use std::collections::BTreeMap;
11
12use crate::error::Result;
13
14pub struct OverlayBuffer {
17 base: Vec<u8>,
18 dirty_pages: BTreeMap<u64, Vec<u8>>,
20}
21
22impl OverlayBuffer {
23 pub fn new(base: Vec<u8>) -> Self {
25 Self {
26 base,
27 dirty_pages: BTreeMap::new(),
28 }
29 }
30
31 pub fn apply_page(&mut self, offset: u64, data: Vec<u8>) {
33 self.dirty_pages.insert(offset, data);
34 }
35
36 pub fn read_at(&self, offset: u64, len: usize) -> Vec<u8> {
38 let mut result = Vec::with_capacity(len);
39 for i in 0..len {
40 let pos = offset + i as u64;
41 let byte = self
43 .dirty_pages
44 .iter()
45 .rev()
46 .find(|(&page_offset, page_data)| {
47 pos >= page_offset && pos < page_offset + page_data.len() as u64
48 })
49 .map_or_else(
50 || {
51 usize::try_from(pos)
52 .ok()
53 .and_then(|idx| self.base.get(idx))
54 .copied()
55 .unwrap_or(0)
56 },
57 |(&page_offset, page_data)| {
58 let idx = usize::try_from(pos - page_offset).unwrap_or(0);
59 page_data[idx]
60 },
61 );
62 result.push(byte);
63 }
64 result
65 }
66
67 pub fn len(&self) -> usize {
69 self.base.len()
70 }
71
72 pub fn is_empty(&self) -> bool {
74 self.base.is_empty()
75 }
76
77 pub fn materialize(&self) -> Vec<u8> {
79 self.read_at(0, self.base.len())
80 }
81
82 pub fn dirty_page_count(&self) -> usize {
84 self.dirty_pages.len()
85 }
86}
87
88pub fn replay_transaction_logs(hive_data: Vec<u8>, log_datas: &[Vec<u8>]) -> Result<OverlayBuffer> {
93 let mut overlay = OverlayBuffer::new(hive_data);
94
95 for log_data in log_datas {
96 if log_data.len() < 512 {
97 continue; }
99
100 if &log_data[0..4] != b"regf" {
102 continue;
103 }
104
105 let file_type = crate::bytes::le_u32(log_data, 0x1C);
106
107 match file_type {
108 1 | 6 => {
109 parse_new_format_log(log_data, &mut overlay);
112 }
113 _ => {}
114 }
115 }
116
117 Ok(overlay)
118}
119
120fn parse_new_format_log(log_data: &[u8], overlay: &mut OverlayBuffer) {
122 let mut pos = 512; while pos + 4 <= log_data.len() {
127 if &log_data[pos..pos + 4] == b"HvLE" {
128 if pos + 40 > log_data.len() {
130 break;
131 }
132
133 let size = crate::bytes::le_u32(log_data, pos + 4);
134 let page_count =
136 crate::bytes::le_u32(log_data, pos + 16) as usize;
137
138 let ref_start = pos + 40;
140 let data_start = ref_start + page_count * 8;
141
142 for i in 0..page_count {
143 let ref_offset = ref_start + i * 8;
144 if ref_offset + 8 > log_data.len() {
145 break;
146 }
147 let page_offset =
148 crate::bytes::le_u32(log_data, ref_offset);
149 let page_size = crate::bytes::le_u32(log_data, ref_offset + 4);
150
151 let accumulated_size: u32 = (0..i)
154 .map(|j| {
155 let r = ref_start + j * 8 + 4;
156 u32::from_le_bytes(log_data[r..r + 4].try_into().unwrap_or([0; 4]))
157 })
158 .sum();
159
160 let data_offset = data_start + accumulated_size as usize;
161 let data_end = data_offset + page_size as usize;
162
163 if data_end <= log_data.len() {
164 let file_offset = 4096u64 + u64::from(page_offset);
166 overlay.apply_page(file_offset, log_data[data_offset..data_end].to_vec());
167 }
168 }
169
170 pos += size as usize;
171 } else {
172 pos += 1; }
174 }
175}
176
177#[cfg(test)]
178mod tests {
179 use super::*;
180
181 #[test]
182 fn overlay_read_through() {
183 let base = vec![1, 2, 3, 4, 5, 6, 7, 8];
184 let overlay = OverlayBuffer::new(base);
185 assert_eq!(overlay.read_at(0, 4), vec![1, 2, 3, 4]);
186 }
187
188 #[test]
189 fn overlay_applies_dirty_page() {
190 let base = vec![0; 16];
191 let mut overlay = OverlayBuffer::new(base);
192 overlay.apply_page(4, vec![0xAA, 0xBB, 0xCC, 0xDD]);
193 let result = overlay.materialize();
194 assert_eq!(result[0..4], [0, 0, 0, 0]);
195 assert_eq!(result[4..8], [0xAA, 0xBB, 0xCC, 0xDD]);
196 assert_eq!(result[8..12], [0, 0, 0, 0]);
197 }
198
199 #[test]
200 fn overlay_multiple_pages() {
201 let base = vec![0; 32];
202 let mut overlay = OverlayBuffer::new(base);
203 overlay.apply_page(0, vec![1, 1, 1, 1]);
204 overlay.apply_page(16, vec![2, 2, 2, 2]);
205 let result = overlay.materialize();
206 assert_eq!(result[0], 1);
207 assert_eq!(result[16], 2);
208 assert_eq!(result[8], 0);
209 }
210
211 #[test]
212 fn overlay_dirty_page_count() {
213 let mut overlay = OverlayBuffer::new(vec![0; 16]);
214 assert_eq!(overlay.dirty_page_count(), 0);
215 overlay.apply_page(0, vec![1]);
216 assert_eq!(overlay.dirty_page_count(), 1);
217 }
218}