From 32b6438c1016bee9ad37ce9497f1593613320030 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 19 Aug 2026 13:56:29 +0000 Subject: [PATCH 1/4] Initial plan From d1bda6d927a37566371d16b6802f02d568e40575 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:00:54 +0000 Subject: [PATCH 2/4] Harden and consolidate wazero guard runtime configuration Co-authored-by: lpcox <15877973+lpcox@users.noreply.github.com> --- internal/guard/wasm_config_test.go | 73 +++++++++++++++++++ internal/guard/wasm_dispatch_test.go | 11 +-- internal/guard/wasm_lifecycle.go | 84 +++++++++++++++------- internal/guard/wasm_parse_coverage_test.go | 12 +--- internal/guard/wasm_test.go | 32 +++------ internal/guard/wasm_testruntime_test.go | 50 +++++++++++++ 6 files changed, 195 insertions(+), 67 deletions(-) create mode 100644 internal/guard/wasm_config_test.go create mode 100644 internal/guard/wasm_testruntime_test.go diff --git a/internal/guard/wasm_config_test.go b/internal/guard/wasm_config_test.go new file mode 100644 index 000000000..379b688dd --- /dev/null +++ b/internal/guard/wasm_config_test.go @@ -0,0 +1,73 @@ +package guard + +import ( + "bytes" + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/tetratelabs/wazero" +) + +// TestNewGuardModuleConfig verifies the shared module config builder applies the +// stdin/stdout isolation guarantees and the default module name. +func TestNewGuardModuleConfig(t *testing.T) { + ctx := context.Background() + + t.Run("uses provided name", func(t *testing.T) { + rt := newTestWasmRuntime(ctx) + defer func() { require.NoError(t, rt.Close(ctx)) }() + + mod, err := rt.InstantiateWithConfig(ctx, minimalGuardWasm, newGuardModuleConfig("named-guard", &bytes.Buffer{}, &bytes.Buffer{})) + require.NoError(t, err) + defer func() { require.NoError(t, mod.Close(ctx)) }() + + assert.Equal(t, "named-guard", mod.Name()) + }) + + t.Run("empty name falls back to guard", func(t *testing.T) { + rt := newTestWasmRuntime(ctx) + defer func() { require.NoError(t, rt.Close(ctx)) }() + + mod, err := rt.InstantiateWithConfig(ctx, minimalGuardWasm, newGuardModuleConfig("", &bytes.Buffer{}, &bytes.Buffer{})) + require.NoError(t, err) + defer func() { require.NoError(t, mod.Close(ctx)) }() + + assert.Equal(t, "guard", mod.Name()) + }) +} + +// TestNewGuardRuntimeConfig verifies the shared runtime config builder produces a +// usable runtime for each compilation cache selection path. +func TestNewGuardRuntimeConfig(t *testing.T) { + ctx := context.Background() + + cacheDir := t.TempDir() + customCache, err := wazero.NewCompilationCacheWithDir(cacheDir) + require.NoError(t, err) + defer func() { require.NoError(t, customCache.Close(ctx)) }() + + testCases := []struct { + name string + opts *WasmGuardOptions + }{ + {name: "nil options uses global cache", opts: nil}, + {name: "custom cache", opts: &WasmGuardOptions{CompilationCache: customCache}}, + {name: "cache disabled", opts: &WasmGuardOptions{DisableCompilationCache: true}}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + config := newGuardRuntimeConfig(tc.opts) + require.NotNil(t, config) + + rt := wazero.NewRuntimeWithConfig(ctx, config) + defer func() { require.NoError(t, rt.Close(ctx)) }() + + mod, err := rt.InstantiateWithConfig(ctx, minimalGuardWasm, newGuardModuleConfig("config-guard", &bytes.Buffer{}, &bytes.Buffer{})) + require.NoError(t, err) + require.NoError(t, mod.Close(ctx)) + }) + } +} diff --git a/internal/guard/wasm_dispatch_test.go b/internal/guard/wasm_dispatch_test.go index e48c15bb1..9a6984ba8 100644 --- a/internal/guard/wasm_dispatch_test.go +++ b/internal/guard/wasm_dispatch_test.go @@ -10,7 +10,6 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "github.com/tetratelabs/wazero" ) // labelResourceReturnsZeroWasm exports "label_resource" and "memory"; the function @@ -68,15 +67,7 @@ var labelResponseReturnsTwoWasm = []byte{ // and returns a WasmGuard wired to it plus a cleanup function. func setupRawWasmModule(t *testing.T, wasmBytes []byte, name string) (*WasmGuard, func()) { t.Helper() - ctx := context.Background() - rt := wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfigInterpreter()) - mod, err := rt.InstantiateWithConfig(ctx, wasmBytes, wazero.NewModuleConfig().WithName(name)) - require.NoError(t, err, "failed to instantiate WASM module %s", name) - g := &WasmGuard{name: name, module: mod} - return g, func() { - require.NoError(t, mod.Close(ctx)) - require.NoError(t, rt.Close(ctx)) - } + return setupTestWasmGuard(t, wasmBytes, name) } // --- TestUnmarshalWasmResponse --- diff --git a/internal/guard/wasm_lifecycle.go b/internal/guard/wasm_lifecycle.go index 3aaac70f2..457f8c40f 100644 --- a/internal/guard/wasm_lifecycle.go +++ b/internal/guard/wasm_lifecycle.go @@ -148,6 +148,15 @@ func ConfigureGlobalCompilationCache(ctx context.Context, dir string) error { return nil } +// getGlobalCompilationCache returns the process-level compilation cache under +// the lock so it cannot be read while ConfigureGlobalCompilationCache or +// CloseGlobalCompilationCache is swapping it out. +func getGlobalCompilationCache() wazero.CompilationCache { + globalCompilationCacheMu.Lock() + defer globalCompilationCacheMu.Unlock() + return globalCompilationCache +} + // CloseGlobalCompilationCache releases JIT resources held by the shared // compilation cache. It should be called during graceful shutdown, after all // WasmGuard runtimes have been closed (i.e., after Registry.Close()). @@ -171,6 +180,55 @@ func CloseGlobalCompilationCache(ctx context.Context) error { return nil } +// guardMemoryLimitPages caps guest memory at 512 pages (32 MiB), which +// accommodates the max input/output buffers plus overhead. +const guardMemoryLimitPages = 512 + +// newGuardRuntimeConfig builds the wazero runtime configuration used for guard +// runtimes. It applies the memory cap, context-cancellation cleanup, and +// compilation cache selection (explicit opt-out, injected cache, or shared +// global cache). +// +// DWARF debug info is disabled unless guard debug logging is enabled: guard +// binaries are untrusted third-party artifacts and the gateway never +// symbolicates guest stack traces, so skipping DWARF parsing saves compile +// time and memory without any functional loss. +func newGuardRuntimeConfig(opts *WasmGuardOptions) wazero.RuntimeConfig { + runtimeConfig := wazero.NewRuntimeConfigCompiler(). + WithCloseOnContextDone(true). + WithMemoryLimitPages(guardMemoryLimitPages). + WithDebugInfoEnabled(logWasm.Enabled()) + + switch { + case opts != nil && opts.DisableCompilationCache: + // Caller explicitly disabled caching + case opts != nil && opts.CompilationCache != nil: + runtimeConfig = runtimeConfig.WithCompilationCache(opts.CompilationCache) + default: + runtimeConfig = runtimeConfig.WithCompilationCache(getGlobalCompilationCache()) + } + + return runtimeConfig +} + +// newGuardModuleConfig builds the wazero module configuration used to +// instantiate guard modules. It keeps the stdin/stdout isolation guarantees +// consistent across every instantiation site. +func newGuardModuleConfig(name string, stdout, stderr io.Writer) wazero.ModuleConfig { + guardName := name + if guardName == "" { + guardName = "guard" + } + return wazero.NewModuleConfig(). + WithName(guardName). + // WithStartFunctions with no args suppresses automatic _start execution + // so guard loading cannot block on stdin or perform unexpected I/O. + WithStartFunctions(). + WithStdin(strings.NewReader("")). // Isolate stdin + WithStdout(stdout). // Keep WASM stdout off gateway stdout (MCP stream) + WithStderr(stderr) +} + // WasmGuardOptions configures optional settings for WASM guard creation type WasmGuardOptions struct { // Stdout is the writer for WASM stdout output. Defaults to os.Stderr if nil. @@ -249,18 +307,7 @@ func NewWasmGuardFromBytes(ctx context.Context, name string, wasmBytes []byte, b func NewWasmGuardWithOptions(ctx context.Context, name string, wasmBytes []byte, backend BackendCaller, opts *WasmGuardOptions) (*WasmGuard, error) { logWasm.Printf("Creating WASM guard from bytes: name=%s, size=%d", name, len(wasmBytes)) - // Select compilation cache: explicit opt-out, injected cache, or shared global. - runtimeConfig := wazero.NewRuntimeConfigCompiler(). - WithCloseOnContextDone(true). - WithMemoryLimitPages(512) // 32 MiB hard cap; accommodates max input/output buffers + overhead - if opts != nil && opts.DisableCompilationCache { - // Caller explicitly disabled caching - } else if opts != nil && opts.CompilationCache != nil { - runtimeConfig = runtimeConfig.WithCompilationCache(opts.CompilationCache) - } else { - runtimeConfig = runtimeConfig.WithCompilationCache(globalCompilationCache) - } - runtime := wazero.NewRuntimeWithConfig(ctx, runtimeConfig) + runtime := wazero.NewRuntimeWithConfig(ctx, newGuardRuntimeConfig(opts)) // Instantiate WASI if _, err := wasi_snapshot_preview1.Instantiate(ctx, runtime); err != nil { @@ -293,18 +340,7 @@ func NewWasmGuardWithOptions(ctx context.Context, name string, wasmBytes []byte, } // WithStdin prevents WASM from accidentally reading gateway's MCP protocol stdin - guardName := name - if guardName == "" { - guardName = "guard" - } - moduleConfig := wazero.NewModuleConfig(). - WithName(guardName). - // WithStartFunctions with no args suppresses automatic _start execution - // so guard loading cannot block on stdin or perform unexpected I/O. - WithStartFunctions(). - WithStdin(strings.NewReader("")). // Isolate stdin - WithStdout(stdoutWriter). // Keep WASM stdout off gateway stdout (MCP stream) - WithStderr(stderrWriter) + moduleConfig := newGuardModuleConfig(name, stdoutWriter, stderrWriter) // Compile and instantiate the WASM module module, err := runtime.InstantiateWithConfig(ctx, wasmBytes, moduleConfig) diff --git a/internal/guard/wasm_parse_coverage_test.go b/internal/guard/wasm_parse_coverage_test.go index 8e0dae1d3..285f0eeba 100644 --- a/internal/guard/wasm_parse_coverage_test.go +++ b/internal/guard/wasm_parse_coverage_test.go @@ -6,7 +6,6 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "github.com/tetratelabs/wazero" ) // allocGuardWasm is a WASM module that exports alloc, dealloc, label_agent, and memory. @@ -93,16 +92,7 @@ var labelReturnsNeg1Wasm = []byte{ // bound to it. The caller must invoke the returned cleanup function. func setupWasmGuard(t *testing.T, wasmBytes []byte, name string) (*WasmGuard, func()) { t.Helper() - ctx := context.Background() - rt := wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfigInterpreter()) - mod, err := rt.InstantiateWithConfig(ctx, wasmBytes, wazero.NewModuleConfig().WithName(name)) - require.NoError(t, err, "failed to instantiate WASM module %s", name) - g := &WasmGuard{name: name, module: mod} - cleanup := func() { - require.NoError(t, mod.Close(ctx)) - require.NoError(t, rt.Close(ctx)) - } - return g, cleanup + return setupTestWasmGuard(t, wasmBytes, name) } // TestParsePathLabeledResponse_NewPathLabeledDataError covers the branch where diff --git a/internal/guard/wasm_test.go b/internal/guard/wasm_test.go index cc903de94..f58b8aa7d 100644 --- a/internal/guard/wasm_test.go +++ b/internal/guard/wasm_test.go @@ -259,7 +259,7 @@ func TestWasmGuardContextPropagation(t *testing.T) { }() // Instantiate the blocking WASM module. - moduleConfig := wazero.NewModuleConfig().WithName("blocking_guard") + moduleConfig := newTestWasmModuleConfig("blocking_guard") mod, err := runtime.InstantiateWithConfig(ctx, blockingGuardWasm, moduleConfig) require.NoError(t, err, "failed to instantiate blocking WASM module") @@ -1191,8 +1191,8 @@ func TestWasmGuardClose(t *testing.T) { t.Run("close ignores caller cancellation during cleanup", func(t *testing.T) { ctx := context.Background() - rt := wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfigInterpreter()) - mod, err := rt.InstantiateWithConfig(ctx, minimalGuardWasm, wazero.NewModuleConfig().WithName("close-guard")) + rt := newTestWasmRuntime(ctx) + mod, err := rt.InstantiateWithConfig(ctx, minimalGuardWasm, newTestWasmModuleConfig("close-guard")) require.NoError(t, err) guard := &WasmGuard{runtime: rt, module: mod} @@ -1315,26 +1315,14 @@ func TestHostCallBackendCallLimit(t *testing.T) { func TestBufferRetryLogic(t *testing.T) { // helper instantiates a module for the retry-logic tests. - setupModule := func(t *testing.T, wasmBytes []byte, moduleName string) (*WasmGuard, func()) { - t.Helper() - ctx := context.Background() - rt := wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfigInterpreter()) - mod, err := rt.InstantiateWithConfig(ctx, wasmBytes, wazero.NewModuleConfig().WithName(moduleName)) - require.NoError(t, err) - g := &WasmGuard{name: moduleName, module: mod} - cleanup := func() { - require.NoError(t, mod.Close(ctx)) - require.NoError(t, rt.Close(ctx)) - } - return g, cleanup - } + setupModule := setupTestWasmGuard t.Run("function not exported from module", func(t *testing.T) { // minimalGuardWasm has no exports at all; ExportedFunction returns nil. ctx := context.Background() - rt := wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfigInterpreter()) + rt := newTestWasmRuntime(ctx) t.Cleanup(func() { require.NoError(t, rt.Close(ctx)) }) - mod, err := rt.InstantiateWithConfig(ctx, minimalGuardWasm, wazero.NewModuleConfig().WithName("minimal-retry")) + mod, err := rt.InstantiateWithConfig(ctx, minimalGuardWasm, newTestWasmModuleConfig("minimal-retry")) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, mod.Close(ctx)) }) g := &WasmGuard{name: "minimal-retry", module: mod} @@ -1465,12 +1453,12 @@ func TestWasmMemoryLayout(t *testing.T) { func TestTryCallWasmFunctionDirectMemoryFallback(t *testing.T) { ctx := context.Background() - runtime := wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfigInterpreter()) + runtime := newTestWasmRuntime(ctx) t.Cleanup(func() { require.NoError(t, runtime.Close(ctx)) }) - module, err := runtime.InstantiateWithConfig(ctx, directMemoryFallbackGuardWasm, wazero.NewModuleConfig().WithName("direct-memory-fallback")) + module, err := runtime.InstantiateWithConfig(ctx, directMemoryFallbackGuardWasm, newTestWasmModuleConfig("direct-memory-fallback")) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, module.Close(ctx)) @@ -1598,7 +1586,7 @@ func TestJSONMarshaling(t *testing.T) { func TestIsWasmTrap(t *testing.T) { t.Run("actual wazero trap still uses wasm error prefix (verified with wazero v1.12.0)", func(t *testing.T) { ctx := context.Background() - runtime := wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfigInterpreter()) + runtime := newTestWasmRuntime(ctx) t.Cleanup(func() { require.NoError(t, runtime.Close(ctx)) }) @@ -1611,7 +1599,7 @@ func TestIsWasmTrap(t *testing.T) { 0x0a, 0x05, 0x01, 0x03, 0x00, 0x00, 0x0b, } - mod, err := runtime.InstantiateWithConfig(ctx, trapWasm, wazero.NewModuleConfig().WithName("trap-check")) + mod, err := runtime.InstantiateWithConfig(ctx, trapWasm, newTestWasmModuleConfig("trap-check")) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, mod.Close(ctx)) diff --git a/internal/guard/wasm_testruntime_test.go b/internal/guard/wasm_testruntime_test.go new file mode 100644 index 000000000..d40257407 --- /dev/null +++ b/internal/guard/wasm_testruntime_test.go @@ -0,0 +1,50 @@ +package guard + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + "github.com/tetratelabs/wazero" + "github.com/tetratelabs/wazero/api" +) + +// newTestWasmRuntime creates an interpreter-backed wazero runtime for tests. +// The interpreter avoids JIT compilation, which keeps the many short-lived +// runtimes used by unit tests fast to create and tear down. +// +// Centralizing runtime construction here keeps test configuration consistent +// and makes future wazero config changes a single-line edit. +func newTestWasmRuntime(ctx context.Context) wazero.Runtime { + return wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfigInterpreter()) +} + +// newTestWasmModuleConfig returns the module configuration used when tests +// instantiate raw WASM binaries directly. +func newTestWasmModuleConfig(name string) wazero.ModuleConfig { + return wazero.NewModuleConfig().WithName(name) +} + +// instantiateTestWasmModule instantiates wasmBytes in rt and registers module +// cleanup with t. +func instantiateTestWasmModule(t *testing.T, ctx context.Context, rt wazero.Runtime, wasmBytes []byte, name string) api.Module { + t.Helper() + mod, err := rt.InstantiateWithConfig(ctx, wasmBytes, newTestWasmModuleConfig(name)) + require.NoError(t, err, "failed to instantiate WASM module %s", name) + return mod +} + +// setupTestWasmGuard instantiates a WASM module directly (bypassing +// NewWasmGuardWithOptions) and returns a WasmGuard wired to it plus a cleanup +// function that closes the module and runtime. +func setupTestWasmGuard(t *testing.T, wasmBytes []byte, name string) (*WasmGuard, func()) { + t.Helper() + ctx := context.Background() + rt := newTestWasmRuntime(ctx) + mod := instantiateTestWasmModule(t, ctx, rt, wasmBytes, name) + g := &WasmGuard{name: name, module: mod} + return g, func() { + require.NoError(t, mod.Close(ctx)) + require.NoError(t, rt.Close(ctx)) + } +} From d001f7d0b734cf293eea380906a9f9fa5eebc873 Mon Sep 17 00:00:00 2001 From: Landon Cox Date: Wed, 19 Aug 2026 07:48:29 -0700 Subject: [PATCH 3/4] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- internal/guard/wasm_testruntime_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/internal/guard/wasm_testruntime_test.go b/internal/guard/wasm_testruntime_test.go index d40257407..5906ac6a7 100644 --- a/internal/guard/wasm_testruntime_test.go +++ b/internal/guard/wasm_testruntime_test.go @@ -25,8 +25,8 @@ func newTestWasmModuleConfig(name string) wazero.ModuleConfig { return wazero.NewModuleConfig().WithName(name) } -// instantiateTestWasmModule instantiates wasmBytes in rt and registers module -// cleanup with t. +// instantiateTestWasmModule instantiates wasmBytes in rt. The caller is +// responsible for closing the returned module. func instantiateTestWasmModule(t *testing.T, ctx context.Context, rt wazero.Runtime, wasmBytes []byte, name string) api.Module { t.Helper() mod, err := rt.InstantiateWithConfig(ctx, wasmBytes, newTestWasmModuleConfig(name)) From 9e9449f0b8736a0b26c12b910843d294101533d9 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:08:20 +0000 Subject: [PATCH 4/4] Remove unused newTestInterpreterRuntime helper to fix lint failure Co-authored-by: lpcox <15877973+lpcox@users.noreply.github.com> --- internal/guard/wasm_test.go | 7 ------- 1 file changed, 7 deletions(-) diff --git a/internal/guard/wasm_test.go b/internal/guard/wasm_test.go index 2c67e05ed..f58b8aa7d 100644 --- a/internal/guard/wasm_test.go +++ b/internal/guard/wasm_test.go @@ -19,13 +19,6 @@ import ( "github.com/tetratelabs/wazero/sys" ) -// newTestInterpreterRuntime creates a wazero interpreter runtime for tests. -// The interpreter engine is preferred over the compiler engine in tests -// because it has faster startup and doesn't require JIT support. -func newTestInterpreterRuntime(ctx context.Context) wazero.Runtime { - return wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfigInterpreter()) -} - func TestMain(m *testing.M) { code := m.Run() if err := globalCompilationCache.Close(context.Background()); err != nil {