From 77beb0b54303c03c0ec015394f39d39b5b27611e Mon Sep 17 00:00:00 2001 From: Dmitry Vyukov Date: Mon, 26 Jan 2026 16:44:09 +0100 Subject: [PATCH 1/2] pkg/aflow: add helper for tool testing Add simple codeeditor tests to test testing. --- pkg/aflow/flow.go | 7 +- pkg/aflow/func_tool.go | 19 +++++ pkg/aflow/test_tool.go | 49 ++++++++++++ pkg/aflow/tool/codeeditor/codeeditor.go | 15 +++- pkg/aflow/tool/codeeditor/codeeditor_test.go | 79 ++++++++++++++++++++ pkg/aflow/verify.go | 9 +++ 6 files changed, 171 insertions(+), 7 deletions(-) create mode 100644 pkg/aflow/test_tool.go create mode 100644 pkg/aflow/tool/codeeditor/codeeditor_test.go diff --git a/pkg/aflow/flow.go b/pkg/aflow/flow.go index 4073a6de13f4..f2b536bc73f4 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 220cb8f2b139..b67830115788 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 000000000000..4b9cb6a4a5c5 --- /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 f67abbd69635..ce2d7afb7ab7 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 000000000000..4ba556f1bd7b --- /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 446b0ac5e2fb..4d1803fcbc6d 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...) } From cdf6ca353b15f42c9cfe0e7eb00425f0b761be0b Mon Sep 17 00:00:00 2001 From: Dmitry Vyukov Date: Mon, 26 Jan 2026 16:44:12 +0100 Subject: [PATCH 2/2] pkg/aflow/tool/codeeditor: add actual implementation --- pkg/aflow/tool/codeeditor/codeeditor.go | 94 ++++++++-- pkg/aflow/tool/codeeditor/codeeditor_test.go | 176 ++++++++++++++++++- 2 files changed, 252 insertions(+), 18 deletions(-) diff --git a/pkg/aflow/tool/codeeditor/codeeditor.go b/pkg/aflow/tool/codeeditor/codeeditor.go index ce2d7afb7ab7..750dba5d8dd2 100644 --- a/pkg/aflow/tool/codeeditor/codeeditor.go +++ b/pkg/aflow/tool/codeeditor/codeeditor.go @@ -4,15 +4,21 @@ package codeeditor import ( + "bytes" + "os" "path/filepath" + "slices" "strings" "github.com/google/syzkaller/pkg/aflow" + "github.com/google/syzkaller/pkg/codesearch" "github.com/google/syzkaller/pkg/osutil" ) var Tool = aflow.NewFuncTool("codeeditor", codeeditor, ` -The tool does one code edit to form the final patch. +The tool does one source code edit to form the final patch by replacing full lines +with new provided lines. If new code is empty, current lines will be deleted. +Provide full lines of code including new line characters. The tool should be called mutiple times to do all required changes one-by-one, but avoid changing the same lines multiple times. Note: You will not see your edits via the codesearch tool. @@ -24,9 +30,9 @@ type state struct { } type args struct { - SourceFile string `jsonschema:"Full source file path."` - CurrentCode string `jsonschema:"The current code to replace verbatim with new lines, but without line numbers."` - NewCode string `jsonschema:"New code to replace the current code snippet."` + SourceFile string `jsonschema:"Full source file path to edit."` + CurrentCode string `jsonschema:"The current code lines to be replaced."` + NewCode string `jsonschema:"New code lines to replace the current code lines."` } func codeeditor(ctx *aflow.Context, state state, args args) (struct{}, error) { @@ -34,17 +40,81 @@ func codeeditor(ctx *aflow.Context, state state, args args) (struct{}, error) { 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) { + // Filter out not source files too (e.g. .git, etc), + // LLM have not seen them and should not be messing with them. + if !osutil.IsExist(file) || !codesearch.IsSourceFile(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, - // ignore white-spaces, etc). - // Should we accept a reference line number, or function name to disambiguate in the case - // of multiple matches? - return struct{}{}, nil + fileData, err := os.ReadFile(file) + if err != nil { + return struct{}{}, err + } + if len(fileData) == 0 || fileData[len(fileData)-1] != '\n' { + // Generally shouldn't happen, but just in case. + fileData = append(fileData, '\n') + } + if args.CurrentCode[len(args.CurrentCode)-1] != '\n' { + args.CurrentCode += "\n" + } + if args.NewCode != "" && args.NewCode[len(args.NewCode)-1] != '\n' { + args.NewCode += "\n" + } + lines := slices.Collect(bytes.Lines(fileData)) + src := slices.Collect(bytes.Lines([]byte(args.CurrentCode))) + dst := slices.Collect(bytes.Lines([]byte(args.NewCode))) + // First, try to match as is. If that fails, try a more permissive matching + // that ignores whitespaces, empty lines, etc. + newLines, matches := replace(lines, src, dst, false) + if matches == 0 { + newLines, matches = replace(lines, src, dst, true) + } + if matches == 0 { + return struct{}{}, aflow.BadCallError("CurrentCode snippet does not match anything in the source file," + + " provide more precise CurrentCode snippet") + } + if matches > 1 { + return struct{}{}, aflow.BadCallError("CurrentCode snippet matched %v places,"+ + " increase context in CurrentCode to avoid ambiguity", matches) + } + err = osutil.WriteFile(file, slices.Concat(newLines...)) + return struct{}{}, err +} + +func replace(lines, src, dst [][]byte, fuzzy bool) (newLines [][]byte, matches int) { + for i := 0; i < len(lines); i++ { + li, si := i, 0 + for li < len(lines) && si < len(src) { + l, s := lines[li], src[si] + if fuzzy { + // Ignore whitespaces and empty lines. + l, s = bytes.TrimSpace(l), bytes.TrimSpace(s) + // Potentially we can remove line numbers from s here if they are present, + // or use them to disambiguate in the case of multiple matches. + if len(s) == 0 { + si++ + continue + } + if len(l) == 0 { + li++ + continue + } + } + if !bytes.Equal(l, s) { + break + } + li++ + si++ + } + if si != len(src) { + newLines = append(newLines, lines[i]) + continue + } + matches++ + newLines = append(newLines, dst...) + i = li - 1 + } + return } diff --git a/pkg/aflow/tool/codeeditor/codeeditor_test.go b/pkg/aflow/tool/codeeditor/codeeditor_test.go index 4ba556f1bd7b..06a97d7afe0f 100644 --- a/pkg/aflow/tool/codeeditor/codeeditor_test.go +++ b/pkg/aflow/tool/codeeditor/codeeditor_test.go @@ -4,6 +4,8 @@ package codeeditor import ( + "fmt" + "os" "path/filepath" "testing" @@ -15,7 +17,7 @@ import ( func TestCodeeditorEscapingPath(t *testing.T) { aflow.TestTool(t, Tool, state{ - KernelScratchSrc: "whatever", + KernelScratchSrc: t.TempDir(), }, args{ SourceFile: "../../passwd", @@ -38,25 +40,187 @@ func TestCodeeditorMissingPath(t *testing.T) { ) } +func TestCodeeditorNonSourceFile(t *testing.T) { + dir := writeTestFile(t, "src", "data") + aflow.TestTool(t, Tool, + state{ + KernelScratchSrc: dir, + }, + args{ + SourceFile: "src", + }, + struct{}{}, + `SourceFile "src" does not exist`, + ) +} + func TestCodeeditorEmptyCurrentCode(t *testing.T) { - dir := writeTestFile(t, "foo", "data") + dir := writeTestFile(t, "src.c", "data") aflow.TestTool(t, Tool, state{ KernelScratchSrc: dir, }, args{ - SourceFile: "foo", + SourceFile: "src.c", }, struct{}{}, `CurrentCode snippet is empty`, ) } +func TestCodeeditorNoMatches(t *testing.T) { + dir := writeTestFile(t, "src.c", "foo") + aflow.TestTool(t, Tool, + state{ + KernelScratchSrc: dir, + }, + args{ + SourceFile: "src.c", + CurrentCode: "foobar", + }, + struct{}{}, + `CurrentCode snippet does not match anything in the source file, provide more precise CurrentCode snippet`, + ) +} + +func TestCodeeditorMultipleMatches(t *testing.T) { + dir := writeTestFile(t, "src.c", ` +linefoo +bar +foo +bar +foo +fooline +foo`) + aflow.TestTool(t, Tool, + state{ + KernelScratchSrc: dir, + }, + args{ + SourceFile: "src.c", + CurrentCode: "foo", + }, + struct{}{}, + `CurrentCode snippet matched 3 places, increase context in CurrentCode to avoid ambiguity`, + ) +} + +func TestCodeeditorReplacement(t *testing.T) { + type Test struct { + curFile string + curCode string + newCode string + newFile string + } + tests := []Test{ + { + curFile: ` +line0 +line1 +lineee2 +lin3 +last line +`, + curCode: `line1 +lineee2 +lin3`, + newCode: `replaced line`, + newFile: ` +line0 +replaced line +last line +`, + }, + { + curFile: ` +line0 +line1 +last line +`, + curCode: `line1 +`, + newCode: `replaced line 1 +replaced line 2 +replaced line 3`, + newFile: ` +line0 +replaced line 1 +replaced line 2 +replaced line 3 +last line +`, + }, + { + curFile: ` +line0 +line1 +line2 +`, + curCode: `line2 +`, + newCode: ``, + newFile: ` +line0 +line1 +`, + }, + { + curFile: `that's it`, + curCode: `that's it`, + newCode: `that's that`, + newFile: `that's that +`, + }, + { + curFile: ` + line0 + line1 + + line2 +line3 + +line4 +`, + curCode: ` +line1 + line2 + + + line3 `, + newCode: ` replacement`, + newFile: ` + line0 + replacement + +line4 +`, + }, + } + for i, test := range tests { + t.Run(fmt.Sprint(i), func(t *testing.T) { + const filename = "src.c" + dir := writeTestFile(t, filename, test.curFile) + aflow.TestTool(t, Tool, + state{ + KernelScratchSrc: dir, + }, + args{ + SourceFile: filename, + CurrentCode: test.curCode, + NewCode: test.newCode, + }, + struct{}{}, + "") + data, err := os.ReadFile(filepath.Join(dir, filename)) + require.NoError(t, err) + require.Equal(t, test.newFile, string(data)) + }) + } +} + 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) - } + require.NoError(t, osutil.WriteFile(filepath.Join(dir, filename), []byte(data))) return dir }