use crate::error::{Error, Result};
use crate::io::Cursor;
use crate::storage::Storage;
const HEAP_SIGNATURE: [u8; 4] = *b"HEAP";
#[derive(Debug, Clone)]
pub struct LocalHeap {
pub data_segment_size: u64,
pub free_list_offset: u64,
pub data_segment_address: u64,
}
impl LocalHeap {
pub fn parse(cursor: &mut Cursor, offset_size: u8, length_size: u8) -> Result<Self> {
let sig = cursor.read_bytes(4)?;
if sig != HEAP_SIGNATURE {
return Err(Error::InvalidLocalHeapSignature);
}
let version = cursor.read_u8()?;
if version != 0 {
return Err(Error::UnsupportedLocalHeapVersion(version));
}
cursor.skip(3)?;
let data_segment_size = cursor.read_length(length_size)?;
let free_list_offset = cursor.read_length(length_size)?;
let data_segment_address = cursor.read_offset(offset_size)?;
Ok(LocalHeap {
data_segment_size,
free_list_offset,
data_segment_address,
})
}
pub fn parse_at_storage(
storage: &dyn Storage,
address: u64,
offset_size: u8,
length_size: u8,
) -> Result<Self> {
let header_len = 4
+ 1
+ 3
+ usize::from(length_size)
+ usize::from(length_size)
+ usize::from(offset_size);
let bytes = storage.read_range(address, header_len)?;
let mut cursor = Cursor::new(bytes.as_ref());
Self::parse(&mut cursor, offset_size, length_size)
}
pub fn get_string(&self, offset: u64, file_data: &[u8]) -> Result<String> {
let abs = self
.data_segment_address
.checked_add(offset)
.ok_or(Error::OffsetOutOfBounds(offset))?;
let start = abs as usize;
if start >= file_data.len() {
return Err(Error::OffsetOutOfBounds(abs));
}
let segment_end = (self.data_segment_address as usize)
.saturating_add(self.data_segment_size as usize)
.min(file_data.len());
let search_region = &file_data[start..segment_end];
let null_pos = search_region.iter().position(|&b| b == 0).ok_or_else(|| {
Error::InvalidData("local heap string missing null terminator".into())
})?;
let s = std::str::from_utf8(&search_region[..null_pos])
.map_err(|e| Error::InvalidData(format!("invalid UTF-8 in local heap string: {e}")))?;
Ok(s.to_string())
}
pub fn get_string_storage(&self, offset: u64, storage: &dyn Storage) -> Result<String> {
if offset >= self.data_segment_size {
return Err(Error::OffsetOutOfBounds(offset));
}
let available = self
.data_segment_size
.checked_sub(offset)
.ok_or(Error::OffsetOutOfBounds(offset))?;
let len = usize::try_from(available).map_err(|_| {
Error::InvalidData("local heap string region exceeds platform usize capacity".into())
})?;
let abs = self
.data_segment_address
.checked_add(offset)
.ok_or(Error::OffsetOutOfBounds(offset))?;
let bytes = storage.read_range(abs, len)?;
let null_pos = bytes.iter().position(|&b| b == 0).ok_or_else(|| {
Error::InvalidData("local heap string missing null terminator".into())
})?;
let s = std::str::from_utf8(&bytes[..null_pos])
.map_err(|e| Error::InvalidData(format!("invalid UTF-8 in local heap string: {e}")))?;
Ok(s.to_string())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn build_heap_header(
data_segment_size: u64,
free_list_offset: u64,
data_segment_address: u64,
) -> Vec<u8> {
let mut buf = Vec::new();
buf.extend_from_slice(b"HEAP");
buf.push(0); buf.extend_from_slice(&[0, 0, 0]); buf.extend_from_slice(&data_segment_size.to_le_bytes());
buf.extend_from_slice(&free_list_offset.to_le_bytes());
buf.extend_from_slice(&data_segment_address.to_le_bytes());
buf
}
#[test]
fn test_parse_local_heap() {
let data = build_heap_header(256, 128, 0x2000);
let mut cursor = Cursor::new(&data);
let heap = LocalHeap::parse(&mut cursor, 8, 8).unwrap();
assert_eq!(heap.data_segment_size, 256);
assert_eq!(heap.free_list_offset, 128);
assert_eq!(heap.data_segment_address, 0x2000);
}
#[test]
fn test_parse_local_heap_4byte() {
let mut buf = Vec::new();
buf.extend_from_slice(b"HEAP");
buf.push(0); buf.extend_from_slice(&[0, 0, 0]); buf.extend_from_slice(&64u32.to_le_bytes()); buf.extend_from_slice(&32u32.to_le_bytes()); buf.extend_from_slice(&0x400u32.to_le_bytes());
let mut cursor = Cursor::new(&buf);
let heap = LocalHeap::parse(&mut cursor, 4, 4).unwrap();
assert_eq!(heap.data_segment_size, 64);
assert_eq!(heap.free_list_offset, 32);
assert_eq!(heap.data_segment_address, 0x400);
}
#[test]
fn test_bad_signature() {
let mut data = build_heap_header(256, 128, 0x2000);
data[0] = b'X'; let mut cursor = Cursor::new(&data);
assert!(matches!(
LocalHeap::parse(&mut cursor, 8, 8),
Err(Error::InvalidLocalHeapSignature)
));
}
#[test]
fn test_bad_version() {
let mut data = build_heap_header(256, 128, 0x2000);
data[4] = 1; let mut cursor = Cursor::new(&data);
assert!(matches!(
LocalHeap::parse(&mut cursor, 8, 8),
Err(Error::UnsupportedLocalHeapVersion(1))
));
}
#[test]
fn test_get_string() {
let mut file_data = vec![0u8; 200];
let seg_start = 100usize;
file_data[seg_start..seg_start + 6].copy_from_slice(b"hello\0");
file_data[seg_start + 6..seg_start + 12].copy_from_slice(b"world\0");
let heap = LocalHeap {
data_segment_size: 100,
free_list_offset: 50,
data_segment_address: seg_start as u64,
};
assert_eq!(heap.get_string(0, &file_data).unwrap(), "hello");
assert_eq!(heap.get_string(6, &file_data).unwrap(), "world");
}
#[test]
fn test_get_string_out_of_bounds() {
let file_data = vec![0u8; 50];
let heap = LocalHeap {
data_segment_size: 100,
free_list_offset: 0,
data_segment_address: 100, };
assert!(heap.get_string(0, &file_data).is_err());
}
#[test]
fn test_get_string_missing_null() {
let mut file_data = vec![0xFFu8; 200];
file_data[100..105].copy_from_slice(b"abcde");
let heap = LocalHeap {
data_segment_size: 5,
free_list_offset: 0,
data_segment_address: 100,
};
assert!(heap.get_string(0, &file_data).is_err());
}
}