Pass exceptions caught at the wasm-to-host transition as inner except…

…ion of the TrapException when the trap bubbles up to the next host-to-wasm transition.

Fixes bytecodealliance#63

src/Function.cs

@@ -2450,12 +2450,27 @@ internal unsafe static IntPtr InvokeCallback(Delegate callback, MethodInfo callb
             }
             catch (Exception ex)
             {
-                var bytes = Encoding.UTF8.GetBytes(ex.Message);
+                return HandleCallbackException(ex);
+            }
+        }
+        internal static unsafe IntPtr HandleCallbackException(Exception ex)
+        {
+            // Store the exception as trap cause, so that we can use it as the TrapException's
+            // InnerException when the trap bubbles up to the next host-to-wasm transition.
+            // If the exception is already a TrapException, we use that one's InnerException,
+            // even if it's null.
+            // Note: This code currently assumes that on every host-to-wasm transition where a
+            // trap can occur, TrapException.FromOwnedTrap() is called when a trap actually occured,
+            // which will then clear this field. If this were not the case, we would need to always
+            // set this field no null when returning at the wasm-to-host transition and no exception
+            // was thrown.
+            CallbackTrapCause = ex is TrapException trapException ? trapException.InnerException : ex;
+
+            var bytes = Encoding.UTF8.GetBytes(ex.Message);
 
-                fixed (byte* ptr = bytes)
-                {
-                    return Native.wasmtime_trap_new(ptr, (UIntPtr)bytes.Length);
-                }
+            fixed (byte* ptr = bytes)
+            {
+                return Native.wasmtime_trap_new(ptr, (UIntPtr)bytes.Length);
             }
         }
 
@@ -2512,6 +2527,19 @@ internal static class Native
         internal readonly List<ValueKind> results = new List<ValueKind>();
         internal static readonly Native.Finalizer Finalizer = (p) => GCHandle.FromIntPtr(p).Free();
 
+        /// <summary>
+        /// Contains the cause for a trap returned by invoking a wasm function, in case
+        /// the trap was caused by the host. 
+        /// </summary>
+        /// <remarks>
+        /// This thread-local field will be set when catching a .NET exception at the
+        /// wasm-to-host transition. When the trap bubbles up to the next host-to-wasm
+        /// transition, the field needs to be cleared, and its value can be used to set
+        /// the inner exception of the created <see cref="TrapException"/>.
+        /// </remarks>
+        [ThreadStatic]
+        internal static Exception? CallbackTrapCause;
+
         private static readonly Function _null = new Function();
         private static readonly object?[] NullParams = new object?[1];
     }

src/TrapException.cs

@@ -259,7 +259,7 @@ public TrapException() { }
         public TrapException(string message) : base(message) { }
 
         /// <inheritdoc/>
-        public TrapException(string message, Exception inner) : base(message, inner) { }
+        public TrapException(string message, Exception? inner) : base(message, inner) { }
 
         /// <summary>
         /// Gets the trap's frames.
@@ -281,18 +281,29 @@ public TrapException(string message, Exception inner) : base(message, inner) { }
         /// <inheritdoc/>
         protected TrapException(SerializationInfo info, StreamingContext context) : base(info, context) { }
 
-        internal TrapException(string message, IReadOnlyList<TrapFrame>? frames, TrapCode type) : base(message)
+        internal TrapException(string message, IReadOnlyList<TrapFrame>? frames, TrapCode type, Exception? innerException = null) 
+            : base(message, innerException)
         {
             Type = type;
             Frames = frames;
         }
 
         internal static TrapException FromOwnedTrap(IntPtr trap, bool delete = true)
         {
+            // Get the cause of the trap if available (in case the trap was caused by a
+            // .NET exception thrown in a callback).
+            var callbackTrapCause = Function.CallbackTrapCause;
+
+            if (callbackTrapCause is not null)
+            {
+                // Clear the field as we consumed the value.
+                Function.CallbackTrapCause = null;
+            }
+
             var accessor = new TrapAccessor(trap);
             try
             {
-                var trappedException = new TrapException(accessor.Message, accessor.GetFrames(), accessor.TrapCode)
+                var trappedException = new TrapException(accessor.Message, accessor.GetFrames(), accessor.TrapCode, callbackTrapCause)
                 {
                     ExitCode = accessor.ExitStatus
                 };

src/WasmtimeException.cs

@@ -18,7 +18,7 @@ public WasmtimeException() { }
         public WasmtimeException(string message) : base(message) { }
 
         /// <inheritdoc/>
-        public WasmtimeException(string message, Exception inner) : base(message, inner) { }
+        public WasmtimeException(string message, Exception? inner) : base(message, inner) { }
 
         /// <inheritdoc/>
         protected WasmtimeException(SerializationInfo info, StreamingContext context) : base(info, context) { }

tests/Modules/Trap.wat

@@ -1,9 +1,14 @@
 (module
+  (import "" "trap_from_host_exception" (func $trap_from_host_exception))
+  (import "" "call_host_callback" (func $call_host_callback))
   (export "ok" (func $ok))
   (export "ok_value" (func $ok_value))
   (export "run" (func $run))
   (export "run_div_zero" (func $run_div_zero))
   (export "run_div_zero_with_result" (func $run_div_zero_with_result))
+  (export "trap_from_host_exception" (func $trap_from_host_exception))
+  (export "call_host_callback" (func $call_host_callback))
+  (export "trap_in_wasm" (func $third))
 
   (func $run
     (call $first)

tests/TrapTests.cs

@@ -1,4 +1,6 @@
 using System;
+using System.IO;
+
 using FluentAssertions;
 using Xunit;
 
@@ -47,11 +49,23 @@ public class TrapTests : IClassFixture<TrapFixture>, IDisposable
 
         private Linker Linker { get; set; }
 
+        private Action TrapFromHostExceptionCallback { get; set; }
+
+        private Action HostCallback { get; set; }
+
         public TrapTests(TrapFixture fixture)
         {
             Fixture = fixture;
             Store = new Store(Fixture.Engine);
             Linker = new Linker(Fixture.Engine);
+
+            Linker.Define("", "trap_from_host_exception", Function.FromCallback(
+                Store,
+                () => TrapFromHostExceptionCallback?.Invoke()));
+
+            Linker.Define("", "call_host_callback", Function.FromCallback(
+                Store,
+                () => HostCallback?.Invoke()));
         }
 
         [Fact]
@@ -190,6 +204,64 @@ public void ItHandlesCustomResultTypeWithTrapResult()
             result.TrapStackDepth.Should().Be(1);
         }
 
+        [Fact]
+        public void ItPassesCallbackTrapCauseAsInnerException()
+        {
+            var instance = Linker.Instantiate(Store, Fixture.Module);
+            var callTrap = instance.GetAction("trap_from_host_exception");
+            var trapInWasm = instance.GetAction("trap_in_wasm");
+
+            var exceptionToThrow = new IOException("My I/O exception.");
+
+            TrapFromHostExceptionCallback = () => throw exceptionToThrow;
+
+            // Verify that the IOException thrown at the host callback is passed as
+            // InnerException to the TrapException thrown on the host-to-wasm transition.
+            var action = callTrap;
+
+            action
+               .Should()
+               .Throw<TrapException>()
+               .Where(e => e.InnerException == exceptionToThrow);
+
+            // After that, ensure that when invoking another function that traps in wasm
+            // (so it cannot have a cause), the TrapException's InnerException is now null.
+            action = trapInWasm;
+            action
+                .Should()
+                .Throw<TrapException>()
+                .Where(e => e.InnerException == null);
+
+            // Also verify the InnerException is set when using an ActionResult.
+            var callTrapAsActionResult = instance.GetFunction<ActionResult>("trap_from_host_exception");
+            var result = callTrapAsActionResult();
+
+            result.Type.Should().Be(ResultType.Trap);
+            result.Trap.InnerException.Should().Be(exceptionToThrow);
+        }
+
+        [Fact]
+        public void ItPassesCallbackTrapCauseAsInnerExceptionOverTwoLevels()
+        {
+            var instance = Linker.Instantiate(Store, Fixture.Module);
+            var callTrap = instance.GetAction("trap_from_host_exception");
+            var callHostCallback = instance.GetAction("call_host_callback");
+
+            var exceptionToThrow = new IOException("My I/O exception.");
+
+            TrapFromHostExceptionCallback = () => throw exceptionToThrow;
+            HostCallback = callTrap;
+
+            // Verify that the IOException is passed as InnerException to the
+            // TrapException even after two levels of wasm-to-host transitions.
+            var action = callHostCallback;
+
+            action
+               .Should()
+               .Throw<TrapException>()
+               .Where(e => ReferenceEquals(e.InnerException, exceptionToThrow));
+        }
+
         public void Dispose()
         {
             Store.Dispose();