restructure StringArrayWriter trait into StringArray struct

crates/wasi-common/src/ctx.rs

@@ -1,7 +1,7 @@
 use crate::entry::{Entry, EntryHandle};
 use crate::fdpool::FdPool;
 use crate::handle::Handle;
-use crate::string_array_writer::StringArrayWriter;
+use crate::string_array::{StringArray, StringArrayError};
 use crate::sys::osdir::OsDir;
 use crate::sys::stdio::NullDevice;
 use crate::sys::stdio::{Stderr, StderrExt, Stdin, StdinExt, Stdout, StdoutExt};
@@ -12,7 +12,7 @@ use std::borrow::Borrow;
 use std::cell::RefCell;
 use std::collections::HashMap;
 use std::convert::TryFrom;
-use std::ffi::{self, CString, OsString};
+use std::ffi::OsString;
 use std::fs::File;
 use std::path::{Path, PathBuf};
 use std::rc::Rc;
@@ -33,9 +33,12 @@ pub enum WasiCtxBuilderError {
     /// This error is expected to only occur on Windows hosts.
     #[error("provided sequence is not valid UTF-16: {0}")]
     InvalidUtf16(#[from] string::FromUtf16Error),
-    /// Provided sequence of bytes contained an unexpected NUL byte.
+    /// Error constructing arguments
-    #[error("provided sequence contained an unexpected NUL byte")]
+    #[error("while constructing arguments: {0}")]
-    UnexpectedNul(#[from] ffi::NulError),
+    Args(#[source] StringArrayError),
+    /// Error constructing environment
+    #[error("while constructing environment: {0}")]
+    Env(#[source] StringArrayError),
     /// The root of a VirtualDirEntry tree must be a VirtualDirEntry::Directory.
     #[error("the root of a VirtualDirEntry tree at {} must be a VirtualDirEntry::Directory", .0.display())]
     VirtualDirEntryRootNotADirectory(PathBuf),
@@ -65,24 +68,24 @@ impl std::fmt::Debug for PendingEntry {
 }
 
 #[derive(Debug, Eq, Hash, PartialEq)]
-enum PendingCString {
+enum PendingString {
     Bytes(Vec<u8>),
     OsString(OsString),
 }
 
-impl From<Vec<u8>> for PendingCString {
+impl From<Vec<u8>> for PendingString {
     fn from(bytes: Vec<u8>) -> Self {
         Self::Bytes(bytes)
     }
 }
 
-impl From<OsString> for PendingCString {
+impl From<OsString> for PendingString {
     fn from(s: OsString) -> Self {
         Self::OsString(s)
     }
 }
 
-impl PendingCString {
+impl PendingString {
     fn into_string(self) -> WasiCtxBuilderResult<String> {
         let res = match self {
             Self::Bytes(v) => String::from_utf8(v)?,
@@ -100,13 +103,6 @@ impl PendingCString {
         };
         Ok(res)
     }
-
-    /// Create a `CString` containing valid UTF-8.
-    fn into_utf8_cstring(self) -> WasiCtxBuilderResult<CString> {
-        let s = self.into_string()?;
-        let s = CString::new(s)?;
-        Ok(s)
-    }
 }
 
 struct PendingPreopen(Box<dyn FnOnce() -> WasiCtxBuilderResult<Box<dyn Handle>>>);
@@ -130,8 +126,8 @@ pub struct WasiCtxBuilder {
     stdout: Option<PendingEntry>,
     stderr: Option<PendingEntry>,
     preopens: Option<Vec<(PathBuf, PendingPreopen)>>,
-    args: Option<Vec<PendingCString>>,
+    args: Option<Vec<PendingString>>,
-    env: Option<HashMap<PendingCString, PendingCString>>,
+    env: Option<HashMap<PendingString, PendingString>>,
 }
 
 impl WasiCtxBuilder {
@@ -180,7 +176,7 @@ impl WasiCtxBuilder {
     pub fn inherit_args(&mut self) -> &mut Self {
         let args = self.args.as_mut().unwrap();
         args.clear();
-        args.extend(env::args_os().map(PendingCString::OsString));
+        args.extend(env::args_os().map(PendingString::OsString));
         self
     }
 
@@ -324,16 +320,16 @@ impl WasiCtxBuilder {
     /// If any of the arguments or environment variables in this builder cannot be converted into
     /// `CString`s, either due to NUL bytes or Unicode conversions, this will fail.
     pub fn build(&mut self) -> WasiCtxBuilderResult<WasiCtx> {
-        // Process arguments and environment variables into `CString`s, failing quickly if they
+        // Process arguments and environment variables into `String`s, failing quickly if they
         // contain any NUL bytes, or if conversion from `OsString` fails.
         let args = self
             .args
             .take()
             .unwrap()
             .into_iter()
-            .map(|arg| arg.into_utf8_cstring())
+            .map(|arg| arg.into_string())
-            .collect::<WasiCtxBuilderResult<Vec<CString>>>()?;
+            .collect::<WasiCtxBuilderResult<Vec<String>>>()?;
-
+        let args = StringArray::new(args).map_err(WasiCtxBuilderError::Args)?;
         let env = self
             .env
             .take()
@@ -344,14 +340,12 @@ impl WasiCtxBuilder {
                     v.into_string().and_then(|v| {
                         pair.push('=');
                         pair.push_str(v.as_str());
-                        // We have valid UTF-8, but the keys and values have not yet been checked
+                        Ok(pair)
-                        // for NULs, so we do a final check here.
-                        let s = CString::new(pair)?;
-                        Ok(s)
                     })
                 })
             })
-            .collect::<WasiCtxBuilderResult<Vec<CString>>>()?;
+            .collect::<WasiCtxBuilderResult<Vec<String>>>()?;
+        let env = StringArray::new(env).map_err(WasiCtxBuilderError::Env)?;
 
         let mut entries = EntryTable::new();
         // Populate the non-preopen entries.
@@ -441,8 +435,8 @@ impl EntryTable {
 
 pub struct WasiCtx {
     entries: RefCell<EntryTable>,
-    pub(crate) args: Vec<CString>,
+    pub(crate) args: StringArray,
-    pub(crate) env: Vec<CString>,
+    pub(crate) env: StringArray,
 }
 
 impl WasiCtx {
@@ -493,6 +487,7 @@ impl WasiCtx {
         self.entries.borrow_mut().remove(fd).ok_or(Error::Badf)
     }
 
+    /*
     pub(crate) fn args(&self) -> &impl StringArrayWriter {
         &self.args
     }
@@ -500,4 +495,5 @@ impl WasiCtx {
     pub(crate) fn env(&self) -> &impl StringArrayWriter {
         &self.env
     }
+    */
 }

crates/wasi-common/src/lib.rs

@@ -32,7 +32,7 @@ mod path;
 mod poll;
 mod sandboxed_tty_writer;
 pub mod snapshots;
-mod string_array_writer;
+mod string_array;
 mod sys;
 pub mod virtfs;
 pub mod wasi;

crates/wasi-common/src/snapshots/wasi_snapshot_preview1.rs

@@ -1,6 +1,5 @@
 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};
@@ -17,12 +16,11 @@ impl<'a> WasiSnapshotPreview1 for WasiCtx {
         argv: &GuestPtr<'b, GuestPtr<'b, u8>>,
         argv_buf: &GuestPtr<'b, u8>,
     ) -> Result<()> {
-        self.args().write_to_guest(argv_buf, argv)
+        self.args.write_to_guest(argv_buf, argv)
     }
 
     fn args_sizes_get(&self) -> Result<(types::Size, types::Size)> {
-        let args = self.args();
+        Ok((self.args.number_elements, self.args.cumulative_size))
-        Ok((args.number_elements()?, args.cumulative_size()?))
     }
 
     fn environ_get<'b>(
@@ -30,12 +28,11 @@ impl<'a> WasiSnapshotPreview1 for WasiCtx {
         environ: &GuestPtr<'b, GuestPtr<'b, u8>>,
         environ_buf: &GuestPtr<'b, u8>,
     ) -> Result<()> {
-        self.env().write_to_guest(environ_buf, environ)
+        self.env.write_to_guest(environ_buf, environ)
     }
 
     fn environ_sizes_get(&self) -> Result<(types::Size, types::Size)> {
-        let args = self.args();
+        Ok((self.env.number_elements, self.env.cumulative_size))
-        Ok((args.number_elements()?, args.cumulative_size()?))
     }
 
     fn clock_res_get(&self, id: types::Clockid) -> Result<types::Timestamp> {

crates/wasi-common/src/string_array.rs

@@ -0,0 +1,80 @@
+use crate::Error;
+use std::convert::TryInto;
+use std::ffi::CString;
+use wiggle::GuestPtr;
+
+pub struct StringArray {
+    elems: Vec<CString>,
+    pub number_elements: u32,
+    pub cumulative_size: u32,
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum StringArrayError {
+    /// Provided sequence of bytes contained an unexpected NUL byte.
+    #[error("provided sequence contained an unexpected NUL byte")]
+    Nul(#[from] std::ffi::NulError),
+    /// Too many elements: must fit into u32
+    #[error("too many elements")]
+    NumElements,
+    /// Element size: must fit into u32
+    #[error("element too big")]
+    ElemSize,
+    /// Cumulative element size: must fit into u32
+    #[error("cumulative element size too big")]
+    CumElemSize,
+}
+
+impl StringArray {
+    pub fn new(elems: Vec<String>) -> Result<Self, StringArrayError> {
+        let elems = elems
+            .into_iter()
+            .map(|s| CString::new(s))
+            .collect::<Result<Vec<CString>, _>>()?;
+        let number_elements = elems
+            .len()
+            .try_into()
+            .map_err(|_| StringArrayError::NumElements)?;
+        let mut cumulative_size: u32 = 0;
+        for elem in elems.iter() {
+            let elem_bytes = elem
+                .as_bytes_with_nul()
+                .len()
+                .try_into()
+                .map_err(|_| StringArrayError::ElemSize)?;
+            cumulative_size = cumulative_size
+                .checked_add(elem_bytes)
+                .ok_or(StringArrayError::CumElemSize)?;
+        }
+        Ok(Self {
+            elems,
+            number_elements,
+            cumulative_size,
+        })
+    }
+
+    pub fn write_to_guest<'a>(
+        &self,
+        buffer: &GuestPtr<'a, u8>,
+        element_heads: &GuestPtr<'a, GuestPtr<'a, u8>>,
+    ) -> Result<(), Error> {
+        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.elems.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::Inval)?; // Elements don't fit in buffer provided
+            elem_buffer.copy_from_slice(bytes)?;
+            head?.write(
+                elem_buffer
+                    .get(0)
+                    .expect("all elem buffers at least length 1"),
+            )?;
+            cursor += len;
+        }
+        Ok(())
+    }
+}

crates/wasi-common/src/string_array_writer.rs

@@ -1,53 +0,0 @@
-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::Inval)?; // Elements don't fit in buffer provided
-            elem_buffer.copy_from_slice(bytes)?;
-            head?.write(
-                elem_buffer
-                    .get(0)
-                    .expect("all elem buffers at least length 1"),
-            )?;
-            cursor += len;
-        }
-        Ok(())
-    }
-}