Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

WritableBuffer improvements #335

Merged
merged 1 commit into from
Jul 25, 2021
Merged

WritableBuffer improvements #335

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
119 changes: 1 addition & 118 deletions src/pod.rs
View file
Open in desktop
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
Open in desktop
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