Skip to content

Commit 381490a

Browse files
heusalagroupbotaibuddy
andauthored
sandbox-wasmrun: add wasm runner validation stub, WASI-deny check, memory bounds helper (#54)
Co-authored-by: aibuddy <aibuddy@dev.hg.fi>
1 parent ea70821 commit 381490a

File tree

5 files changed

+586
-0
lines changed

5 files changed

+586
-0
lines changed

internal/tools/wasmrun/handler.go

Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
package wasmrun
2+
3+
import (
4+
"bytes"
5+
"encoding/base64"
6+
"encoding/json"
7+
"errors"
8+
"os"
9+
"path/filepath"
10+
"time"
11+
)
12+
13+
// Input models the expected stdin JSON for code.sandbox.wasm.run
14+
type Input struct {
15+
ModuleB64 string `json:"module_b64"`
16+
Entry string `json:"entry"`
17+
Input string `json:"input"`
18+
Limits struct {
19+
WallMS int `json:"wall_ms"`
20+
MemPages int `json:"mem_pages"`
21+
OutputKB int `json:"output_kb"`
22+
} `json:"limits"`
23+
}
24+
25+
// Output is the successful stdout JSON shape
26+
type Output struct {
27+
Output string `json:"output"`
28+
}
29+
30+
// Error represents a structured error payload for stderr JSON
31+
type Error struct {
32+
Code string `json:"code"`
33+
Message string `json:"message"`
34+
}
35+
36+
var (
37+
errInvalidInput = errors.New("INVALID_INPUT")
38+
)
39+
40+
// Run parses input and (in future) executes the provided WebAssembly module.
41+
// Returns (stdoutJSON, stderrJSON, err). For now, only input validation is implemented.
42+
func Run(raw []byte) ([]byte, []byte, error) {
43+
start := time.Now()
44+
var in Input
45+
if err := json.Unmarshal(raw, &in); err != nil {
46+
// audit invalid input
47+
_ = appendAudit(map[string]any{ //nolint:errcheck // best-effort audit
48+
"ts": time.Now().UTC().Format(time.RFC3339Nano),
49+
"tool": "code.sandbox.wasm.run",
50+
"span": "tools.wasm.run",
51+
"ms": time.Since(start).Milliseconds(),
52+
"module_bytes": 0,
53+
"wall_ms": 0,
54+
"mem_pages_used": 0,
55+
"bytes_out": 0,
56+
"event": "INVALID_INPUT",
57+
})
58+
return nil, mustMarshalError("INVALID_INPUT", "invalid JSON: "+err.Error()), errInvalidInput
59+
}
60+
if in.ModuleB64 == "" {
61+
_ = appendAudit(map[string]any{ //nolint:errcheck // best-effort audit
62+
"ts": time.Now().UTC().Format(time.RFC3339Nano),
63+
"tool": "code.sandbox.wasm.run",
64+
"span": "tools.wasm.run",
65+
"ms": time.Since(start).Milliseconds(),
66+
"module_bytes": 0,
67+
"wall_ms": in.Limits.WallMS,
68+
"mem_pages_used": 0,
69+
"bytes_out": 0,
70+
"event": "INVALID_INPUT",
71+
})
72+
return nil, mustMarshalError("INVALID_INPUT", "missing module_b64"), errInvalidInput
73+
}
74+
// Validate base64 early to surface errors deterministically
75+
modBytes, err := base64.StdEncoding.DecodeString(in.ModuleB64)
76+
if err != nil {
77+
_ = appendAudit(map[string]any{ //nolint:errcheck // best-effort audit
78+
"ts": time.Now().UTC().Format(time.RFC3339Nano),
79+
"tool": "code.sandbox.wasm.run",
80+
"span": "tools.wasm.run",
81+
"ms": time.Since(start).Milliseconds(),
82+
"module_bytes": 0,
83+
"wall_ms": in.Limits.WallMS,
84+
"mem_pages_used": 0,
85+
"bytes_out": 0,
86+
"event": "INVALID_INPUT",
87+
})
88+
return nil, mustMarshalError("INVALID_INPUT", "module_b64 is not valid base64: "+err.Error()), errInvalidInput
89+
}
90+
// Validate limits
91+
if in.Limits.OutputKB <= 0 {
92+
_ = appendAudit(map[string]any{ //nolint:errcheck // best-effort audit
93+
"ts": time.Now().UTC().Format(time.RFC3339Nano),
94+
"tool": "code.sandbox.wasm.run",
95+
"span": "tools.wasm.run",
96+
"ms": time.Since(start).Milliseconds(),
97+
"module_bytes": len(modBytes),
98+
"wall_ms": in.Limits.WallMS,
99+
"mem_pages_used": 0,
100+
"bytes_out": 0,
101+
"event": "INVALID_INPUT",
102+
})
103+
return nil, mustMarshalError("INVALID_INPUT", "limits.output_kb must be > 0"), errInvalidInput
104+
}
105+
if in.Limits.WallMS <= 0 {
106+
_ = appendAudit(map[string]any{ //nolint:errcheck // best-effort audit
107+
"ts": time.Now().UTC().Format(time.RFC3339Nano),
108+
"tool": "code.sandbox.wasm.run",
109+
"span": "tools.wasm.run",
110+
"ms": time.Since(start).Milliseconds(),
111+
"module_bytes": len(modBytes),
112+
"wall_ms": in.Limits.WallMS,
113+
"mem_pages_used": 0,
114+
"bytes_out": 0,
115+
"event": "INVALID_INPUT",
116+
})
117+
return nil, mustMarshalError("INVALID_INPUT", "limits.wall_ms must be > 0"), errInvalidInput
118+
}
119+
if in.Limits.MemPages <= 0 {
120+
_ = appendAudit(map[string]any{ //nolint:errcheck // best-effort audit
121+
"ts": time.Now().UTC().Format(time.RFC3339Nano),
122+
"tool": "code.sandbox.wasm.run",
123+
"span": "tools.wasm.run",
124+
"ms": time.Since(start).Milliseconds(),
125+
"module_bytes": len(modBytes),
126+
"wall_ms": in.Limits.WallMS,
127+
"mem_pages_used": 0,
128+
"bytes_out": 0,
129+
"event": "INVALID_INPUT",
130+
})
131+
return nil, mustMarshalError("INVALID_INPUT", "limits.mem_pages must be > 0"), errInvalidInput
132+
}
133+
134+
// Deny WASI by default: detect imports of wasi_snapshot_preview1 and fail fast.
135+
// This is a conservative check prior to implementing full wasm execution.
136+
if bytes.Contains(modBytes, []byte("wasi_snapshot_preview1")) {
137+
_ = appendAudit(map[string]any{ //nolint:errcheck // best-effort audit
138+
"ts": time.Now().UTC().Format(time.RFC3339Nano),
139+
"tool": "code.sandbox.wasm.run",
140+
"span": "tools.wasm.run",
141+
"ms": time.Since(start).Milliseconds(),
142+
"module_bytes": len(modBytes),
143+
"wall_ms": in.Limits.WallMS,
144+
"mem_pages_used": 0,
145+
"bytes_out": 0,
146+
"event": "MISSING_IMPORT",
147+
})
148+
return nil, mustMarshalError("MISSING_IMPORT", "WASI is not available by default; modules requiring 'wasi_snapshot_preview1' are unsupported"), errors.New("missing import: wasi_snapshot_preview1")
149+
}
150+
151+
// Not yet implemented: actual wasm execution. Return a stable stub error.
152+
_ = appendAudit(map[string]any{ //nolint:errcheck // best-effort audit
153+
"ts": time.Now().UTC().Format(time.RFC3339Nano),
154+
"tool": "code.sandbox.wasm.run",
155+
"span": "tools.wasm.run",
156+
"ms": time.Since(start).Milliseconds(),
157+
"module_bytes": len(modBytes),
158+
"wall_ms": in.Limits.WallMS,
159+
"mem_pages_used": 0,
160+
"bytes_out": 0,
161+
"event": "UNIMPLEMENTED",
162+
})
163+
return nil, mustMarshalError("UNIMPLEMENTED", "wasm execution not yet implemented"), errors.New("unimplemented")
164+
}
165+
166+
func mustMarshalError(code, msg string) []byte {
167+
b, err := json.Marshal(Error{Code: code, Message: msg})
168+
if err != nil {
169+
return []byte("{\"code\":\"" + code + "\",\"message\":\"" + msg + "\"}")
170+
}
171+
return b
172+
}
173+
174+
// appendAudit writes an NDJSON line under .goagent/audit/YYYYMMDD.log at the repo root.
175+
func appendAudit(entry any) error {
176+
b, err := json.Marshal(entry)
177+
if err != nil {
178+
return err
179+
}
180+
root := moduleRoot()
181+
dir := filepath.Join(root, ".goagent", "audit")
182+
if err := os.MkdirAll(dir, 0o755); err != nil {
183+
return err
184+
}
185+
fname := time.Now().UTC().Format("20060102") + ".log"
186+
path := filepath.Join(dir, fname)
187+
f, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644)
188+
if err != nil {
189+
return err
190+
}
191+
defer func() { _ = f.Close() }() //nolint:errcheck // best-effort close
192+
if _, err := f.Write(append(b, '\n')); err != nil {
193+
return err
194+
}
195+
return nil
196+
}
197+
198+
// moduleRoot walks upward from CWD to the directory containing go.mod; falls back to CWD.
199+
func moduleRoot() string {
200+
cwd, err := os.Getwd()
201+
if err != nil || cwd == "" {
202+
return "."
203+
}
204+
dir := cwd
205+
for {
206+
if _, err := os.Stat(filepath.Join(dir, "go.mod")); err == nil {
207+
return dir
208+
}
209+
parent := filepath.Dir(dir)
210+
if parent == dir {
211+
return cwd
212+
}
213+
dir = parent
214+
}
215+
}
Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
package wasmrun
2+
3+
import (
4+
"encoding/base64"
5+
"encoding/json"
6+
"testing"
7+
)
8+
9+
func TestRun_InvalidJSON(t *testing.T) {
10+
stdout, stderr, err := Run([]byte("not-json"))
11+
if err == nil {
12+
t.Fatalf("expected error for invalid JSON")
13+
}
14+
if len(stdout) != 0 {
15+
t.Fatalf("expected no stdout, got: %s", string(stdout))
16+
}
17+
var e struct{ Code, Message string }
18+
if jerr := json.Unmarshal(stderr, &e); jerr != nil {
19+
t.Fatalf("stderr not JSON: %v: %s", jerr, string(stderr))
20+
}
21+
if e.Code != "INVALID_INPUT" {
22+
t.Fatalf("expected INVALID_INPUT, got %q (%s)", e.Code, e.Message)
23+
}
24+
}
25+
26+
func TestRun_MissingModuleB64(t *testing.T) {
27+
req := map[string]any{
28+
"entry": "main",
29+
"input": "",
30+
"limits": map[string]any{"output_kb": 1, "wall_ms": 10, "mem_pages": 1},
31+
}
32+
b, merr := json.Marshal(req)
33+
if merr != nil {
34+
t.Fatalf("marshal failed: %v", merr)
35+
}
36+
stdout, stderr, err := Run(b)
37+
if err == nil {
38+
t.Fatalf("expected error for missing module_b64")
39+
}
40+
if len(stdout) != 0 {
41+
t.Fatalf("expected no stdout, got: %s", string(stdout))
42+
}
43+
var e struct{ Code, Message string }
44+
if jerr := json.Unmarshal(stderr, &e); jerr != nil {
45+
t.Fatalf("stderr not JSON: %v: %s", jerr, string(stderr))
46+
}
47+
if e.Code != "INVALID_INPUT" {
48+
t.Fatalf("expected INVALID_INPUT, got %q (%s)", e.Code, e.Message)
49+
}
50+
}
51+
52+
func TestRun_BadBase64(t *testing.T) {
53+
req := map[string]any{
54+
"module_b64": "!!!not-base64!!!",
55+
"entry": "main",
56+
"input": "",
57+
"limits": map[string]any{"output_kb": 1, "wall_ms": 10, "mem_pages": 1},
58+
}
59+
b, merr := json.Marshal(req)
60+
if merr != nil {
61+
t.Fatalf("marshal failed: %v", merr)
62+
}
63+
stdout, stderr, err := Run(b)
64+
if err == nil {
65+
t.Fatalf("expected error for invalid base64")
66+
}
67+
if len(stdout) != 0 {
68+
t.Fatalf("expected no stdout, got: %s", string(stdout))
69+
}
70+
var e struct{ Code, Message string }
71+
if jerr := json.Unmarshal(stderr, &e); jerr != nil {
72+
t.Fatalf("stderr not JSON: %v: %s", jerr, string(stderr))
73+
}
74+
if e.Code != "INVALID_INPUT" {
75+
t.Fatalf("expected INVALID_INPUT, got %q (%s)", e.Code, e.Message)
76+
}
77+
}
78+
79+
func TestRun_UnimplementedOnValidInput(t *testing.T) {
80+
// module_b64 is valid base64 but not necessarily a valid wasm; current stub only validates base64
81+
req := map[string]any{
82+
"module_b64": "AA==", // base64 for single zero byte
83+
"entry": "main",
84+
"input": "",
85+
"limits": map[string]any{"output_kb": 1, "wall_ms": 10, "mem_pages": 1},
86+
}
87+
b, merr := json.Marshal(req)
88+
if merr != nil {
89+
t.Fatalf("marshal failed: %v", merr)
90+
}
91+
stdout, stderr, err := Run(b)
92+
if err == nil {
93+
t.Fatalf("expected unimplemented error")
94+
}
95+
if len(stdout) != 0 {
96+
t.Fatalf("expected no stdout, got: %s", string(stdout))
97+
}
98+
var e struct{ Code, Message string }
99+
if jerr := json.Unmarshal(stderr, &e); jerr != nil {
100+
t.Fatalf("stderr not JSON: %v: %s", jerr, string(stderr))
101+
}
102+
if e.Code != "UNIMPLEMENTED" {
103+
t.Fatalf("expected UNIMPLEMENTED, got %q (%s)", e.Code, e.Message)
104+
}
105+
}
106+
107+
func TestRun_InvalidLimits(t *testing.T) {
108+
cases := []map[string]any{
109+
{"module_b64": "AA==", "entry": "main", "input": "", "limits": map[string]any{"output_kb": 0, "wall_ms": 10, "mem_pages": 1}},
110+
{"module_b64": "AA==", "entry": "main", "input": "", "limits": map[string]any{"output_kb": 1, "wall_ms": 0, "mem_pages": 1}},
111+
{"module_b64": "AA==", "entry": "main", "input": "", "limits": map[string]any{"output_kb": 1, "wall_ms": 10, "mem_pages": 0}},
112+
}
113+
for i, req := range cases {
114+
b, merr := json.Marshal(req)
115+
if merr != nil {
116+
t.Fatalf("case %d: marshal failed: %v", i, merr)
117+
}
118+
stdout, stderr, err := Run(b)
119+
if err == nil {
120+
t.Fatalf("case %d: expected error for invalid limits", i)
121+
}
122+
if len(stdout) != 0 {
123+
t.Fatalf("case %d: expected no stdout, got: %s", i, string(stdout))
124+
}
125+
var e struct{ Code, Message string }
126+
if jerr := json.Unmarshal(stderr, &e); jerr != nil {
127+
t.Fatalf("case %d: stderr not JSON: %v: %s", i, jerr, string(stderr))
128+
}
129+
if e.Code != "INVALID_INPUT" {
130+
t.Fatalf("case %d: expected INVALID_INPUT, got %q (%s)", i, e.Code, e.Message)
131+
}
132+
}
133+
}
134+
135+
func TestRun_DenyWASIByDefault(t *testing.T) {
136+
// Any bytes containing the string "wasi_snapshot_preview1" should be denied
137+
// even before actual execution is implemented.
138+
wasmLike := base64.StdEncoding.EncodeToString([]byte("xxwasi_snapshot_preview1xx"))
139+
req := map[string]any{
140+
"module_b64": wasmLike,
141+
"entry": "main",
142+
"input": "",
143+
"limits": map[string]any{"output_kb": 1, "wall_ms": 10, "mem_pages": 1},
144+
}
145+
b, merr := json.Marshal(req)
146+
if merr != nil {
147+
t.Fatalf("marshal failed: %v", merr)
148+
}
149+
stdout, stderr, err := Run(b)
150+
if err == nil {
151+
t.Fatalf("expected error for WASI-dependent module")
152+
}
153+
if len(stdout) != 0 {
154+
t.Fatalf("expected no stdout, got: %s", string(stdout))
155+
}
156+
var e struct{ Code, Message string }
157+
if jerr := json.Unmarshal(stderr, &e); jerr != nil {
158+
t.Fatalf("stderr not JSON: %v: %s", jerr, string(stderr))
159+
}
160+
if e.Code != "MISSING_IMPORT" {
161+
t.Fatalf("expected MISSING_IMPORT, got %q (%s)", e.Code, e.Message)
162+
}
163+
}

0 commit comments

Comments
 (0)