From 51d88f7899ff8d1c5dc3032a7aac114d265cc905 Mon Sep 17 00:00:00 2001 From: Pat Hickey Date: Mon, 24 Aug 2020 16:38:40 -0700 Subject: [PATCH] wasi-common: move implementation of args, env methods into StringArrayWriter --- crates/wasi-common/src/ctx.rs | 9 ++++ crates/wasi-common/src/lib.rs | 1 + .../src/snapshots/wasi_snapshot_preview1.rs | 47 +++------------- crates/wasi-common/src/string_array_writer.rs | 53 +++++++++++++++++++ 4 files changed, 70 insertions(+), 40 deletions(-) create mode 100644 crates/wasi-common/src/string_array_writer.rs diff --git a/crates/wasi-common/src/ctx.rs b/crates/wasi-common/src/ctx.rs index 08584dbaf5..346af12c28 100644 --- a/crates/wasi-common/src/ctx.rs +++ b/crates/wasi-common/src/ctx.rs @@ -1,6 +1,7 @@ use crate::entry::{Entry, EntryHandle}; use crate::fdpool::FdPool; use crate::handle::Handle; +use crate::string_array_writer::StringArrayWriter; use crate::sys::osdir::OsDir; use crate::sys::stdio::NullDevice; use crate::sys::stdio::{Stderr, StderrExt, Stdin, StdinExt, Stdout, StdoutExt}; @@ -491,4 +492,12 @@ impl WasiCtx { pub(crate) fn remove_entry(&self, fd: types::Fd) -> Result> { self.entries.borrow_mut().remove(fd).ok_or(Error::Badf) } + + pub(crate) fn args(&self) -> &impl StringArrayWriter { + &self.args + } + + pub(crate) fn env(&self) -> &impl StringArrayWriter { + &self.env + } } diff --git a/crates/wasi-common/src/lib.rs b/crates/wasi-common/src/lib.rs index 76e64bb3c0..c2beccdeb6 100644 --- a/crates/wasi-common/src/lib.rs +++ b/crates/wasi-common/src/lib.rs @@ -32,6 +32,7 @@ mod path; mod poll; mod sandboxed_tty_writer; pub mod snapshots; +mod string_array_writer; mod sys; pub mod virtfs; pub mod wasi; diff --git a/crates/wasi-common/src/snapshots/wasi_snapshot_preview1.rs b/crates/wasi-common/src/snapshots/wasi_snapshot_preview1.rs index 3643cf9107..be8994c24d 100644 --- a/crates/wasi-common/src/snapshots/wasi_snapshot_preview1.rs +++ b/crates/wasi-common/src/snapshots/wasi_snapshot_preview1.rs @@ -1,5 +1,6 @@ use crate::entry::{Entry, EntryHandle}; use crate::handle::HandleRights; +use crate::string_array_writer::StringArrayWriter; use crate::sys::clock; use crate::wasi::wasi_snapshot_preview1::WasiSnapshotPreview1; use crate::wasi::{types, AsBytes}; @@ -16,29 +17,12 @@ impl<'a> WasiSnapshotPreview1 for WasiCtx { argv: &GuestPtr<'b, GuestPtr<'b, u8>>, argv_buf: &GuestPtr<'b, u8>, ) -> Result<()> { - let mut argv = argv.clone(); - let mut argv_buf = argv_buf.clone(); - - for arg in &self.args { - let arg_bytes = arg.as_bytes_with_nul(); - let elems = arg_bytes.len().try_into()?; - argv_buf.as_array(elems).copy_from_slice(arg_bytes)?; - argv.write(argv_buf)?; - argv = argv.add(1)?; - argv_buf = argv_buf.add(elems)?; - } - - Ok(()) + self.args().write_to_guest(argv_buf, argv) } fn args_sizes_get(&self) -> Result<(types::Size, types::Size)> { - let argc = self.args.len().try_into()?; - let mut argv_size: types::Size = 0; - for arg in &self.args { - let arg_len = arg.as_bytes_with_nul().len().try_into()?; - argv_size = argv_size.checked_add(arg_len).ok_or(Error::Overflow)?; - } - Ok((argc, argv_size)) + let args = self.args(); + Ok((args.number_elements()?, args.cumulative_size()?)) } fn environ_get<'b>( @@ -46,29 +30,12 @@ impl<'a> WasiSnapshotPreview1 for WasiCtx { environ: &GuestPtr<'b, GuestPtr<'b, u8>>, environ_buf: &GuestPtr<'b, u8>, ) -> Result<()> { - let mut environ = environ.clone(); - let mut environ_buf = environ_buf.clone(); - - for e in &self.env { - let environ_bytes = e.as_bytes_with_nul(); - let elems = environ_bytes.len().try_into()?; - environ_buf.as_array(elems).copy_from_slice(environ_bytes)?; - environ.write(environ_buf)?; - environ = environ.add(1)?; - environ_buf = environ_buf.add(elems)?; - } - - Ok(()) + self.env().write_to_guest(environ_buf, environ) } fn environ_sizes_get(&self) -> Result<(types::Size, types::Size)> { - let environ_count = self.env.len().try_into()?; - let mut environ_size: types::Size = 0; - for environ in &self.env { - let env_len = environ.as_bytes_with_nul().len().try_into()?; - environ_size = environ_size.checked_add(env_len).ok_or(Error::Overflow)?; - } - Ok((environ_count, environ_size)) + let args = self.args(); + Ok((args.number_elements()?, args.cumulative_size()?)) } fn clock_res_get(&self, id: types::Clockid) -> Result { diff --git a/crates/wasi-common/src/string_array_writer.rs b/crates/wasi-common/src/string_array_writer.rs new file mode 100644 index 0000000000..8f44253858 --- /dev/null +++ b/crates/wasi-common/src/string_array_writer.rs @@ -0,0 +1,53 @@ +use crate::{Error, Result}; +use std::convert::TryInto; +use std::ffi::CString; +use wiggle::GuestPtr; + +pub trait StringArrayWriter { + fn number_elements(&self) -> Result; + fn cumulative_size(&self) -> Result; + fn write_to_guest<'a>( + &self, + buffer: &GuestPtr<'a, u8>, + elements: &GuestPtr<'a, GuestPtr<'a, u8>>, + ) -> Result<()>; +} + +impl StringArrayWriter for Vec { + fn number_elements(&self) -> Result { + let elems = self.len().try_into()?; + Ok(elems) + } + fn cumulative_size(&self) -> Result { + let mut total: u32 = 0; + for elem in self.iter() { + let elem_bytes = elem.as_bytes_with_nul().len().try_into()?; + total = total.checked_add(elem_bytes).ok_or(Error::Overflow)?; + } + Ok(total) + } + fn write_to_guest<'a>( + &self, + buffer: &GuestPtr<'a, u8>, + element_heads: &GuestPtr<'a, GuestPtr<'a, u8>>, + ) -> Result<()> { + let element_heads = element_heads.as_array(self.number_elements()?); + let buffer = buffer.as_array(self.cumulative_size()?); + let mut cursor = 0; + for (elem, head) in self.iter().zip(element_heads.iter()) { + let bytes = elem.as_bytes_with_nul(); + let len: u32 = bytes.len().try_into()?; + let elem_buffer = buffer + .get_range(cursor..(cursor + len)) + .ok_or(Error::Overflow)?; + elem_buffer.copy_from_slice(bytes)?; + head?.write( + elem_buffer + .get(0) + .expect("all elem buffers at least length 1"), + )?; + cursor += len; + } + Ok(()) + } +}