wasi-common: move implementation of args, env methods into StringArrayWriter

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<Rc<Entry>> {
         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
+    }
 }

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;

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<types::Timestamp> {

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<u32>;
+    fn cumulative_size(&self) -> Result<u32>;
+    fn write_to_guest<'a>(
+        &self,
+        buffer: &GuestPtr<'a, u8>,
+        elements: &GuestPtr<'a, GuestPtr<'a, u8>>,
+    ) -> Result<()>;
+}
+
+impl StringArrayWriter for Vec<CString> {
+    fn number_elements(&self) -> Result<u32> {
+        let elems = self.len().try_into()?;
+        Ok(elems)
+    }
+    fn cumulative_size(&self) -> Result<u32> {
+        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(())
+    }
+}