use anyhow::{anyhow, Result};
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use std::collections::HashMap;
use std::io::{Cursor, Seek, SeekFrom};
use std::vec::Vec;
use super::SerializationState;
pub struct BinArchive {
pub(super) data: Vec<u8>,
pub(super) text_pointers: HashMap<usize, String>,
pub(super) labels: HashMap<usize, Vec<String>>,
pub(super) internal_pointers: HashMap<usize, usize>,
}
fn adjust_pointer(orig: usize, addr: usize, amount: isize) -> usize {
if orig >= addr {
((orig as isize) + amount) as usize
} else {
orig
}
}
impl BinArchive {
pub fn new() -> Self {
BinArchive {
data: Vec::new(),
text_pointers: HashMap::new(),
labels: HashMap::new(),
internal_pointers: HashMap::new(),
}
}
fn label_count(&self) -> u32 {
let mut label_count = 0;
for (_, value) in &self.labels {
label_count += value.len();
}
label_count as u32
}
fn label_start(&self) -> u32 {
let labels_start =
self.data.len() + self.text_pointers.len() * 4 + self.internal_pointers.len() * 4;
(labels_start + (self.label_count() * 8) as usize) as u32
}
fn adjust_pointers(&mut self, addr: usize, amount: isize, anchor: bool) {
let mut adjusted_internal_pointers: HashMap<usize, usize> = HashMap::new();
for (key, value) in &self.internal_pointers {
if anchor && *value == addr {
adjusted_internal_pointers.insert(*key, *value);
} else {
let adjusted_pointer = adjust_pointer(*key, addr, amount);
let adjusted_data_pointer = adjust_pointer(*value, addr, amount);
adjusted_internal_pointers.insert(adjusted_pointer, adjusted_data_pointer);
}
}
let mut adjusted_labels: HashMap<usize, Vec<String>> = HashMap::new();
for (key, value) in &self.labels {
let adjusted = adjust_pointer(*key, addr, amount);
adjusted_labels.insert(adjusted, value.clone());
}
let mut adjusted_text_pointers: HashMap<usize, String> = HashMap::new();
for (key, value) in &self.text_pointers {
let adjusted = adjust_pointer(*key, addr, amount);
adjusted_text_pointers.insert(adjusted, value.clone());
}
self.internal_pointers = adjusted_internal_pointers;
self.text_pointers = adjusted_text_pointers;
self.labels = adjusted_labels;
}
pub fn serialize(&mut self) -> Result<Vec<u8>> {
let label_start = self.label_start();
let mut state = SerializationState::new();
state.load_internal_pointers(&mut self.data, &self.internal_pointers)?;
state.load_labels(&self.labels);
state.load_text_pointers(&mut self.data, &self.text_pointers, label_start)?;
state.assemble(&self.data)
}
pub fn put_u8(&mut self, addr: usize, value: u8) -> Result<()> {
self.validate_address(addr, false)?;
self.data[addr] = value;
Ok(())
}
pub fn put_u16(&mut self, addr: usize, value: u16) -> Result<()> {
self.validate_address(addr, false)?;
self.validate_address(addr + 2, true)?;
self.validate_alignment(addr, 2)?;
let mut cursor = Cursor::new(&mut self.data);
cursor.seek(SeekFrom::Start(addr as u64))?;
cursor.write_u16::<LittleEndian>(value)?;
Ok(())
}
pub fn put_u32(&mut self, addr: usize, value: u32) -> Result<()> {
self.validate_address(addr, false)?;
self.validate_address(addr + 4, true)?;
self.validate_alignment(addr, 4)?;
let mut cursor = Cursor::new(&mut self.data);
cursor.seek(SeekFrom::Start(addr as u64))?;
cursor.write_u32::<LittleEndian>(value)?;
Ok(())
}
pub fn put_f32(&mut self, addr: usize, value: f32) -> Result<()> {
self.validate_address(addr, false)?;
self.validate_address(addr + 4, true)?;
self.validate_alignment(addr, 4)?;
let mut cursor = Cursor::new(&mut self.data);
cursor.seek(SeekFrom::Start(addr as u64))?;
cursor.write_f32::<LittleEndian>(value)?;
Ok(())
}
pub fn put_text(&mut self, addr: usize, text: Option<&String>) -> Result<()> {
self.validate_address(addr, false)?;
self.validate_address(addr + 4, true)?;
self.validate_alignment(addr, 4)?;
match text {
Some(value) => {
self.text_pointers.insert(addr, value.to_string());
}
None => self.clear_text(addr)?,
}
Ok(())
}
pub fn put_label(&mut self, addr: usize, value: &str) -> Result<()> {
self.validate_address(addr, false)?;
self.validate_address(addr + 4, true)?;
self.validate_alignment(addr, 4)?;
let bucket = self.labels.get_mut(&addr);
match bucket {
Some(bucket_value) => {
let string_value = value.to_string();
if !bucket_value.contains(&string_value) {
bucket_value.push(string_value);
}
}
None => {
let mut bucket: Vec<String> = Vec::new();
bucket.push(value.to_string());
self.labels.insert(addr, bucket);
}
}
Ok(())
}
pub fn put_pointer(&mut self, addr: usize, dest: usize) -> Result<()> {
self.validate_address(addr, false)?;
self.validate_address(addr + 4, true)?;
self.validate_alignment(addr, 4)?;
self.internal_pointers.insert(addr, dest);
Ok(())
}
pub fn read_u8(&self, addr: usize) -> Result<u8> {
self.validate_address(addr, false)?;
let mut cursor = Cursor::new(&self.data);
cursor.seek(SeekFrom::Start(addr as u64))?;
let value = cursor.read_u8()?;
Ok(value)
}
pub fn read_u16(&self, addr: usize) -> Result<u16> {
self.validate_address(addr, false)?;
self.validate_address(addr + 2, true)?;
self.validate_alignment(addr, 2)?;
let mut cursor = Cursor::new(&self.data);
cursor.seek(SeekFrom::Start(addr as u64))?;
let value = cursor.read_u16::<LittleEndian>()?;
Ok(value)
}
pub fn read_u32(&self, addr: usize) -> Result<u32> {
self.validate_address(addr, false)?;
self.validate_address(addr + 4, true)?;
self.validate_alignment(addr, 4)?;
let mut cursor = Cursor::new(&self.data);
cursor.seek(SeekFrom::Start(addr as u64))?;
let value = cursor.read_u32::<LittleEndian>()?;
Ok(value)
}
pub fn read_f32(&self, addr: usize) -> Result<f32> {
self.validate_address(addr, false)?;
self.validate_address(addr + 4, true)?;
self.validate_alignment(addr, 4)?;
let mut cursor = Cursor::new(&self.data);
cursor.seek(SeekFrom::Start(addr as u64))?;
let value = cursor.read_f32::<LittleEndian>()?;
Ok(value)
}
pub fn read_string(&self, addr: usize) -> Result<Option<String>> {
self.validate_address(addr, false)?;
self.validate_address(addr + 4, true)?;
self.validate_alignment(addr, 4)?;
Ok(self.text_pointers.get(&addr).cloned())
}
pub fn read_label(&self, addr: usize, index: usize) -> Result<Option<String>> {
self.validate_address(addr, false)?;
self.validate_address(addr + 4, true)?;
self.validate_alignment(addr, 4)?;
Ok(match self.labels.get(&addr) {
Some(value) => value.get(index).cloned(),
None => None,
})
}
pub fn read_pointer(&self, addr: usize) -> Result<Option<usize>> {
self.validate_address(addr, false)?;
self.validate_address(addr + 4, true)?;
self.validate_alignment(addr, 4)?;
Ok(match self.internal_pointers.get(&addr) {
Some(value) => Some(*value),
None => None,
})
}
pub fn clear_pointer(&mut self, addr: usize) -> Result<()> {
self.put_u32(addr, 0)?;
self.internal_pointers.remove(&addr);
Ok(())
}
pub fn clear_text(&mut self, addr: usize) -> Result<()> {
self.put_u32(addr, 0)?;
self.text_pointers.remove(&addr);
Ok(())
}
pub fn clear_labels(&mut self, addr: usize) -> Result<()> {
self.labels.remove(&addr);
Ok(())
}
pub fn clear_label(&mut self, addr: usize, value: &str) {
match self.labels.get_mut(&addr) {
Some(bucket) => {
let index_optional = bucket.iter().position(|x| *x == value);
match index_optional {
Some(index) => {
bucket.remove(index);
}
None => {}
}
}
None => {}
}
}
pub fn size(&self) -> usize {
self.data.len()
}
pub fn allocate_at_end(&mut self, amount:usize) -> Result<()> {
self.validate_alignment(amount, 4)?;
for _ in 0..amount {
self.data.push(0)
}
Ok(())
}
pub fn allocate(&mut self, addr: usize, amount: isize, anchor: bool) {
self.adjust_pointers(addr, amount, anchor);
let mut range_to_add: Vec<u8> = Vec::new();
for _ in 0..amount {
range_to_add.push(0);
}
self.data.splice(addr..addr, range_to_add);
}
pub fn deallocate(&mut self, addr: usize, amount: isize, anchor: bool) -> Result<()> {
if addr >= self.size() || addr + (amount as usize) > self.size() {
return Err(anyhow!(
"Bad deallocation addr:{} amount:{} size:{}",
addr,
amount,
self.size()
));
}
let end = addr + (amount as usize);
self.text_pointers.retain(|&k, _| !(k >= addr && k < end));
self.labels.retain(|&k, _| !(k >= addr && k < end));
self.internal_pointers.retain(|&k, &mut v| {
!(k >= addr && k < end) && !(v >= addr && v < end && !(anchor && v == addr))
});
let empty: Vec<u8> = Vec::new();
self.data.splice(addr..end, empty);
self.adjust_pointers(addr, -amount, anchor);
Ok(())
}
pub fn addr_of_label(&self, label: &str) -> Option<usize> {
for (key, value) in &self.labels {
for string in value {
if string == label {
return Some(*key);
}
}
}
None
}
fn validate_address(&self, address: usize, end_is_valid: bool) -> Result<()> {
if (!end_is_valid && address >= self.data.len())
|| (end_is_valid && address > self.data.len())
{
Err(anyhow!("Bad address 0x{:x}", address))
} else {
Ok(())
}
}
fn validate_alignment(&self, value: usize, bytes: usize) -> Result<()> {
if value % bytes != 0 {
Err(anyhow!(
"Bad alignment for value:{} expected alignment:{}",
value,
bytes
))
} else {
Ok(())
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn allocate_at_end_validation() {
let mut archive = BinArchive::new();
let result = archive.allocate_at_end(6);
assert!(result.is_err());
}
#[test]
fn allocate_at_end_success() {
let mut archive = BinArchive::new();
let result = archive.allocate_at_end(12);
assert!(result.is_ok());
assert_eq!(12, archive.size());
}
#[test]
fn put_read_u32_validation() {
let mut archive = get_test_archive();
let result = archive.put_u32(8, 12);
assert!(result.is_err());
let result = archive.put_u32(3, 12);
assert!(result.is_err());
let result = archive.read_u32(2);
assert!(result.is_err());
let result = archive.read_u32(8);
assert!(result.is_err());
}
#[test]
fn put_read_u32_success() {
let mut archive = get_test_archive();
let result = archive.put_u32(4, 12);
assert!(result.is_ok());
let result = archive.read_u32(4);
assert!(result.is_ok());
assert_eq!(12, result.unwrap());
}
#[test]
fn put_read_short_validation() {
let mut archive = get_test_archive();
let result = archive.put_u16(8, 12);
assert!(result.is_err());
let result = archive.put_u16(3, 12);
assert!(result.is_err());
let result = archive.read_u16(1);
assert!(result.is_err());
let result = archive.read_u16(8);
assert!(result.is_err());
}
#[test]
fn put_read_short_success() {
let mut archive = get_test_archive();
let result = archive.put_u16(6, 12);
assert!(result.is_ok());
let result = archive.read_u16(6);
assert!(result.is_ok());
assert_eq!(12, result.unwrap());
}
#[test]
fn put_read_byte_validation() {
let mut archive = get_test_archive();
let result = archive.put_u8(8, 12);
assert!(result.is_err());
let result = archive.read_u8(8);
assert!(result.is_err());
}
#[test]
fn put_read_byte_success() {
let mut archive = get_test_archive();
let result = archive.put_u8(3, 12);
assert!(result.is_ok());
let result = archive.read_u8(3);
assert!(result.is_ok());
assert_eq!(12, result.unwrap());
}
fn get_test_archive() -> BinArchive {
let mut archive = BinArchive::new();
archive.allocate_at_end(8).unwrap();
archive
}
}