dtmt/lib/sdk/src/binary.rs

286 lines
8.3 KiB
Rust

use std::io::{Cursor, Read, Seek, Write};
use color_eyre::Result;
use self::sync::{ReadExt, WriteExt};
pub trait FromBinary: Sized {
fn from_binary<R: Read + Seek>(r: &mut R) -> Result<Self>;
}
pub trait ToBinary {
fn to_binary(&self) -> Result<Vec<u8>>;
}
impl<T: ToBinary> ToBinary for Vec<T> {
fn to_binary(&self) -> Result<Vec<u8>> {
// TODO: Allocations for the vector could be optimized by first
// serializing one value, then calculating the size from that.
let mut bin = Cursor::new(Vec::new());
bin.write_u32(self.len() as u32)?;
for val in self.iter() {
let buf = val.to_binary()?;
bin.write_all(&buf)?;
}
Ok(bin.into_inner())
}
}
impl<T: FromBinary> FromBinary for Vec<T> {
fn from_binary<R: Read + Seek>(r: &mut R) -> Result<Self> {
let size = r.read_u32()? as usize;
let mut list = Vec::with_capacity(size);
for _ in 0..size {
list.push(T::from_binary(r)?);
}
Ok(list)
}
}
pub fn flags_from_bits<T: bitflags::Flags>(bits: T::Bits) -> T
where
<T as bitflags::Flags>::Bits: std::fmt::Binary,
{
if let Some(flags) = T::from_bits(bits) {
flags
} else {
let unknown = bits & !T::all().bits();
tracing::warn!(
"Unknown bits found for '{}': known = {:0b}, unknown = {:0b}",
std::any::type_name::<T>(),
T::all().bits(),
unknown
);
T::from_bits_truncate(bits)
}
}
pub mod sync {
use std::ffi::CStr;
use std::io::{self, Read, Seek, SeekFrom, Write};
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use color_eyre::eyre::{self, WrapErr};
use color_eyre::{Help, Report, Result, SectionExt};
macro_rules! make_read {
($func:ident, $read:ident, $type:ty) => {
fn $read(&mut self) -> io::Result<$type> {
ReadBytesExt::$func::<LittleEndian>(self)
}
fn $func(&mut self) -> Result<$type> {
let res =
ReadExt::$read(self).wrap_err(concat!("failed to read ", stringify!($type)));
if res.is_ok() {
return res;
}
let pos = self.stream_position();
if pos.is_ok() {
res.with_section(|| {
format!("{pos:#X} ({pos})", pos = pos.unwrap()).header("Position: ")
})
} else {
res
}
}
};
}
macro_rules! make_write {
($func:ident, $write:ident, $type:ty) => {
fn $write(&mut self, val: $type) -> io::Result<()> {
WriteBytesExt::$func::<LittleEndian>(self, val)
}
fn $func(&mut self, val: $type) -> Result<()> {
let res = WriteExt::$write(self, val)
.wrap_err(concat!("failed to write ", stringify!($type)));
if res.is_ok() {
return res;
}
let pos = self.stream_position();
if pos.is_ok() {
res.with_section(|| {
format!("{pos:#X} ({pos})", pos = pos.unwrap()).header("Position: ")
})
} else {
res
}
}
};
}
macro_rules! make_skip {
($func:ident, $read:ident, $type:ty) => {
fn $func(&mut self, cmp: $type) -> Result<()> {
let val = ReadExt::$read(self)?;
if val != cmp {
let pos = self.stream_position().unwrap_or(u64::MAX);
tracing::debug!(
pos,
expected = cmp,
actual = val,
"Unexpected value for skipped {}",
stringify!($type)
);
}
Ok(())
}
};
}
pub trait ReadExt: Read + Seek {
fn read_u8(&mut self) -> io::Result<u8> {
ReadBytesExt::read_u8(self)
}
make_read!(read_u16, read_u16_le, u16);
make_read!(read_u32, read_u32_le, u32);
make_read!(read_u64, read_u64_le, u64);
make_skip!(skip_u16, read_u16, u16);
make_skip!(skip_u32, read_u32, u32);
// Implementation based on https://en.wikipedia.com/wiki/LEB128
fn read_uleb128(&mut self) -> io::Result<u64> {
let mut result: u64 = 0;
let mut shift: u64 = 0;
loop {
let byte = ReadExt::read_u8(self)? as u64;
result |= (byte & 0x7f) << shift;
if byte < 0x80 {
return Ok(result);
}
shift += 7;
}
}
fn skip_padding(&mut self) -> io::Result<()> {
let pos = self.stream_position()?;
let padding_size = 16 - (pos % 16);
if padding_size < 16 && padding_size > 0 {
tracing::trace!(pos, padding_size, "Skipping padding");
self.seek(SeekFrom::Current(padding_size as i64))?;
} else {
tracing::trace!(pos, padding_size, "No padding to skip");
}
Ok(())
}
fn read_string_len(&mut self, len: usize) -> Result<String> {
let pos = self.stream_position();
let res = read_string_len(self, len);
if res.is_ok() {
return res;
}
if pos.is_ok() {
res.with_section(|| {
format!("{pos:#X} ({pos})", pos = pos.unwrap()).header("Position: ")
})
} else {
res
}
}
fn read_bool(&mut self) -> Result<bool> {
match ReadExt::read_u8(self)? {
0 => Ok(false),
1 => Ok(true),
v => eyre::bail!("Invalid value for boolean '{}'", v),
}
}
}
pub trait WriteExt: Write + Seek {
fn write_u8(&mut self, val: u8) -> io::Result<()> {
WriteBytesExt::write_u8(self, val)
}
make_write!(write_u32, write_u32_le, u32);
make_write!(write_u64, write_u64_le, u64);
fn write_bool(&mut self, val: bool) -> io::Result<()> {
WriteBytesExt::write_u8(self, if val { 1 } else { 0 })
}
fn write_padding(&mut self) -> io::Result<usize> {
let pos = self.stream_position()?;
let size = 16 - (pos % 16) as usize;
tracing::trace!(padding_size = size, "Writing padding");
if size > 0 && size < 16 {
let buf = vec![0; size];
self.write_all(&buf)?;
Ok(size)
} else {
Ok(0)
}
}
}
impl<R: Read + Seek + ?Sized> ReadExt for R {}
impl<W: Write + Seek + ?Sized> WriteExt for W {}
pub(crate) fn _read_up_to<R>(r: &mut R, buf: &mut Vec<u8>) -> Result<usize>
where
R: Read + Seek,
{
let pos = r.stream_position()?;
let err = {
match r.read_exact(buf) {
Ok(_) => return Ok(buf.len()),
Err(err) if err.kind() == std::io::ErrorKind::UnexpectedEof => {
r.seek(SeekFrom::Start(pos))?;
match r.read_to_end(buf) {
Ok(read) => return Ok(read),
Err(err) => err,
}
}
Err(err) => err,
}
};
Err(err).with_section(|| format!("{pos:#X} ({pos})").header("Position: "))
}
fn read_string_len(mut r: impl Read, len: usize) -> Result<String> {
let mut buf = vec![0; len];
r.read_exact(&mut buf)
.wrap_err_with(|| format!("Failed to read {} bytes", len))?;
let res = match CStr::from_bytes_until_nul(&buf) {
Ok(s) => {
let s = s.to_str()?;
Ok(s.to_string())
}
Err(_) => String::from_utf8(buf.clone()).map_err(Report::new),
};
res.wrap_err("Invalid binary for UTF8 string")
.with_section(|| format!("{}", String::from_utf8_lossy(&buf)).header("ASCI:"))
.with_section(|| format!("{:x?}", buf).header("Bytes:"))
}
}