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 2e8dea988..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 := newTestInterpreterRuntime(ctx) - 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 1d03fbffe..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 := newTestInterpreterRuntime(ctx) - 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 f5de8abe9..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 { @@ -266,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") @@ -1198,8 +1191,8 @@ func TestWasmGuardClose(t *testing.T) { t.Run("close ignores caller cancellation during cleanup", func(t *testing.T) { ctx := context.Background() - rt := newTestInterpreterRuntime(ctx) - 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} @@ -1322,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 := newTestInterpreterRuntime(ctx) - 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 := newTestInterpreterRuntime(ctx) + 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} @@ -1472,12 +1453,12 @@ func TestWasmMemoryLayout(t *testing.T) { func TestTryCallWasmFunctionDirectMemoryFallback(t *testing.T) { ctx := context.Background() - runtime := newTestInterpreterRuntime(ctx) + 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)) @@ -1605,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 := newTestInterpreterRuntime(ctx) + runtime := newTestWasmRuntime(ctx) t.Cleanup(func() { require.NoError(t, runtime.Close(ctx)) }) @@ -1618,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..5906ac6a7 --- /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. 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)) + 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)) + } +}