From 01336d9a4b75d14dc512e89d423171494f09a565 Mon Sep 17 00:00:00 2001 From: rrai35 Date: Thu, 7 Dec 2023 10:29:03 -0500 Subject: [PATCH] feat(closeOnContextDone): Add the ability to close the wasm execution based on context --- engine.go | 11 +++++++++++ engines/wazero/wazero.go | 14 +++++++++++--- 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/engine.go b/engine.go index 94e7161..cfbd4ab 100644 --- a/engine.go +++ b/engine.go @@ -35,6 +35,17 @@ type ( Stdout io.Writer // Stderr is the writer WASI uses for `fd_write` to file descriptor 2. Stderr io.Writer + // CloseOnContextDone ensures the executions of functions to be closed under one of the following circumstances: + // + // - context.Context passed to the Call method of api.Function is canceled during execution. (i.e. ctx by context.WithCancel) + // - context.Context passed to the Call method of api.Function reaches timeout during execution. (i.e. ctx by context.WithTimeout or context.WithDeadline) + // - Close or CloseWithExitCode of api.Module is explicitly called during execution. + // + // This is especially useful when one wants to run untrusted Wasm binaries since otherwise, any invocation of + // api.Function can potentially block the corresponding Goroutine forever. Moreover, it might block the + // entire underlying OS thread which runs the api.Function call. See "Why it's safe to execute runtime-generated + // machine codes against async Goroutine preemption" section in internal/engine/compiler/RATIONALE.md for detail. + CloseOnContextDone bool } // Module is a WebAssembly Module. diff --git a/engines/wazero/wazero.go b/engines/wazero/wazero.go index 9f6b46d..2cae14f 100644 --- a/engines/wazero/wazero.go +++ b/engines/wazero/wazero.go @@ -117,9 +117,17 @@ func DefaultRuntime(ctx context.Context) (wazero.Runtime, error) { // New implements the same method as documented on wapc.Engine. func (e *engine) New(ctx context.Context, host wapc.HostCallHandler, guest []byte, config *wapc.ModuleConfig) (mod wapc.Module, err error) { - r, err := e.newRuntime(ctx) - if err != nil { - return nil, err + //runtime conf changes + var r wazero.Runtime + if config.CloseOnContextDone { + runtimeConfig := wazero.NewRuntimeConfig().WithCloseOnContextDone(true) + r = wazero.NewRuntimeWithConfig(ctx, runtimeConfig) + wasi_snapshot_preview1.MustInstantiate(ctx, r) + } else { + r, err = e.newRuntime(ctx) + if err != nil { + return nil, err + } } m := &Module{runtime: r, wapcHostCallHandler: host}