From 5fd49afedb66bbb1a04209ac583285517d8050ff Mon Sep 17 00:00:00 2001
From: Christopher Serr <christopher.serr@gmail.com>
Date: Thu, 30 Apr 2020 18:19:56 +0200
Subject: [PATCH] Allow using virtual files for stdin/out/err

This adds the ability to use virtual files not only for the file system
in WASI, but also for stdin, stdout and stderr.
---
 crates/wasi-common/src/ctx.rs | 54 +++++++++++++++++++++--------------
 1 file changed, 33 insertions(+), 21 deletions(-)

diff --git a/crates/wasi-common/src/ctx.rs b/crates/wasi-common/src/ctx.rs
index de66d81814c3..d53c315874df 100644
--- a/crates/wasi-common/src/ctx.rs
+++ b/crates/wasi-common/src/ctx.rs
@@ -2,9 +2,8 @@ use crate::entry::{Entry, EntryHandle};
 use crate::fdpool::FdPool;
 use crate::handle::Handle;
 use crate::sys::oshandle::{OsHandle, OsHandleExt};
-use crate::virtfs::{VirtualDir, VirtualDirEntry};
-use crate::wasi::types;
-use crate::wasi::{Errno, Result};
+use crate::virtfs::{FileContents, InMemoryFile, VirtualDir, VirtualDirEntry};
+use crate::wasi::{types, Errno, Result};
 use std::borrow::Borrow;
 use std::cell::RefCell;
 use std::collections::HashMap;
@@ -45,6 +44,7 @@ type WasiCtxBuilderResult<T> = std::result::Result<T, WasiCtxBuilderError>;
 enum PendingEntry {
     Thunk(fn() -> io::Result<OsHandle>),
     File(File),
+    Virtual(Box<dyn FileContents>),
 }
 
 impl std::fmt::Debug for PendingEntry {
@@ -56,6 +56,7 @@ impl std::fmt::Debug for PendingEntry {
                 f as *const fn() -> io::Result<OsHandle>
             ),
             Self::File(f) => write!(fmt, "PendingEntry::File({:?})", f),
+            Self::Virtual(_) => write!(fmt, "PendingEntry::Virtual(...)"),
         }
     }
 }
@@ -229,24 +230,42 @@ impl WasiCtxBuilder {
         self
     }
 
-    /// Provide a File to use as stdin
+    /// Provide a File to use as stdin.
     pub fn stdin(&mut self, file: File) -> &mut Self {
         self.stdin = Some(PendingEntry::File(file));
         self
     }
 
-    /// Provide a File to use as stdout
+    /// Provide a File to use as stdout.
     pub fn stdout(&mut self, file: File) -> &mut Self {
         self.stdout = Some(PendingEntry::File(file));
         self
     }
 
-    /// Provide a File to use as stderr
+    /// Provide a File to use as stderr.
     pub fn stderr(&mut self, file: File) -> &mut Self {
         self.stderr = Some(PendingEntry::File(file));
         self
     }
 
+    /// Provide a virtual file to use as stdin.
+    pub fn stdin_virt(&mut self, file: Box<dyn FileContents>) -> &mut Self {
+        self.stdin = Some(PendingEntry::Virtual(file));
+        self
+    }
+
+    /// Provide a virtual file to use as stdout.
+    pub fn stdout_virt(&mut self, file: Box<dyn FileContents>) -> &mut Self {
+        self.stdout = Some(PendingEntry::Virtual(file));
+        self
+    }
+
+    /// Provide a virtual file to use as stderr.
+    pub fn stderr_virt(&mut self, file: Box<dyn FileContents>) -> &mut Self {
+        self.stderr = Some(PendingEntry::Virtual(file));
+        self
+    }
+
     /// Add a preopened directory.
     pub fn preopened_dir<P: AsRef<Path>>(&mut self, dir: File, guest_path: P) -> &mut Self {
         self.preopens.as_mut().unwrap().push((
@@ -334,22 +353,15 @@ impl WasiCtxBuilder {
             self.stderr.take().unwrap(),
         ] {
             log::debug!("WasiCtx inserting entry {:?}", pending);
-            let fd = match pending {
-                PendingEntry::Thunk(f) => {
-                    let handle = EntryHandle::new(f()?);
-                    let entry = Entry::from(handle)?;
-                    entries
-                        .insert(entry)
-                        .ok_or(WasiCtxBuilderError::TooManyFilesOpen)?
-                }
-                PendingEntry::File(f) => {
-                    let handle = EntryHandle::new(OsHandle::from(f));
-                    let entry = Entry::from(handle)?;
-                    entries
-                        .insert(entry)
-                        .ok_or(WasiCtxBuilderError::TooManyFilesOpen)?
-                }
+            let handle = match pending {
+                PendingEntry::Thunk(f) => EntryHandle::new(f()?),
+                PendingEntry::File(f) => EntryHandle::new(OsHandle::from(f)),
+                PendingEntry::Virtual(f) => EntryHandle::new(InMemoryFile::new(f)),
             };
+            let entry = Entry::from(handle)?;
+            let fd = entries
+                .insert(entry)
+                .ok_or(WasiCtxBuilderError::TooManyFilesOpen)?;
             log::debug!("WasiCtx inserted at {:?}", fd);
         }
         // Then add the preopen entries.