diff --git a/src/write/coff.rs b/src/write/coff.rs index c8516c10..1bfb10ac 100644 --- a/src/write/coff.rs +++ b/src/write/coff.rs @@ -579,7 +579,7 @@ impl Object { debug_assert!(aux_len >= symbol.name.len()); let old_len = buffer.len(); buffer.write_bytes(&symbol.name); - buffer.resize(old_len + aux_len, 0); + buffer.resize(old_len + aux_len); } SymbolKind::Section => { debug_assert_eq!(number_of_aux_symbols, 1); diff --git a/src/write/elf/object.rs b/src/write/elf/object.rs index 82ddd531..a11c55a1 100644 --- a/src/write/elf/object.rs +++ b/src/write/elf/object.rs @@ -154,7 +154,9 @@ impl Object { .sections .iter() .map(|section| { - let mut reloc_name = Vec::new(); + let mut reloc_name = Vec::with_capacity( + if is_rela { ".rela".len() } else { ".rel".len() } + section.name.len(), + ); if !section.relocations.is_empty() { reloc_name.extend_from_slice(if is_rela { &b".rela"[..] diff --git a/src/write/elf/writer.rs b/src/write/elf/writer.rs index fba719ba..eb815856 100644 --- a/src/write/elf/writer.rs +++ b/src/write/elf/writer.rs @@ -246,7 +246,7 @@ impl<'a> Writer<'a> { /// Write padding up to the given file offset. pub fn pad_until(&mut self, offset: usize) { debug_assert!(self.buffer.len() <= offset); - self.buffer.resize(offset, 0); + self.buffer.resize(offset); } fn file_header_size(&self) -> usize { @@ -882,6 +882,7 @@ impl<'a> Writer<'a> { return; } self.symtab_shndx_offset = self.reserve(self.symtab_num as usize * 4, 4); + self.symtab_shndx_data.reserve(self.symtab_num as usize * 4); } /// Write the extended section indices for the symbol table. diff --git a/src/write/mod.rs b/src/write/mod.rs index 7bfbe1d9..2ae5fe47 100644 --- a/src/write/mod.rs +++ b/src/write/mod.rs @@ -1,9 +1,10 @@ //! Interface for writing object files. +use std::boxed::Box; use std::collections::HashMap; use std::string::String; use std::vec::Vec; -use std::{error, fmt, result, str}; +use std::{error, fmt, io, result, str}; use crate::endian::{Endianness, U32, U64}; use crate::{ @@ -530,6 +531,17 @@ impl Object { Ok(buffer) } + /// Write the object to a `Write` implementation. + /// + /// It is advisable to use a buffered writer like [`BufWriter`](std::io::BufWriter) + /// instead of an unbuffered writer like [`File`](std::fs::File). + pub fn write_stream(&self, w: W) -> result::Result<(), Box> { + let mut stream = StreamingBuffer::new(w); + self.emit(&mut stream)?; + stream.result()?; + Ok(()) + } + /// Write the object to a `WritableBuffer`. pub fn emit(&self, buffer: &mut dyn WritableBuffer) -> Result<()> { match self.format { diff --git a/src/write/util.rs b/src/write/util.rs index bb6b8597..9d3761a3 100644 --- a/src/write/util.rs +++ b/src/write/util.rs @@ -1,3 +1,5 @@ +use std::io; +use std::mem; use std::vec::Vec; use crate::pod::{bytes_of, bytes_of_slice, Pod}; @@ -6,14 +8,19 @@ use crate::pod::{bytes_of, bytes_of_slice, Pod}; #[allow(clippy::len_without_is_empty)] pub trait WritableBuffer { /// Returns position/offset for data to be written at. + /// + /// Should only be used in debug assertions fn len(&self) -> usize; /// Reserves specified number of bytes in the buffer. - fn reserve(&mut self, additional: usize) -> Result<(), ()>; + /// + /// This will be called exactly once before writing anything to the buffer, + /// and the given size is the exact total number of bytes that will be written. + fn reserve(&mut self, size: usize) -> Result<(), ()>; - /// Writes the specified value at the end of the buffer - /// until the buffer has the specified length. - fn resize(&mut self, new_len: usize, value: u8); + /// Writes zero bytes at the end of the buffer until the buffer + /// has the specified length. + fn resize(&mut self, new_len: usize); /// Writes the specified slice of bytes at the end of the buffer. fn write_bytes(&mut self, val: &[u8]); @@ -54,22 +61,88 @@ impl WritableBuffer for Vec { } #[inline] - fn reserve(&mut self, additional: usize) -> Result<(), ()> { - self.reserve(additional); + fn reserve(&mut self, size: usize) -> Result<(), ()> { + debug_assert!(self.is_empty()); + self.reserve(size); Ok(()) } #[inline] - fn resize(&mut self, new_len: usize, value: u8) { - self.resize(new_len, value); + fn resize(&mut self, new_len: usize) { + debug_assert!(new_len >= self.len()); + self.resize(new_len, 0); } #[inline] fn write_bytes(&mut self, val: &[u8]) { + debug_assert!(self.len() + val.len() <= self.capacity()); self.extend_from_slice(val) } } +/// A [`WritableBuffer`] that streams data to a [`Write`](std::io::Write) implementation. +/// +/// [`Self::result`] must be called to determine if an I/O error occurred during writing. +/// +/// It is advisable to use a buffered writer like [`BufWriter`](std::io::BufWriter) +/// instead of an unbuffered writer like [`File`](std::fs::File). +#[derive(Debug)] +pub struct StreamingBuffer { + writer: W, + len: usize, + result: Result<(), io::Error>, +} + +impl StreamingBuffer { + /// Create a new `StreamingBuffer` backed by the given writer. + pub fn new(writer: W) -> Self { + StreamingBuffer { + writer, + len: 0, + result: Ok(()), + } + } + + /// Unwraps this [`StreamingBuffer`] giving back the original writer. + pub fn into_inner(self) -> W { + self.writer + } + + /// Returns any error that occurred during writing. + pub fn result(&mut self) -> Result<(), io::Error> { + mem::replace(&mut self.result, Ok(())) + } +} + +impl WritableBuffer for StreamingBuffer { + #[inline] + fn len(&self) -> usize { + self.len + } + + #[inline] + fn reserve(&mut self, _size: usize) -> Result<(), ()> { + Ok(()) + } + + #[inline] + fn resize(&mut self, new_len: usize) { + debug_assert!(self.len <= new_len); + while self.len < new_len { + let write_amt = (new_len - self.len - 1) % 1024 + 1; + self.write_bytes(&[0; 1024][..write_amt]); + } + } + + #[inline] + fn write_bytes(&mut self, val: &[u8]) { + if self.result.is_ok() { + self.result = self.writer.write_all(val); + } + self.len += val.len(); + } +} + /// A trait for mutable byte slices. /// /// It provides convenience methods for `Pod` types. @@ -99,7 +172,7 @@ pub(crate) fn align_u64(offset: u64, size: u64) -> u64 { pub(crate) fn write_align(buffer: &mut dyn WritableBuffer, size: usize) { let new_len = align(buffer.len(), size); - buffer.resize(new_len, 0); + buffer.resize(new_len); } #[cfg(test)]