From 4a574c14ebbc11364f736859ea6d0b747d352425 Mon Sep 17 00:00:00 2001
From: Pat Hickey <pat@moreproductive.org>
Date: Tue, 5 Jan 2021 18:21:01 -0800
Subject: [PATCH] wasi-common: port to use wiggle::Trap

---
 .../src/snapshots/wasi_snapshot_preview1.rs   | 11 +++++----
 .../src/snapshots/wasi_unstable.rs            |  4 ++--
 crates/wasi-common/src/wasi.rs                |  2 +-
 crates/wasi/src/lib.rs                        | 23 +------------------
 crates/wasi/src/old/snapshot_0.rs             |  9 +-------
 5 files changed, 12 insertions(+), 37 deletions(-)

diff --git a/crates/wasi-common/src/snapshots/wasi_snapshot_preview1.rs b/crates/wasi-common/src/snapshots/wasi_snapshot_preview1.rs
index 3e0bd6d4fc..3d905159c7 100644
--- a/crates/wasi-common/src/snapshots/wasi_snapshot_preview1.rs
+++ b/crates/wasi-common/src/snapshots/wasi_snapshot_preview1.rs
@@ -714,10 +714,13 @@ impl<'a> WasiSnapshotPreview1 for WasiCtx {
         Ok(nevents)
     }
 
-    fn proc_exit(&self, _rval: types::Exitcode) -> std::result::Result<(), ()> {
-        // proc_exit is special in that it's expected to unwind the stack, which
-        // typically requires runtime-specific logic.
-        unimplemented!("runtimes are expected to override this implementation")
+    fn proc_exit(&self, status: types::Exitcode) -> wiggle::Trap {
+        // Check that the status is within WASI's range.
+        if status < 126 {
+            wiggle::Trap::I32(status as i32)
+        } else {
+            wiggle::Trap::String("exit with invalid exit status outside of [0..126)".to_owned())
+        }
     }
 
     fn proc_raise(&self, _sig: types::Signal) -> Result<()> {
diff --git a/crates/wasi-common/src/snapshots/wasi_unstable.rs b/crates/wasi-common/src/snapshots/wasi_unstable.rs
index 6c05c21427..9029dd6ae4 100644
--- a/crates/wasi-common/src/snapshots/wasi_unstable.rs
+++ b/crates/wasi-common/src/snapshots/wasi_unstable.rs
@@ -23,7 +23,7 @@ impl types::GuestErrorConversion for WasiCtx {
 }
 
 impl types::UserErrorConversion for WasiCtx {
-    fn errno_from_error(&self, e: Error) -> Result<Errno, String> {
+    fn errno_from_error(&self, e: Error) -> Result<Errno, wiggle::Trap> {
         tracing::debug!("Error: {:?}", e);
         Ok(e.into())
     }
@@ -346,7 +346,7 @@ impl wasi_unstable::WasiUnstable for WasiCtx {
         Ok(nevents)
     }
 
-    fn proc_exit(&self, rval: Exitcode) -> Result<(), ()> {
+    fn proc_exit(&self, rval: Exitcode) -> wiggle::Trap {
         WasiSnapshotPreview1::proc_exit(self, rval)
     }
 
diff --git a/crates/wasi-common/src/wasi.rs b/crates/wasi-common/src/wasi.rs
index 214057e766..1877bbf094 100644
--- a/crates/wasi-common/src/wasi.rs
+++ b/crates/wasi-common/src/wasi.rs
@@ -23,7 +23,7 @@ impl types::GuestErrorConversion for WasiCtx {
 }
 
 impl types::UserErrorConversion for WasiCtx {
-    fn errno_from_error(&self, e: Error) -> Result<Errno, String> {
+    fn errno_from_error(&self, e: Error) -> Result<Errno, wiggle::Trap> {
         debug!("Error: {:?}", e);
         Ok(e.into())
     }
diff --git a/crates/wasi/src/lib.rs b/crates/wasi/src/lib.rs
index ba4f7b6c2e..2276c1db2e 100644
--- a/crates/wasi/src/lib.rs
+++ b/crates/wasi/src/lib.rs
@@ -1,5 +1,3 @@
-use wasmtime::Trap;
-
 pub mod old;
 
 pub use wasi_common::virtfs;
@@ -26,12 +24,7 @@ modules. This structure exports all that various fields of the wasi instance
 as fields which can be used to implement your own instantiation logic, if
 necessary. Additionally [`Wasi::get_export`] can be used to do name-based
 resolution.",
-        // Don't use the wiggle generated code to implement proc_exit, we need
-        // to hook directly into the runtime there:
-          function_override: {
-            proc_exit => wasi_proc_exit
-          }
-        },
+        }
     },
 });
 
@@ -41,17 +34,3 @@ pub fn is_wasi_module(name: &str) -> bool {
     // trick.
     name.starts_with("wasi")
 }
-
-/// Implement the WASI `proc_exit` function. This function is implemented here
-/// instead of in wasi-common so that we can use the runtime to perform an
-/// unwind rather than exiting the host process.
-fn wasi_proc_exit(status: i32) -> Result<(), Trap> {
-    // Check that the status is within WASI's range.
-    if status >= 0 && status < 126 {
-        Err(Trap::i32_exit(status))
-    } else {
-        Err(Trap::new(
-            "exit with invalid exit status outside of [0..126)",
-        ))
-    }
-}
diff --git a/crates/wasi/src/old/snapshot_0.rs b/crates/wasi/src/old/snapshot_0.rs
index 826acb995a..f71dee7cdd 100644
--- a/crates/wasi/src/old/snapshot_0.rs
+++ b/crates/wasi/src/old/snapshot_0.rs
@@ -1,8 +1,6 @@
 pub use wasi_common::virtfs;
 pub use wasi_common::{WasiCtx, WasiCtxBuilder};
 
-use crate::wasi_proc_exit;
-
 // Defines a `struct Wasi` with member fields and appropriate APIs for dealing
 // with all the various WASI exports.
 wasmtime_wiggle::wasmtime_integration!({
@@ -25,12 +23,7 @@ modules. This structure exports all that various fields of the wasi instance
 as fields which can be used to implement your own instantiation logic, if
 necessary. Additionally [`Wasi::get_export`] can be used to do name-based
 resolution.",
-        // Don't use the wiggle generated code to implement proc_exit, we need
-        // to hook directly into the runtime there:
-          function_override: {
-            proc_exit => wasi_proc_exit
-          }
-        },
+        }
     },
 });