diff --git a/pkg/itr/wasm/runtime.go b/pkg/itr/wasm/runtime.go new file mode 100644 index 000000000..4c0b469ec --- /dev/null +++ b/pkg/itr/wasm/runtime.go @@ -0,0 +1,197 @@ +// Package wasm provides sandboxed code execution via wazero (pure Go, no CGO). +// Each invocation runs in an isolated WASM instance with its own linear memory, +// capability-gated WASI imports, and configurable resource limits. +// +// This is the Layer 5 isolation mechanism for untrusted tool code within the +// Isolated Tool Runtime. It integrates with the SecureBus via the CmdCodeExec +// command type. +package wasm + +import ( + "context" + "fmt" + "io" + "strings" + "sync" + "time" + + "github.com/tetratelabs/wazero" + "github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1" +) + +// RuntimeConfig controls resource limits for WASM execution. +type RuntimeConfig struct { + MaxMemoryPages uint32 // max WASM memory pages (64KiB each); 0 → 256 (16MiB) + ExecTimeout time.Duration // per-invocation timeout; 0 → 30s + MaxOutputBytes int // max stdout capture; 0 → 1MiB + EnableWASI bool // mount WASI snapshot_preview1 imports +} + +// DefaultRuntimeConfig returns safe defaults for sandboxed execution. +func DefaultRuntimeConfig() RuntimeConfig { + return RuntimeConfig{ + MaxMemoryPages: 256, + ExecTimeout: 30 * time.Second, + MaxOutputBytes: 1 << 20, + EnableWASI: true, + } +} + +// Runtime manages a pool of wazero runtimes for executing WASM modules. +// Each invocation gets a fresh module instance with private linear memory. +type Runtime struct { + mu sync.Mutex + config RuntimeConfig + engine wazero.Runtime + closed bool +} + +// NewRuntime creates a WASM runtime backed by wazero's compiler. +func NewRuntime(ctx context.Context, cfg RuntimeConfig) (*Runtime, error) { + if cfg.MaxMemoryPages == 0 { + cfg.MaxMemoryPages = 256 + } + if cfg.ExecTimeout == 0 { + cfg.ExecTimeout = 30 * time.Second + } + if cfg.MaxOutputBytes == 0 { + cfg.MaxOutputBytes = 1 << 20 + } + + runtimeCfg := wazero.NewRuntimeConfig(). + WithMemoryLimitPages(cfg.MaxMemoryPages) + + engine := wazero.NewRuntimeWithConfig(ctx, runtimeCfg) + + if cfg.EnableWASI { + if _, err := wasi_snapshot_preview1.Instantiate(ctx, engine); err != nil { + engine.Close(ctx) + return nil, fmt.Errorf("instantiate WASI: %w", err) + } + } + + return &Runtime{ + config: cfg, + engine: engine, + }, nil +} + +// ExecResult holds the output of a WASM execution. +type ExecResult struct { + Stdout string + Stderr string + ExitCode uint32 + Duration time.Duration +} + +// Execute compiles and runs a WASM binary with the given stdin input. +// Each invocation gets a fresh module instance — no state leaks between calls. +func (r *Runtime) Execute(ctx context.Context, wasmBinary []byte, stdin string) (*ExecResult, error) { + r.mu.Lock() + if r.closed { + r.mu.Unlock() + return nil, fmt.Errorf("runtime closed") + } + r.mu.Unlock() + + execCtx, cancel := context.WithTimeout(ctx, r.config.ExecTimeout) + defer cancel() + + compiled, err := r.engine.CompileModule(execCtx, wasmBinary) + if err != nil { + return nil, fmt.Errorf("compile WASM module: %w", err) + } + defer compiled.Close(execCtx) + + var stdout, stderr limitedBuffer + stdout.max = r.config.MaxOutputBytes + stderr.max = r.config.MaxOutputBytes + + moduleCfg := wazero.NewModuleConfig(). + WithStdout(&stdout). + WithStderr(&stderr). + WithStdin(strings.NewReader(stdin)). + WithName("") + + start := time.Now() + mod, err := r.engine.InstantiateModule(execCtx, compiled, moduleCfg) + duration := time.Since(start) + + result := &ExecResult{ + Stdout: stdout.String(), + Stderr: stderr.String(), + Duration: duration, + } + + if err != nil { + if exitErr, ok := extractExitCode(err); ok { + result.ExitCode = exitErr + return result, nil + } + return result, fmt.Errorf("execute WASM: %w", err) + } + + if mod != nil { + _ = mod.Close(execCtx) + } + + return result, nil +} + +// Close releases all wazero resources. +func (r *Runtime) Close(ctx context.Context) error { + r.mu.Lock() + defer r.mu.Unlock() + + if r.closed { + return nil + } + r.closed = true + return r.engine.Close(ctx) +} + +// limitedBuffer captures output up to a maximum size. +type limitedBuffer struct { + buf strings.Builder + max int + n int +} + +func (lb *limitedBuffer) Write(p []byte) (int, error) { + total := len(p) + remaining := lb.max - lb.n + if remaining <= 0 { + return total, nil + } + toWrite := p + if len(toWrite) > remaining { + toWrite = toWrite[:remaining] + } + n, err := lb.buf.Write(toWrite) + lb.n += n + if err != nil { + return n, err + } + return total, nil +} + +func (lb *limitedBuffer) String() string { + return lb.buf.String() +} + +var _ io.Writer = (*limitedBuffer)(nil) + +// extractExitCode checks if err is a sys.ExitError and returns the code. +func extractExitCode(err error) (uint32, bool) { + if err == nil { + return 0, false + } + errStr := err.Error() + if strings.Contains(errStr, "exit_code(") { + var code uint32 + if _, scanErr := fmt.Sscanf(errStr, "module closed with exit_code(%d)", &code); scanErr == nil { + return code, true + } + } + return 0, false +} diff --git a/pkg/itr/wasm/runtime_test.go b/pkg/itr/wasm/runtime_test.go new file mode 100644 index 000000000..2d58c79b5 --- /dev/null +++ b/pkg/itr/wasm/runtime_test.go @@ -0,0 +1,135 @@ +package wasm + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// minimalWASM is a valid WASM module that immediately returns (no-op _start). +// Produced by hand: (module (func (export "_start"))) +var minimalWASM = []byte{ + 0x00, 0x61, 0x73, 0x6d, // magic + 0x01, 0x00, 0x00, 0x00, // version 1 + + // Type section: one function type () -> () + 0x01, // section id + 0x04, // section size + 0x01, // one type + 0x60, // func + 0x00, 0x00, // no params, no results + + // Function section: one function of type 0 + 0x03, // section id + 0x02, // section size + 0x01, // one function + 0x00, // type index 0 + + // Export section: export "_start" as function 0 + 0x07, // section id + 0x0a, // section size + 0x01, // one export + 0x06, // name length + '_', 's', 't', 'a', 'r', 't', // name + 0x00, // export kind: function + 0x00, // function index + + // Code section: one function body + 0x0a, // section id + 0x04, // section size + 0x01, // one body + 0x02, // body size + 0x00, // no locals + 0x0b, // end +} + +func TestNewRuntime(t *testing.T) { + ctx := context.Background() + rt, err := NewRuntime(ctx, DefaultRuntimeConfig()) + require.NoError(t, err) + defer rt.Close(ctx) +} + +func TestExecuteMinimalModule(t *testing.T) { + ctx := context.Background() + rt, err := NewRuntime(ctx, DefaultRuntimeConfig()) + require.NoError(t, err) + defer rt.Close(ctx) + + result, err := rt.Execute(ctx, minimalWASM, "") + require.NoError(t, err) + assert.Equal(t, uint32(0), result.ExitCode) + assert.Empty(t, result.Stderr) + assert.True(t, result.Duration > 0) +} + +func TestExecuteTimeout(t *testing.T) { + ctx := context.Background() + cfg := DefaultRuntimeConfig() + cfg.ExecTimeout = 1 * time.Millisecond + rt, err := NewRuntime(ctx, cfg) + require.NoError(t, err) + defer rt.Close(ctx) + + result, err := rt.Execute(ctx, minimalWASM, "") + // Minimal module is fast enough to succeed even with 1ms timeout. + // This test validates that the timeout machinery doesn't break normal execution. + if err == nil { + assert.Equal(t, uint32(0), result.ExitCode) + } +} + +func TestExecuteInvalidWASM(t *testing.T) { + ctx := context.Background() + rt, err := NewRuntime(ctx, DefaultRuntimeConfig()) + require.NoError(t, err) + defer rt.Close(ctx) + + _, err = rt.Execute(ctx, []byte("not wasm"), "") + assert.Error(t, err) + assert.Contains(t, err.Error(), "compile WASM module") +} + +func TestRuntimeCloseIdempotent(t *testing.T) { + ctx := context.Background() + rt, err := NewRuntime(ctx, DefaultRuntimeConfig()) + require.NoError(t, err) + + assert.NoError(t, rt.Close(ctx)) + assert.NoError(t, rt.Close(ctx)) +} + +func TestExecuteAfterClose(t *testing.T) { + ctx := context.Background() + rt, err := NewRuntime(ctx, DefaultRuntimeConfig()) + require.NoError(t, err) + rt.Close(ctx) + + _, err = rt.Execute(ctx, minimalWASM, "") + assert.Error(t, err) + assert.Contains(t, err.Error(), "runtime closed") +} + +func TestLimitedBuffer(t *testing.T) { + lb := &limitedBuffer{max: 5} + n, err := lb.Write([]byte("hello world")) + assert.NoError(t, err) + assert.Equal(t, 11, n) + assert.Equal(t, "hello", lb.String()) +} + +func TestLimitedBufferExactFit(t *testing.T) { + lb := &limitedBuffer{max: 5} + n, err := lb.Write([]byte("hello")) + assert.NoError(t, err) + assert.Equal(t, 5, n) + assert.Equal(t, "hello", lb.String()) + + n, err = lb.Write([]byte("more")) + assert.NoError(t, err) + assert.Equal(t, 4, n) + assert.Equal(t, "hello", lb.String()) +} diff --git a/pkg/itr/wasm/transport.go b/pkg/itr/wasm/transport.go new file mode 100644 index 000000000..76b8ed443 --- /dev/null +++ b/pkg/itr/wasm/transport.go @@ -0,0 +1,90 @@ +package wasm + +import ( + "context" + "fmt" + + "github.com/sipeed/picoclaw/pkg/itr" +) + +// Transport implements the SecureBus transport interface for WASM-isolated +// tool execution. It handles CmdCodeExec requests by compiling and running +// the provided WASM binary in an isolated instance. +// +// Non-CodeExec requests are forwarded to a fallback transport. +type Transport struct { + runtime *Runtime + fallback func(ctx context.Context, req itr.ToolRequest) (itr.ToolResponse, error) +} + +// NewTransport creates a WasmTransport backed by the given Runtime. +// fallback handles non-CodeExec requests. +func NewTransport(rt *Runtime, fallback func(ctx context.Context, req itr.ToolRequest) (itr.ToolResponse, error)) *Transport { + return &Transport{ + runtime: rt, + fallback: fallback, + } +} + +// Send dispatches a request. CodeExec requests are handled by the WASM runtime; +// all others are forwarded to the fallback. +func (t *Transport) Send(ctx context.Context, req itr.ToolRequest) (itr.ToolResponse, error) { + if req.Type != itr.CmdCodeExec { + if t.fallback != nil { + return t.fallback(ctx, req) + } + return itr.ToolResponse{ + ID: req.ID, + IsError: true, + Result: "WasmTransport: unsupported command type: " + string(req.Type), + }, nil + } + + codeExec, ok := req.Payload.(itr.CodeExec) + if !ok { + return itr.ToolResponse{ + ID: req.ID, + IsError: true, + Result: fmt.Sprintf("expected CodeExec payload, got %T", req.Payload), + }, nil + } + + if len(codeExec.Code) == 0 { + return itr.ToolResponse{ + ID: req.ID, + IsError: true, + Result: "CodeExec: empty code", + }, nil + } + + result, err := t.runtime.Execute(ctx, []byte(codeExec.Code), "") + if err != nil { + return itr.ToolResponse{ + ID: req.ID, + IsError: true, + Result: fmt.Sprintf("WASM execution failed: %v", err), + }, nil + } + + output := result.Stdout + if result.Stderr != "" { + output += "\n[stderr]\n" + result.Stderr + } + if result.ExitCode != 0 { + return itr.ToolResponse{ + ID: req.ID, + IsError: true, + Result: fmt.Sprintf("exit code %d: %s", result.ExitCode, output), + }, nil + } + + return itr.ToolResponse{ + ID: req.ID, + Result: output, + }, nil +} + +// Close releases the WASM runtime resources. +func (t *Transport) Close() error { + return t.runtime.Close(context.Background()) +} diff --git a/pkg/itr/wasm/transport_test.go b/pkg/itr/wasm/transport_test.go new file mode 100644 index 000000000..f06ee9412 --- /dev/null +++ b/pkg/itr/wasm/transport_test.go @@ -0,0 +1,65 @@ +package wasm + +import ( + "context" + "testing" + + "github.com/sipeed/picoclaw/pkg/itr" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestTransportNonCodeExecForwarded(t *testing.T) { + rt, err := NewRuntime(context.Background(), DefaultRuntimeConfig()) + require.NoError(t, err) + defer rt.Close(context.Background()) + + forwarded := false + transport := NewTransport(rt, func(ctx context.Context, req itr.ToolRequest) (itr.ToolResponse, error) { + forwarded = true + return itr.NewSuccessResponse(req.ID, "forwarded", 0), nil + }) + + req := itr.NewToolExecRequest("id-1", "sess", "tc", "read_file", `{"path":"/tmp"}`) + resp, err := transport.Send(context.Background(), req) + require.NoError(t, err) + assert.True(t, forwarded) + assert.Equal(t, "forwarded", resp.Result) + assert.False(t, resp.IsError) +} + +func TestTransportNonCodeExecNoFallback(t *testing.T) { + rt, err := NewRuntime(context.Background(), DefaultRuntimeConfig()) + require.NoError(t, err) + defer rt.Close(context.Background()) + + transport := NewTransport(rt, nil) + + req := itr.NewToolExecRequest("id-1", "sess", "tc", "read_file", `{"path":"/tmp"}`) + resp, err := transport.Send(context.Background(), req) + require.NoError(t, err) + assert.True(t, resp.IsError) + assert.Contains(t, resp.Result, "unsupported command type") +} + +func TestTransportCodeExecEmptyCode(t *testing.T) { + rt, err := NewRuntime(context.Background(), DefaultRuntimeConfig()) + require.NoError(t, err) + defer rt.Close(context.Background()) + + transport := NewTransport(rt, nil) + + req := itr.NewCodeExecRequest("id-1", "sess", "", "javascript") + resp, err := transport.Send(context.Background(), req) + require.NoError(t, err) + assert.True(t, resp.IsError) + assert.Contains(t, resp.Result, "empty code") +} + +func TestTransportClose(t *testing.T) { + rt, err := NewRuntime(context.Background(), DefaultRuntimeConfig()) + require.NoError(t, err) + + transport := NewTransport(rt, nil) + assert.NoError(t, transport.Close()) +}