diff --git a/crates/c-api/src/wasi.rs b/crates/c-api/src/wasi.rs index bc5456479e65..5c71d8619402 100644 --- a/crates/c-api/src/wasi.rs +++ b/crates/c-api/src/wasi.rs @@ -47,20 +47,16 @@ pub unsafe extern "C" fn wasi_config_set_argv( argc: c_int, argv: *const *const c_char, ) { - let mut config = Box::from_raw(config); - config.builder = config.builder.args( + (*config).builder.args( slice::from_raw_parts(argv, argc as usize) .iter() .map(|a| slice::from_raw_parts(*a as *const u8, CStr::from_ptr(*a).to_bytes().len())), ); - std::mem::forget(config); } #[no_mangle] pub unsafe extern "C" fn wasi_config_inherit_argv(config: *mut wasi_config_t) { - let mut config = Box::from_raw(config); - config.builder = config.builder.inherit_args(); - std::mem::forget(config); + (*config).builder.inherit_args(); } #[no_mangle] @@ -72,23 +68,18 @@ pub unsafe extern "C" fn wasi_config_set_env( ) { let names = slice::from_raw_parts(names, envc as usize); let values = slice::from_raw_parts(values, envc as usize); - let mut config = Box::from_raw(config); - - for i in 0..envc as usize { - config.builder = config.builder.env( - CStr::from_ptr(names[i]).to_bytes(), - CStr::from_ptr(values[i]).to_bytes(), - ); - } - std::mem::forget(config); + (*config).builder.envs( + names + .iter() + .map(|p| CStr::from_ptr(*p).to_bytes()) + .zip(values.iter().map(|p| CStr::from_ptr(*p).to_bytes())), + ); } #[no_mangle] pub unsafe extern "C" fn wasi_config_inherit_env(config: *mut wasi_config_t) { - let mut config = Box::from_raw(config); - config.builder = config.builder.inherit_env(); - std::mem::forget(config); + (*config).builder.inherit_env(); } #[no_mangle] @@ -101,18 +92,14 @@ pub unsafe extern "C" fn wasi_config_set_stdin( None => return false, }; - let mut config = Box::from_raw(config); - config.builder = config.builder.stdin(file); - std::mem::forget(config); + (*config).builder.stdin(file); true } #[no_mangle] pub unsafe extern "C" fn wasi_config_inherit_stdin(config: *mut wasi_config_t) { - let mut config = Box::from_raw(config); - config.builder = config.builder.inherit_stdin(); - std::mem::forget(config); + (*config).builder.inherit_stdin(); } #[no_mangle] @@ -125,18 +112,14 @@ pub unsafe extern "C" fn wasi_config_set_stdout( None => return false, }; - let mut config = Box::from_raw(config); - config.builder = config.builder.stdout(file); - std::mem::forget(config); + (*config).builder.stdout(file); true } #[no_mangle] pub unsafe extern "C" fn wasi_config_inherit_stdout(config: *mut wasi_config_t) { - let mut config = Box::from_raw(config); - config.builder = config.builder.inherit_stdout(); - std::mem::forget(config); + (*config).builder.inherit_stdout(); } #[no_mangle] @@ -149,18 +132,14 @@ pub unsafe extern "C" fn wasi_config_set_stderr( None => return false, }; - let mut config = Box::from_raw(config); - config.builder = config.builder.stderr(file); - std::mem::forget(config); + (*config).builder.stderr(file); true } #[no_mangle] pub unsafe extern "C" fn wasi_config_inherit_stderr(config: *mut wasi_config_t) { - let mut config = Box::from_raw(config); - config.builder = config.builder.inherit_stderr(); - std::mem::forget(config); + (*config).builder.inherit_stderr(); } #[no_mangle] @@ -182,9 +161,7 @@ pub unsafe extern "C" fn wasi_config_preopen_dir( None => return false, }; - let mut config = Box::from_raw(config); - config.builder = config.builder.preopened_dir(dir, guest_path); - std::mem::forget(config); + (*config).builder.preopened_dir(dir, guest_path); true } @@ -202,7 +179,7 @@ pub unsafe extern "C" fn wasi_instance_new( trap: *mut *mut wasm_trap_t, ) -> *mut wasi_instance_t { let store = &(*store).store.borrow(); - let config = Box::from_raw(config); + let mut config = Box::from_raw(config); match config.builder.build() { Ok(ctx) => Box::into_raw(Box::new(wasi_instance_t { diff --git a/crates/wasi-common/src/ctx.rs b/crates/wasi-common/src/ctx.rs index 4c2cd65e6bc2..e9e64a386610 100644 --- a/crates/wasi-common/src/ctx.rs +++ b/crates/wasi-common/src/ctx.rs @@ -60,38 +60,38 @@ impl PendingCString { /// A builder allowing customizable construction of `WasiCtx` instances. pub struct WasiCtxBuilder { - fds: HashMap, - preopens: Vec<(PathBuf, File)>, - args: Vec, - env: HashMap, + fds: Option>, + preopens: Option>, + args: Option>, + env: Option>, } impl WasiCtxBuilder { /// Builder for a new `WasiCtx`. pub fn new() -> Self { - let mut builder = Self { - fds: HashMap::new(), - preopens: Vec::new(), - args: vec![], - env: HashMap::new(), - }; - - builder.fds.insert(0, PendingFdEntry::Thunk(FdEntry::null)); - builder.fds.insert(1, PendingFdEntry::Thunk(FdEntry::null)); - builder.fds.insert(2, PendingFdEntry::Thunk(FdEntry::null)); - - builder + let mut fds = HashMap::new(); + + fds.insert(0, PendingFdEntry::Thunk(FdEntry::null)); + fds.insert(1, PendingFdEntry::Thunk(FdEntry::null)); + fds.insert(2, PendingFdEntry::Thunk(FdEntry::null)); + + Self { + fds: Some(fds), + preopens: Some(Vec::new()), + args: Some(Vec::new()), + env: Some(HashMap::new()), + } } /// Add arguments to the command-line arguments list. /// /// Arguments must be valid UTF-8 with no NUL bytes, or else `WasiCtxBuilder::build()` will fail /// with `Error::EILSEQ`. - pub fn args>(mut self, args: impl IntoIterator) -> Self { - self.args = args - .into_iter() - .map(|arg| arg.as_ref().to_vec().into()) - .collect(); + pub fn args>(&mut self, args: impl IntoIterator) -> &mut Self { + self.args + .as_mut() + .unwrap() + .extend(args.into_iter().map(|a| a.as_ref().to_vec().into())); self } @@ -99,8 +99,11 @@ impl WasiCtxBuilder { /// /// Arguments must be valid UTF-8 with no NUL bytes, or else `WasiCtxBuilder::build()` will fail /// with `Error::EILSEQ`. - pub fn arg>(mut self, arg: S) -> Self { - self.args.push(arg.as_ref().to_vec().into()); + pub fn arg>(&mut self, arg: S) -> &mut Self { + self.args + .as_mut() + .unwrap() + .push(arg.as_ref().to_vec().into()); self } @@ -108,40 +111,46 @@ impl WasiCtxBuilder { /// /// If any arguments from the host process contain invalid UTF-8, `WasiCtxBuilder::build()` will /// fail with `Error::EILSEQ`. - pub fn inherit_args(mut self) -> Self { - self.args = env::args_os().map(PendingCString::OsString).collect(); + 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)); self } /// Inherit stdin from the host process. - pub fn inherit_stdin(mut self) -> Self { + pub fn inherit_stdin(&mut self) -> &mut Self { self.fds + .as_mut() + .unwrap() .insert(0, PendingFdEntry::Thunk(FdEntry::duplicate_stdin)); self } /// Inherit stdout from the host process. - pub fn inherit_stdout(mut self) -> Self { + pub fn inherit_stdout(&mut self) -> &mut Self { self.fds + .as_mut() + .unwrap() .insert(1, PendingFdEntry::Thunk(FdEntry::duplicate_stdout)); self } /// Inherit stdout from the host process. - pub fn inherit_stderr(mut self) -> Self { + pub fn inherit_stderr(&mut self) -> &mut Self { self.fds + .as_mut() + .unwrap() .insert(2, PendingFdEntry::Thunk(FdEntry::duplicate_stderr)); self } /// Inherit the stdin, stdout, and stderr streams from the host process. - pub fn inherit_stdio(mut self) -> Self { - self.fds - .insert(0, PendingFdEntry::Thunk(FdEntry::duplicate_stdin)); - self.fds - .insert(1, PendingFdEntry::Thunk(FdEntry::duplicate_stdout)); - self.fds - .insert(2, PendingFdEntry::Thunk(FdEntry::duplicate_stderr)); + pub fn inherit_stdio(&mut self) -> &mut Self { + let fds = self.fds.as_mut().unwrap(); + fds.insert(0, PendingFdEntry::Thunk(FdEntry::duplicate_stdin)); + fds.insert(1, PendingFdEntry::Thunk(FdEntry::duplicate_stdout)); + fds.insert(2, PendingFdEntry::Thunk(FdEntry::duplicate_stderr)); self } @@ -150,10 +159,10 @@ impl WasiCtxBuilder { /// If any environment variables from the host process contain invalid Unicode (UTF-16 for /// Windows, UTF-8 for other platforms), `WasiCtxBuilder::build()` will fail with /// `Error::EILSEQ`. - pub fn inherit_env(mut self) -> Self { - self.env = std::env::vars_os() - .map(|(k, v)| (k.into(), v.into())) - .collect(); + pub fn inherit_env(&mut self) -> &mut Self { + let env = self.env.as_mut().unwrap(); + env.clear(); + env.extend(std::env::vars_os().map(|(k, v)| (k.into(), v.into()))); self } @@ -161,8 +170,10 @@ impl WasiCtxBuilder { /// /// Environment variable keys and values must be valid UTF-8 with no NUL bytes, or else /// `WasiCtxBuilder::build()` will fail with `Error::EILSEQ`. - pub fn env>(mut self, k: S, v: S) -> Self { + pub fn env>(&mut self, k: S, v: S) -> &mut Self { self.env + .as_mut() + .unwrap() .insert(k.as_ref().to_vec().into(), v.as_ref().to_vec().into()); self } @@ -172,40 +183,49 @@ impl WasiCtxBuilder { /// Environment variable keys and values must be valid UTF-8 with no NUL bytes, or else /// `WasiCtxBuilder::build()` will fail with `Error::EILSEQ`. pub fn envs, T: Borrow<(S, S)>>( - mut self, + &mut self, envs: impl IntoIterator, - ) -> Self { - self.env = envs - .into_iter() - .map(|t| { - let (k, v) = t.borrow(); - (k.as_ref().to_vec().into(), v.as_ref().to_vec().into()) - }) - .collect(); + ) -> &mut Self { + self.env.as_mut().unwrap().extend(envs.into_iter().map(|t| { + let (k, v) = t.borrow(); + (k.as_ref().to_vec().into(), v.as_ref().to_vec().into()) + })); self } /// Provide a File to use as stdin - pub fn stdin(mut self, file: File) -> Self { - self.fds.insert(0, PendingFdEntry::File(file)); + pub fn stdin(&mut self, file: File) -> &mut Self { + self.fds + .as_mut() + .unwrap() + .insert(0, PendingFdEntry::File(file)); self } /// Provide a File to use as stdout - pub fn stdout(mut self, file: File) -> Self { - self.fds.insert(1, PendingFdEntry::File(file)); + pub fn stdout(&mut self, file: File) -> &mut Self { + self.fds + .as_mut() + .unwrap() + .insert(1, PendingFdEntry::File(file)); self } /// Provide a File to use as stderr - pub fn stderr(mut self, file: File) -> Self { - self.fds.insert(2, PendingFdEntry::File(file)); + pub fn stderr(&mut self, file: File) -> &mut Self { + self.fds + .as_mut() + .unwrap() + .insert(2, PendingFdEntry::File(file)); self } /// Add a preopened directory. - pub fn preopened_dir>(mut self, dir: File, guest_path: P) -> Self { - self.preopens.push((guest_path.as_ref().to_owned(), dir)); + pub fn preopened_dir>(&mut self, dir: File, guest_path: P) -> &mut Self { + self.preopens + .as_mut() + .unwrap() + .push((guest_path.as_ref().to_owned(), dir)); self } @@ -213,17 +233,21 @@ 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 returns `Error::EILSEQ`. - pub fn build(self) -> Result { + pub fn build(&mut self) -> Result { // Process arguments and environment variables into `CString`s, failing quickly if they // contain any NUL bytes, or if conversion from `OsString` fails. let args = self .args + .take() + .ok_or(Error::EINVAL)? .into_iter() .map(|arg| arg.into_utf8_cstring()) .collect::>>()?; let env = self .env + .take() + .ok_or(Error::EINVAL)? .into_iter() .map(|(k, v)| { k.into_string().and_then(|mut pair| { @@ -240,7 +264,7 @@ impl WasiCtxBuilder { let mut fds: HashMap = HashMap::new(); // Populate the non-preopen fds. - for (fd, pending) in self.fds { + for (fd, pending) in self.fds.take().ok_or(Error::EINVAL)? { log::debug!("WasiCtx inserting ({:?}, {:?})", fd, pending); match pending { PendingFdEntry::Thunk(f) => { @@ -255,7 +279,7 @@ impl WasiCtxBuilder { // so we start from there. This variable is initially 2, though, because the loop // immediately does the increment and check for overflow. let mut preopen_fd: wasi::__wasi_fd_t = 2; - for (guest_path, dir) in self.preopens { + for (guest_path, dir) in self.preopens.take().ok_or(Error::EINVAL)? { // We do the increment at the beginning of the loop body, so that we don't overflow // unnecessarily if we have exactly the maximum number of file descriptors. preopen_fd = preopen_fd.checked_add(1).ok_or(Error::ENFILE)?; diff --git a/src/commands/run.rs b/src/commands/run.rs index 3e1cb2c5caa6..d5fa7ee0a886 100644 --- a/src/commands/run.rs +++ b/src/commands/run.rs @@ -303,22 +303,25 @@ impl ModuleRegistry { argv: &[String], vars: &[(String, String)], ) -> Result { - let mut cx1 = wasi_common::WasiCtxBuilder::new() - .inherit_stdio() - .args(argv) - .envs(vars); + let mut cx1 = wasi_common::WasiCtxBuilder::new(); + + cx1.inherit_stdio().args(argv).envs(vars); + for (name, file) in preopen_dirs { - cx1 = cx1.preopened_dir(file.try_clone()?, name); + cx1.preopened_dir(file.try_clone()?, name); } + let cx1 = cx1.build()?; let mut cx2 = wasi_common::old::snapshot_0::WasiCtxBuilder::new() .inherit_stdio() .args(argv) .envs(vars); + for (name, file) in preopen_dirs { cx2 = cx2.preopened_dir(file.try_clone()?, name); } + let cx2 = cx2.build()?; Ok(ModuleRegistry {