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.

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

Add write::Object::write_stream #369

Merged
merged 4 commits into from
Aug 30, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
2 changes: 1 addition & 1 deletion src/write/coff.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
4 changes: 3 additions & 1 deletion src/write/elf/object.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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"[..]
Expand Down
3 changes: 2 additions & 1 deletion src/write/elf/writer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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.
Expand Down
14 changes: 13 additions & 1 deletion src/write/mod.rs
Original file line number Diff line number Diff line change
@@ -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::{
Expand Down Expand Up @@ -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<W: io::Write>(&self, w: W) -> result::Result<(), Box<dyn error::Error>> {
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 {
Expand Down
91 changes: 82 additions & 9 deletions src/write/util.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
use std::io;
use std::mem;
use std::vec::Vec;

use crate::pod::{bytes_of, bytes_of_slice, Pod};
Expand All @@ -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]);
Expand Down Expand Up @@ -54,22 +61,88 @@ impl WritableBuffer for Vec<u8> {
}

#[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<W> {
writer: W,
len: usize,
result: Result<(), io::Error>,
}

impl<W> StreamingBuffer<W> {
/// 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<W: io::Write> WritableBuffer for StreamingBuffer<W> {
#[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.
Expand Down Expand Up @@ -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)]
Expand Down