feat(itr): add wazero WASM runtime for sandboxed code execution
Introduce Layer 5 isolation using wazero (pure Go, no CGO). Each invocation runs in an isolated WASM instance with configurable memory limits, execution timeouts, and capability-gated WASI imports. Includes a SecureBus transport adapter that routes CmdCodeExec requests to the WASM runtime and forwards all others to a fallback.
This commit is contained in:
ZanzyTHEbar
parent
c86f4a53de
commit
dc0559b617
4 changed files with 487 additions and 0 deletions
197
pkg/itr/wasm/runtime.go
Normal file
197
pkg/itr/wasm/runtime.go
Normal file
|
|
@ -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
|
||||
}
|
||||
135
pkg/itr/wasm/runtime_test.go
Normal file
135
pkg/itr/wasm/runtime_test.go
Normal file
|
|
@ -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())
|
||||
}
|
||||
90
pkg/itr/wasm/transport.go
Normal file
90
pkg/itr/wasm/transport.go
Normal file
|
|
@ -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())
|
||||
}
|
||||
65
pkg/itr/wasm/transport_test.go
Normal file
65
pkg/itr/wasm/transport_test.go
Normal file
|
|
@ -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())
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue