1#[cfg(not(feature = "std"))]
4use alloc::string::String;
5
6use crate::error::FormatError;
7
8#[derive(Debug, Clone)]
10pub struct LocalHeap {
11 pub data_segment_size: u64,
13 pub free_list_head_offset: u64,
15 pub data_segment_address: u64,
17}
18
19fn read_offset(data: &[u8], pos: usize, size: u8) -> Result<u64, FormatError> {
20 let s = size as usize;
21 if pos + s > data.len() {
22 return Err(FormatError::UnexpectedEof {
23 expected: pos + s,
24 available: data.len(),
25 });
26 }
27 let slice = &data[pos..pos + s];
28 Ok(match size {
29 2 => u16::from_le_bytes([slice[0], slice[1]]) as u64,
30 4 => u32::from_le_bytes([slice[0], slice[1], slice[2], slice[3]]) as u64,
31 8 => u64::from_le_bytes([
32 slice[0], slice[1], slice[2], slice[3], slice[4], slice[5], slice[6], slice[7],
33 ]),
34 _ => return Err(FormatError::InvalidOffsetSize(size)),
35 })
36}
37
38impl LocalHeap {
39 pub fn parse(
41 file_data: &[u8],
42 offset: usize,
43 offset_size: u8,
44 length_size: u8,
45 ) -> Result<LocalHeap, FormatError> {
46 let ls = length_size as usize;
48 let os = offset_size as usize;
49 let total = 8 + ls * 2 + os;
50 if offset + total > file_data.len() {
51 return Err(FormatError::UnexpectedEof {
52 expected: offset + total,
53 available: file_data.len(),
54 });
55 }
56
57 if &file_data[offset..offset + 4] != b"HEAP" {
58 return Err(FormatError::InvalidLocalHeapSignature);
59 }
60
61 let version = file_data[offset + 4];
62 if version != 0 {
63 return Err(FormatError::InvalidLocalHeapVersion(version));
64 }
65
66 let mut pos = offset + 8;
67 let data_segment_size = read_offset(file_data, pos, length_size)?;
68 pos += ls;
69 let free_list_head_offset = read_offset(file_data, pos, length_size)?;
70 pos += ls;
71 let data_segment_address = read_offset(file_data, pos, offset_size)?;
72
73 Ok(LocalHeap {
74 data_segment_size,
75 free_list_head_offset,
76 data_segment_address,
77 })
78 }
79
80 pub fn read_string(&self, file_data: &[u8], string_offset: u64) -> Result<String, FormatError> {
82 let seg_addr = self.data_segment_address as usize;
83 let str_start = seg_addr + string_offset as usize;
84 let seg_end = seg_addr + self.data_segment_size as usize;
85
86 if str_start >= file_data.len() || str_start >= seg_end {
87 return Err(FormatError::UnexpectedEof {
88 expected: str_start + 1,
89 available: file_data.len(),
90 });
91 }
92
93 let search_end = seg_end.min(file_data.len());
95 let mut end = str_start;
96 while end < search_end && file_data[end] != 0 {
97 end += 1;
98 }
99
100 if end >= search_end {
101 return Err(FormatError::UnexpectedEof {
102 expected: end + 1,
103 available: search_end,
104 });
105 }
106
107 let s = core::str::from_utf8(&file_data[str_start..end])
108 .map_err(|_| FormatError::InvalidLocalHeapSignature)?;
109 Ok(String::from(s))
110 }
111}
112
113#[cfg(test)]
114mod tests {
115 use super::*;
116
117 fn build_heap_file(
118 heap_offset: usize,
119 data_seg_offset: usize,
120 strings: &[&str],
121 offset_size: u8,
122 length_size: u8,
123 ) -> Vec<u8> {
124 let mut data_seg = Vec::new();
126 for s in strings {
127 data_seg.extend_from_slice(s.as_bytes());
128 data_seg.push(0); }
130 let data_seg_size = data_seg.len();
131
132 let total_size = data_seg_offset + data_seg_size + 64;
133 let mut file = vec![0u8; total_size];
134
135 let mut pos = heap_offset;
137 file[pos..pos + 4].copy_from_slice(b"HEAP");
138 pos += 4;
139 file[pos] = 0; pos += 1;
141 pos += 3;
143 write_val(&mut file, pos, data_seg_size as u64, length_size);
145 pos += length_size as usize;
146 write_val(&mut file, pos, 0xFFFFFFFF, length_size);
148 pos += length_size as usize;
149 write_val(&mut file, pos, data_seg_offset as u64, offset_size);
151
152 file[data_seg_offset..data_seg_offset + data_seg_size].copy_from_slice(&data_seg);
154
155 file
156 }
157
158 fn write_val(buf: &mut [u8], pos: usize, val: u64, size: u8) {
159 match size {
160 4 => buf[pos..pos + 4].copy_from_slice(&(val as u32).to_le_bytes()),
161 8 => buf[pos..pos + 8].copy_from_slice(&val.to_le_bytes()),
162 _ => panic!("test"),
163 }
164 }
165
166 #[test]
167 fn parse_heap_header() {
168 let file = build_heap_file(0, 100, &["hello", "world"], 8, 8);
169 let heap = LocalHeap::parse(&file, 0, 8, 8).unwrap();
170 assert_eq!(heap.data_segment_address, 100);
171 assert_eq!(heap.data_segment_size, 12); }
173
174 #[test]
175 fn read_string_at_offset_0() {
176 let file = build_heap_file(0, 100, &["hello", "world"], 8, 8);
177 let heap = LocalHeap::parse(&file, 0, 8, 8).unwrap();
178 let s = heap.read_string(&file, 0).unwrap();
179 assert_eq!(s, "hello");
180 }
181
182 #[test]
183 fn read_string_at_offset_6() {
184 let file = build_heap_file(0, 100, &["hello", "world"], 8, 8);
185 let heap = LocalHeap::parse(&file, 0, 8, 8).unwrap();
186 let s = heap.read_string(&file, 6).unwrap();
187 assert_eq!(s, "world");
188 }
189
190 #[test]
191 fn invalid_signature() {
192 let mut file = build_heap_file(0, 100, &["x"], 8, 8);
193 file[0] = b'X';
194 let err = LocalHeap::parse(&file, 0, 8, 8).unwrap_err();
195 assert_eq!(err, FormatError::InvalidLocalHeapSignature);
196 }
197
198 #[test]
199 fn read_string_past_segment() {
200 let file = build_heap_file(0, 100, &["hi"], 8, 8);
201 let heap = LocalHeap::parse(&file, 0, 8, 8).unwrap();
202 let err = heap.read_string(&file, 100).unwrap_err();
203 assert!(matches!(err, FormatError::UnexpectedEof { .. }));
204 }
205
206 #[test]
207 fn parse_heap_4byte_offsets() {
208 let file = build_heap_file(0, 80, &["test"], 4, 4);
209 let heap = LocalHeap::parse(&file, 0, 4, 4).unwrap();
210 assert_eq!(heap.data_segment_address, 80);
211 let s = heap.read_string(&file, 0).unwrap();
212 assert_eq!(s, "test");
213 }
214
215 #[test]
216 fn invalid_version() {
217 let mut file = build_heap_file(0, 100, &["x"], 8, 8);
218 file[4] = 1; let err = LocalHeap::parse(&file, 0, 8, 8).unwrap_err();
220 assert_eq!(err, FormatError::InvalidLocalHeapVersion(1));
221 }
222}