pkg/aflow: add helper for tool testing

Add simple codeeditor tests to test testing.

6 files changed, 171 insertions, 7 deletions

diff --git a/pkg/aflow/flow.go b/pkg/aflow/flow.go
index 4073a6de1..f2b536bc7 100644
--- a/pkg/aflow/flow.go
+++ b/pkg/aflow/flow.go
@@ -87,12 +87,7 @@ func registerOne[Inputs, Outputs any](all map[string]*Flow, flow *Flow) error {
 	if all[flow.Name] != nil {
 		return fmt.Errorf("flow %v is already registered", flow.Name)
 	}
-	ctx := &verifyContext{
-		inputs:  true,
-		outputs: true,
-		state:   make(map[string]*varState),
-		models:  make(map[string]bool),
-	}
+	ctx := newVerifyContext()
 	provideOutputs[Inputs](ctx, "flow inputs")
 	flow.Root.verify(ctx)
 	requireInputs[Outputs](ctx, "flow outputs")
diff --git a/pkg/aflow/func_tool.go b/pkg/aflow/func_tool.go
index 220cb8f2b..b67830115 100644
--- a/pkg/aflow/func_tool.go
+++ b/pkg/aflow/func_tool.go
@@ -5,7 +5,10 @@ package aflow
 
 import (
 	"fmt"
+	"reflect"
+	"testing"
 
+	"github.com/stretchr/testify/require"
 	"google.golang.org/genai"
 )
 
@@ -77,3 +80,19 @@ func (t *funcTool[State, Args, Results]) verify(ctx *verifyContext) {
 	requireSchema[Results](ctx, t.Name, "Results")
 	requireInputs[State](ctx, t.Name)
 }
+
+func (*funcTool[State, Args, Results]) checkTestTypes(t *testing.T, ctx *verifyContext, state, args, results any) (
+	map[string]any, map[string]any, map[string]any) {
+	require.Equal(t, reflect.TypeFor[State](), reflect.TypeOf(state))
+	require.Equal(t, reflect.TypeFor[Args](), reflect.TypeOf(args))
+	require.Equal(t, reflect.TypeFor[Results](), reflect.TypeOf(results))
+	provideOutputs[State](ctx, "state")
+	return convertToMap(state.(State)), convertToMap(args.(Args)), convertToMap(results.(Results))
+}
+
+func (*funcTool[State, Args, Results]) checkFuzzTypes(t *testing.T, state, args any) (
+	map[string]any, map[string]any) {
+	require.Equal(t, reflect.TypeFor[State](), reflect.TypeOf(state))
+	require.Equal(t, reflect.TypeFor[Args](), reflect.TypeOf(args))
+	return convertToMap(state.(State)), convertToMap(args.(Args))
+}
diff --git a/pkg/aflow/test_tool.go b/pkg/aflow/test_tool.go
new file mode 100644
index 000000000..4b9cb6a4a
--- /dev/null
+++ b/pkg/aflow/test_tool.go
@@ -0,0 +1,49 @@
+// Copyright 2026 syzkaller project authors. All rights reserved.
+// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
+
+package aflow
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestTool(t *testing.T, tool Tool, initState, initArgs, wantResults any, wantError string) {
+	type toolTester interface {
+		checkTestTypes(t *testing.T, ctx *verifyContext, state, args, results any) (
+			map[string]any, map[string]any, map[string]any)
+	}
+	vctx := newVerifyContext()
+	state, args, results := tool.(toolTester).checkTestTypes(t, vctx, initState, initArgs, wantResults)
+	tool.verify(vctx)
+	if err := vctx.finalize(); err != nil {
+		t.Fatal(err)
+	}
+	// Just ensure it does not crash.
+	_ = tool.declaration()
+	// We don't init all fields, init more, if necessary.
+	ctx := &Context{
+		state: state,
+	}
+	defer ctx.close()
+	gotResults, err := tool.execute(ctx, args)
+	gotError := ""
+	if err != nil {
+		gotError = err.Error()
+	}
+	require.Equal(t, wantError, gotError)
+	require.Equal(t, results, gotResults)
+}
+
+func FuzzTool(t *testing.T, tool Tool, initState, initArgs any) (map[string]any, error) {
+	type toolFuzzer interface {
+		checkFuzzTypes(t *testing.T, state, args any) (map[string]any, map[string]any)
+	}
+	state, args := tool.(toolFuzzer).checkFuzzTypes(t, initState, initArgs)
+	ctx := &Context{
+		state: state,
+	}
+	defer ctx.close()
+	return tool.execute(ctx, args)
+}
diff --git a/pkg/aflow/tool/codeeditor/codeeditor.go b/pkg/aflow/tool/codeeditor/codeeditor.go
index f67abbd69..ce2d7afb7 100644
--- a/pkg/aflow/tool/codeeditor/codeeditor.go
+++ b/pkg/aflow/tool/codeeditor/codeeditor.go
@@ -4,7 +4,11 @@
 package codeeditor
 
 import (
+	"path/filepath"
+	"strings"
+
 	"github.com/google/syzkaller/pkg/aflow"
+	"github.com/google/syzkaller/pkg/osutil"
 )
 
 var Tool = aflow.NewFuncTool("codeeditor", codeeditor, `
@@ -26,7 +30,16 @@ type args struct {
 }
 
 func codeeditor(ctx *aflow.Context, state state, args args) (struct{}, error) {
-	// TODO: check that the SourceFile is not escaping.
+	if strings.Contains(filepath.Clean(args.SourceFile), "..") {
+		return struct{}{}, aflow.BadCallError("SourceFile %q is outside of the source tree", args.SourceFile)
+	}
+	file := filepath.Join(state.KernelScratchSrc, args.SourceFile)
+	if !osutil.IsExist(file) {
+		return struct{}{}, aflow.BadCallError("SourceFile %q does not exist", args.SourceFile)
+	}
+	if strings.TrimSpace(args.CurrentCode) == "" {
+		return struct{}{}, aflow.BadCallError("CurrentCode snippet is empty")
+	}
 	// If SourceFile is incorrect, or CurrentCode is not matched, return aflow.BadCallError
 	// with an explanation. Say that it needs to increase context if CurrentCode is not matched.
 	// Try to do as fuzzy match for CurrentCode as possible (strip line numbers,
diff --git a/pkg/aflow/tool/codeeditor/codeeditor_test.go b/pkg/aflow/tool/codeeditor/codeeditor_test.go
new file mode 100644
index 000000000..4ba556f1b
--- /dev/null
+++ b/pkg/aflow/tool/codeeditor/codeeditor_test.go
@@ -0,0 +1,79 @@
+// Copyright 2026 syzkaller project authors. All rights reserved.
+// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
+
+package codeeditor
+
+import (
+	"path/filepath"
+	"testing"
+
+	"github.com/google/syzkaller/pkg/aflow"
+	"github.com/google/syzkaller/pkg/osutil"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCodeeditorEscapingPath(t *testing.T) {
+	aflow.TestTool(t, Tool,
+		state{
+			KernelScratchSrc: "whatever",
+		},
+		args{
+			SourceFile: "../../passwd",
+		},
+		struct{}{},
+		`SourceFile "../../passwd" is outside of the source tree`,
+	)
+}
+
+func TestCodeeditorMissingPath(t *testing.T) {
+	aflow.TestTool(t, Tool,
+		state{
+			KernelScratchSrc: t.TempDir(),
+		},
+		args{
+			SourceFile: "missing-file",
+		},
+		struct{}{},
+		`SourceFile "missing-file" does not exist`,
+	)
+}
+
+func TestCodeeditorEmptyCurrentCode(t *testing.T) {
+	dir := writeTestFile(t, "foo", "data")
+	aflow.TestTool(t, Tool,
+		state{
+			KernelScratchSrc: dir,
+		},
+		args{
+			SourceFile: "foo",
+		},
+		struct{}{},
+		`CurrentCode snippet is empty`,
+	)
+}
+
+func writeTestFile(t *testing.T, filename, data string) string {
+	dir := t.TempDir()
+	if err := osutil.WriteFile(filepath.Join(dir, filename), []byte(data)); err != nil {
+		t.Fatal(err)
+	}
+	return dir
+}
+
+func Fuzz(f *testing.F) {
+	dir := f.TempDir()
+	const filename = "src.c"
+	fullFilename := filepath.Join(dir, filename)
+	f.Fuzz(func(t *testing.T, fileData []byte, curCode, newCode string) {
+		require.NoError(t, osutil.WriteFile(fullFilename, fileData))
+		aflow.FuzzTool(t, Tool,
+			state{
+				KernelScratchSrc: dir,
+			},
+			args{
+				SourceFile:  filename,
+				CurrentCode: curCode,
+				NewCode:     newCode,
+			})
+	})
+}
diff --git a/pkg/aflow/verify.go b/pkg/aflow/verify.go
index 446b0ac5e..4d1803fcb 100644
--- a/pkg/aflow/verify.go
+++ b/pkg/aflow/verify.go
@@ -24,6 +24,15 @@ type varState struct {
 	used   bool
 }
 
+func newVerifyContext() *verifyContext {
+	return &verifyContext{
+		inputs:  true,
+		outputs: true,
+		state:   make(map[string]*varState),
+		models:  make(map[string]bool),
+	}
+}
+
 func (ctx *verifyContext) errorf(who, msg string, args ...any) {
 	noteError(&ctx.err, fmt.Sprintf("action %v: %v", who, msg), args...)
 }