Skip to content

Commit

Permalink
WritableBuffer improvements (#335)
Browse files Browse the repository at this point in the history 
- implement `WritableBuffer` for `Vec`
- rename `WritableBuffer::extend` to `write_bytes` to avoid conflict with `Extend::extend`
- add methods for writing `Pod` types
- replace `BytesMut` with a trait
- move `WritableBuffer` and `BytesMut` to `write::util`
  • Loading branch information
@philipc
philipc authored Jul 25, 2021
1 parent dd132a0 commit 5a3bf7e
Show file tree
Hide file tree
Showing 6 changed files with 173 additions and 176 deletions.
119 changes: 1 addition & 118 deletions src/pod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
)]

use alloc::string::String;
use alloc::vec::Vec;
use core::{fmt, mem, result, slice};

type Result<T> = result::Result<T, ()>;
Expand Down Expand Up @@ -308,90 +307,7 @@ impl<'data> Bytes<'data> {
}
}

/// Trait for writable buffer.
pub trait WritableBuffer {
/// Returns position/offset for data to be written at.
fn len(&self) -> usize;
/// Returns true if buffer contains no data.
fn is_empty(&self) -> bool;
/// Reserves specified number of bytes in the buffer.
fn reserve(&mut self, additional: usize) -> Result<()>;
/// Resizes buffer to the specified length, fills new items
/// with the specified value.
fn resize(&mut self, new_len: usize, value: u8);
/// Extends buffer with the specified slice of bytes.
fn extend(&mut self, val: &[u8]);
}

/// A newtype for byte vectors.
///
/// It provides convenience methods for `Pod` types.
// TODO: should this be an extension trait for `Vec<u8>` instead?
#[derive(Default, Clone, PartialEq, Eq)]
pub(crate) struct BytesMut(pub Vec<u8>);

impl fmt::Debug for BytesMut {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
debug_list_bytes(&self.0, fmt)
}
}

impl BytesMut {
#[inline]
pub fn new() -> Self {
BytesMut(Vec::new())
}

#[inline]
#[allow(dead_code)]
pub fn write<T: Pod>(&mut self, val: &T) {
self.0.extend_from_slice(bytes_of(val))
}

#[inline]
pub fn write_at<T: Pod>(&mut self, offset: usize, val: &T) -> Result<()> {
let src = bytes_of(val);
let dest = self.0.get_mut(offset..).ok_or(())?;
let dest = dest.get_mut(..src.len()).ok_or(())?;
dest.copy_from_slice(src);
Ok(())
}

#[inline]
pub fn as_slice(&self) -> &[u8] {
self.0.as_slice()
}
}

impl WritableBuffer for BytesMut {
#[inline]
fn len(&self) -> usize {
self.0.len()
}

#[inline]
fn is_empty(&self) -> bool {
self.0.is_empty()
}

#[inline]
fn reserve(&mut self, additional: usize) -> Result<()> {
self.0.reserve(additional);
Ok(())
}

#[inline]
fn resize(&mut self, new_len: usize, value: u8) {
self.0.resize(new_len, value);
}

#[inline]
fn extend(&mut self, val: &[u8]) {
self.0.extend_from_slice(val)
}
}

// Only for Debug impl of `Bytes/BytesMut`.
// Only for Debug impl of `Bytes`.
fn debug_list_bytes(bytes: &[u8], fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
let mut list = fmt.debug_list();
list.entries(bytes.iter().take(8).copied().map(DebugByte));
Expand Down Expand Up @@ -605,30 +521,6 @@ mod tests {
assert_eq!(data.read_string_at(3), Err(()));
}

#[test]
fn bytes_mut() {
let data = BytesMut(vec![0x01, 0x23, 0x45, 0x67]);

let mut bytes = data.clone();
bytes.write(&u16::to_be(0x89ab));
assert_eq!(bytes.0, [0x01, 0x23, 0x45, 0x67, 0x89, 0xab]);

let mut bytes = data.clone();
assert_eq!(bytes.write_at(0, &u16::to_be(0x89ab)), Ok(()));
assert_eq!(bytes.0, [0x89, 0xab, 0x45, 0x67]);

let mut bytes = data.clone();
assert_eq!(bytes.write_at(2, &u16::to_be(0x89ab)), Ok(()));
assert_eq!(bytes.0, [0x01, 0x23, 0x89, 0xab]);

assert_eq!(bytes.write_at(3, &u16::to_be(0x89ab)), Err(()));
assert_eq!(bytes.write_at(4, &u16::to_be(0x89ab)), Err(()));
assert_eq!(
BytesMut::default().write_at(0, &u32::to_be(0x89ab)),
Err(())
);
}

#[test]
fn bytes_debug() {
assert_eq!(format!("{:?}", Bytes(&[])), "[]");
Expand All @@ -647,14 +539,5 @@ mod tests {
),
"[0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, ...; 9]"
);

assert_eq!(format!("{:?}", BytesMut(vec![])), "[]");
assert_eq!(
format!(
"{:?}",
BytesMut(vec![0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09])
),
"[0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, ...; 9]"
);
}
}
25 changes: 12 additions & 13 deletions src/write/coff.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@ use std::vec::Vec;

use crate::endian::{LittleEndian as LE, U16Bytes, U32Bytes, U16, U32};
use crate::pe as coff;
use crate::pod::{bytes_of, WritableBuffer};
use crate::write::string::*;
use crate::write::util::*;
use crate::write::*;
Expand Down Expand Up @@ -58,7 +57,7 @@ impl Object {
pub(crate) fn coff_subsection_name(&self, section: &[u8], value: &[u8]) -> Vec<u8> {
let mut name = section.to_vec();
name.push(b'$');
name.extend(value);
name.extend_from_slice(value);
name
}

Expand Down Expand Up @@ -108,7 +107,7 @@ impl Object {
let stub_size = self.architecture.address_size().unwrap().bytes();

let mut name = b".rdata$.refptr.".to_vec();
name.extend(&self.symbols[symbol_id.0].name);
name.extend_from_slice(&self.symbols[symbol_id.0].name);
let section_id = self.add_section(Vec::new(), name, SectionKind::ReadOnlyData);
let section = self.section_mut(section_id);
section.set_data(vec![0; stub_size as usize], u64::from(stub_size));
Expand All @@ -122,7 +121,7 @@ impl Object {
}];

let mut name = b".refptr.".to_vec();
name.extend(&self.symbol(symbol_id).name);
name.extend_from_slice(&self.symbol(symbol_id).name);
let stub_id = self.add_raw_symbol(Symbol {
name,
value: 0,
Expand Down Expand Up @@ -289,7 +288,7 @@ impl Object {
_ => U16::default(),
},
};
buffer.extend(bytes_of(&header));
buffer.write(&header);

// Write section headers.
for (index, section) in self.sections.iter().enumerate() {
Expand Down Expand Up @@ -417,7 +416,7 @@ impl Object {
return Err(Error(format!("invalid section name offset {}", str_offset)));
}
}
buffer.extend(bytes_of(&coff_section));
buffer.write(&coff_section);
}

// Write section data and relocations.
Expand All @@ -426,7 +425,7 @@ impl Object {
if len != 0 {
write_align(buffer, 4);
debug_assert_eq!(section_offsets[index].offset, buffer.len());
buffer.extend(section.data.as_slice());
buffer.write_bytes(&section.data);
}

if !section.relocations.is_empty() {
Expand Down Expand Up @@ -481,7 +480,7 @@ impl Object {
),
typ: U16Bytes::new(LE, typ),
};
buffer.extend(bytes_of(&coff_relocation));
buffer.write(&coff_relocation);
}
}
}
Expand Down Expand Up @@ -571,15 +570,15 @@ impl Object {
let str_offset = strtab.get_offset(symbol_offsets[index].str_id.unwrap());
coff_symbol.name[4..8].copy_from_slice(&u32::to_le_bytes(str_offset as u32));
}
buffer.extend(bytes_of(&coff_symbol));
buffer.write(&coff_symbol);

// Write auxiliary symbols.
match symbol.kind {
SymbolKind::File => {
let aux_len = number_of_aux_symbols as usize * coff::IMAGE_SIZEOF_SYMBOL;
debug_assert!(aux_len >= symbol.name.len());
let old_len = buffer.len();
buffer.extend(&symbol.name);
buffer.write_bytes(&symbol.name);
buffer.resize(old_len + aux_len, 0);
}
SymbolKind::Section => {
Expand All @@ -600,7 +599,7 @@ impl Object {
// TODO: bigobj
high_number: U16Bytes::default(),
};
buffer.extend(bytes_of(&aux));
buffer.write(&aux);
}
_ => {
debug_assert_eq!(number_of_aux_symbols, 0);
Expand All @@ -610,8 +609,8 @@ impl Object {

// Write strtab section.
debug_assert_eq!(strtab_offset, buffer.len());
buffer.extend(&u32::to_le_bytes(strtab_len as u32));
buffer.extend(&strtab_data);
buffer.write_bytes(&u32::to_le_bytes(strtab_len as u32));
buffer.write_bytes(&strtab_data);

debug_assert_eq!(offset, buffer.len());

Expand Down
Loading

0 comments on commit 5a3bf7e

Please sign in to comment.