dotnetdll/binary/
heap.rs

1use super::{metadata::index, signature::compressed};
2use crate::utils::hash;
3use scroll::{ctx::StrCtx, Pread, Pwrite, Result};
4use std::collections::HashMap;
5
6// TODO: seal these traits
7
8pub trait Reader<'a> {
9    type Index;
10    type Value;
11
12    const NAME: &'static str;
13
14    fn new(bytes: &'a [u8]) -> Self;
15
16    fn at_index(&self, idx: Self::Index) -> Result<Self::Value>;
17}
18
19macro_rules! heap_reader {
20    ($name:ident, $heap:literal, $index:ty, $value:ty, |$s:ident, $val:ident| $e:expr) => {
21        pub struct $name<'a> {
22            bytes: &'a [u8],
23        }
24
25        impl<'a> Reader<'a> for $name<'a> {
26            type Index = $index;
27            type Value = $value;
28
29            const NAME: &'static str = $heap;
30
31            fn new(bytes: &'a [u8]) -> $name<'a> {
32                $name {
33                    bytes: &bytes,
34                }
35            }
36
37            fn at_index(&$s, $val: Self::Index) -> Result<Self::Value> {
38                $e
39            }
40        }
41
42        impl std::fmt::Debug for $name<'_> {
43            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44                f.write_str(stringify!($name))
45            }
46        }
47    };
48}
49
50fn read_bytes(bytes: &[u8], idx: usize) -> Result<&[u8]> {
51    let mut offset = idx;
52
53    let compressed::Unsigned(size) = bytes.gread(&mut offset)?;
54
55    bytes.pread_with(offset, size as usize)
56}
57
58heap_reader!(StringsReader, "#Strings", index::String, &'a str, |self, idx| self
59    .bytes
60    .pread_with(idx.0, StrCtx::Delimiter(0)));
61heap_reader!(BlobReader, "#Blob", index::Blob, &'a [u8], |self, idx| {
62    read_bytes(self.bytes, idx.0)
63});
64heap_reader!(GUIDReader, "#GUID", index::GUID, [u8; 16], |self, idx| {
65    let mut buf = [0_u8; 16];
66    self.bytes
67        .gread_inout_with(&mut ((idx.0 - 1) * 16), &mut buf, scroll::LE)?;
68    Ok(buf)
69});
70heap_reader!(UserStringReader, "#US", usize, Vec<u16>, |self, idx| {
71    let bytes = read_bytes(self.bytes, idx)?;
72
73    let num_utf16 = (bytes.len() - 1) / 2;
74    let offset = &mut 0;
75    let chars = (0..num_utf16)
76        .map(|_| bytes.gread_with::<u16>(offset, scroll::LE))
77        .collect::<Result<_>>()?;
78
79    Ok(chars)
80});
81
82fn write_bytes(bytes: &[u8]) -> Result<Vec<u8>> {
83    let len = bytes.len();
84
85    let mut buf = vec![0_u8; len];
86
87    let len_size = buf.pwrite(compressed::Unsigned(len as u32), 0)?;
88
89    buf.extend(vec![0_u8; len_size]);
90
91    buf.pwrite(bytes, len_size)?;
92
93    Ok(buf)
94}
95
96pub trait Writer {
97    type Index;
98    type Value: ?Sized;
99
100    fn new() -> Self;
101
102    fn write(&mut self, value: &Self::Value) -> Result<Self::Index>;
103
104    fn into_vec(self) -> Vec<u8>;
105}
106
107macro_rules! heap_writer {
108    ($name:ident, ($buf:expr, $map:expr), $index:ty, $value:ty, |$s:ident, $n:ident| $e:expr) => {
109        pub struct $name {
110            buffer: Vec<u8>,
111            index_cache: HashMap<u64, <Self as Writer>::Index>,
112        }
113
114        impl Writer for $name {
115            type Index = $index;
116            type Value = $value;
117
118            fn new() -> Self {
119                $name {
120                    buffer: $buf,
121                    index_cache: $map,
122                }
123            }
124
125            fn into_vec(self) -> Vec<u8> {
126                self.buffer
127            }
128
129            fn write(&mut $s, $n: &Self::Value) -> Result<Self::Index> {
130                let h = hash($n);
131
132                Ok(match $s.index_cache.get(&h) {
133                    Some(&i) => i,
134                    None => {
135                        let idx = $e;
136                        $s.index_cache.insert(h, idx);
137                        idx
138                    }
139                })
140            }
141        }
142
143        impl std::fmt::Debug for $name {
144            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
145                f.write_str(stringify!($name))
146            }
147        }
148    };
149}
150
151heap_writer!(
152    StringsWriter,
153    (vec![0], HashMap::from([(hash(""), 0.into())])),
154    index::String,
155    str,
156    |self, value| {
157        let start = self.buffer.len();
158        self.buffer.extend(value.as_bytes());
159        self.buffer.push(0_u8);
160        index::String(start)
161    }
162);
163heap_writer!(
164    BlobWriter,
165    (vec![0], HashMap::from([(hash(&[] as &Self::Value), 0.into())])),
166    index::Blob,
167    [u8],
168    |self, value| {
169        let start = self.buffer.len();
170        self.buffer.extend(write_bytes(value)?);
171        index::Blob(start)
172    }
173);
174heap_writer!(
175    GUIDWriter,
176    (vec![], HashMap::new()),
177    index::GUID,
178    [u8; 16],
179    |self, value| {
180        let start = self.buffer.len();
181        self.buffer.extend(value);
182        index::GUID(((start + 1) / 16) + 1)
183    }
184);
185heap_writer!(
186    UserStringWriter,
187    (vec![0], HashMap::from([(hash(&[] as &Self::Value), 0)])),
188    usize,
189    [u16],
190    |self, value| {
191        let final_byte = value.iter().any(|u| {
192            let [high, low] = u.to_be_bytes();
193            high != 0 || matches!(low, 0x01..=0x08 | 0x0E..=0x1F | 0x27 | 0x2D | 0x7F)
194        }) as u8;
195
196        let start = self.buffer.len();
197        self.buffer.extend(write_bytes(
198            &value
199                .iter()
200                .flat_map(|&u| u.to_le_bytes())
201                .chain(std::iter::once(final_byte))
202                .collect::<Vec<_>>(),
203        )?);
204        start
205    }
206);
dotnetdll/binary/heap.rs

dotnetdll/binary/
heap.rs
