Merge pull request #4220 from wasmerio/tainting

Ability to detect a tainted instance
wasmerio · Mar 12, 2024 · 7d2226f · 7d2226f
2 parents 4b57317 + 9ddbb2a
commit 7d2226f
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 6 deletions.
diff --git a/lib/wasix/src/bin_factory/exec.rs b/lib/wasix/src/bin_factory/exec.rs
@@ -2,8 +2,11 @@ use std::{pin::Pin, sync::Arc};
 
 use crate::{
     os::task::{thread::WasiThreadRunGuard, TaskJoinHandle},
-    runtime::task_manager::{
-        TaskWasm, TaskWasmRecycle, TaskWasmRecycleProperties, TaskWasmRunProperties,
+    runtime::{
+        task_manager::{
+            TaskWasm, TaskWasmRecycle, TaskWasmRecycleProperties, TaskWasmRunProperties,
+        },
+        TaintReason,
     },
     syscalls::rewind_ext,
     RewindState, SpawnError, WasiError, WasiRuntimeError,
@@ -188,6 +191,7 @@ fn call_module(
     let pid = env.pid();
     let tasks = env.tasks().clone();
     handle.thread.set_status_running();
+    let runtime = env.runtime.clone();
 
     // If we need to rewind then do so
     if let Some((rewind_state, rewind_result)) = rewind_state {
@@ -237,7 +241,10 @@ fn call_module(
         if let Err(err) = call_ret {
             match err.downcast::<WasiError>() {
                 Ok(WasiError::Exit(code)) if code.is_success() => Ok(Errno::Success),
-                Ok(err @ WasiError::Exit(_)) => Err(err.into()),
+                Ok(WasiError::Exit(code)) => {
+                    runtime.on_taint(TaintReason::NonZeroExitCode(code));
+                    Err(WasiError::Exit(code).into())
+                }
                 Ok(WasiError::DeepSleep(deep)) => {
                     // Create the callback that will be invoked when the thread respawns after a deep sleep
                     let rewind = deep.rewind;
@@ -264,9 +271,13 @@ fn call_module(
                 }
                 Ok(WasiError::UnknownWasiVersion) => {
                     debug!("failed as wasi version is unknown",);
+                    runtime.on_taint(TaintReason::UnknownWasiVersion);
                     Ok(Errno::Noexec)
                 }
-                Err(err) => Err(WasiRuntimeError::from(err)),
+                Err(err) => {
+                    runtime.on_taint(TaintReason::RuntimeError(err.clone()));
+                    Err(WasiRuntimeError::from(err))
+                }
             }
         } else {
             Ok(Errno::Success)

diff --git a/lib/wasix/src/runtime/mod.rs b/lib/wasix/src/runtime/mod.rs
@@ -18,7 +18,8 @@ use std::{
 use derivative::Derivative;
 use futures::future::BoxFuture;
 use virtual_net::{DynVirtualNetworking, VirtualNetworking};
-use wasmer::Module;
+use wasmer::{Module, RuntimeError};
+use wasmer_wasix_types::wasi::ExitCode;
 
 #[cfg(feature = "journal")]
 use crate::journal::DynJournal;
@@ -33,6 +34,13 @@ use crate::{
     WasiTtyState,
 };
 
+#[derive(Clone)]
+pub enum TaintReason {
+    UnknownWasiVersion,
+    NonZeroExitCode(ExitCode),
+    RuntimeError(RuntimeError),
+}
+
 /// Runtime components used when running WebAssembly programs.
 ///
 /// Think of this as the "System" in "WebAssembly Systems Interface".
@@ -109,6 +117,10 @@ where
         InlineWaker::block_on(self.load_module(wasm))
     }
 
+    /// Callback thats invokes whenever the instance is tainted, tainting can occur
+    /// for multiple reasons however the most common is a panic within the process
+    fn on_taint(&self, _reason: TaintReason) {}
+
     /// The list of journals which will be used to restore the state of the
     /// runtime at a particular point in time
     #[cfg(feature = "journal")]

diff --git a/lib/wasix/src/syscalls/wasix/thread_spawn.rs b/lib/wasix/src/syscalls/wasix/thread_spawn.rs
@@ -6,7 +6,10 @@ use crate::journal::JournalEffector;
 use crate::{
     capture_instance_snapshot,
     os::task::thread::WasiMemoryLayout,
-    runtime::task_manager::{TaskWasm, TaskWasmRunProperties},
+    runtime::{
+        task_manager::{TaskWasm, TaskWasmRunProperties},
+        TaintReason,
+    },
     syscalls::*,
     WasiThreadHandle,
 };
@@ -177,6 +180,9 @@ fn call_module<M: MemorySize>(
                     ret = if code.is_success() {
                         Errno::Success
                     } else {
+                        env.data(&store)
+                            .runtime
+                            .on_taint(TaintReason::NonZeroExitCode(code));
                         Errno::Noexec
                     };
                 }
@@ -186,10 +192,16 @@ fn call_module<M: MemorySize>(
                 }
                 Ok(WasiError::UnknownWasiVersion) => {
                     debug!("failed as wasi version is unknown",);
+                    env.data(&store)
+                        .runtime
+                        .on_taint(TaintReason::UnknownWasiVersion);
                     ret = Errno::Noexec;
                 }
                 Err(err) => {
                     debug!("failed with runtime error: {}", err);
+                    env.data(&store)
+                        .runtime
+                        .on_taint(TaintReason::RuntimeError(err));
                     ret = Errno::Noexec;
                 }
             }