pkg/aflow/action/kernel: keep build files that codesearch will need

We currently duplicate list of source extensions in the build action
and codesearch tool. Unify the lists.

2 files changed, 26 insertions, 8 deletions

diff --git a/pkg/aflow/action/kernel/build.go b/pkg/aflow/action/kernel/build.go
index 594b33fb5..a24f9359b 100644
--- a/pkg/aflow/action/kernel/build.go
+++ b/pkg/aflow/action/kernel/build.go
@@ -14,6 +14,7 @@ import (
 
 	"github.com/google/syzkaller/pkg/aflow"
 	"github.com/google/syzkaller/pkg/build"
+	"github.com/google/syzkaller/pkg/codesearch"
 	"github.com/google/syzkaller/pkg/hash"
 	"github.com/google/syzkaller/pkg/osutil"
 	"github.com/google/syzkaller/sys/targets"
@@ -50,17 +51,23 @@ func buildKernel(ctx *aflow.Context, args buildArgs) (buildResult, error) {
 			return aflow.FlowError(err)
 		}
 		// Remove main intermediate build files, we don't need them anymore
-		// and they take lots of space. Keep generated source files.
-		keepExt := map[string]bool{".h": true, ".c": true, ".s": true, ".S": true}
+		// and they take lots of space. But keep generated source files.
 		keepFiles := map[string]bool{
-			filepath.Join(dir, image):               true,
-			filepath.Join(dir, target.KernelObject): true,
-			filepath.Join(dir, compileCommnads):     true,
+			image:               true,
+			target.KernelObject: true,
+			compileCommnads:     true,
 		}
 		return filepath.WalkDir(dir, func(path string, d fs.DirEntry, err error) error {
-			if err != nil || d.IsDir() || keepFiles[path] || keepExt[filepath.Ext(d.Name())] {
+			if err != nil {
 				return err
 			}
+			relative, err := filepath.Rel(dir, path)
+			if err != nil {
+				return err
+			}
+			if d.IsDir() || keepFiles[relative] || codesearch.IsSourceFile(relative) {
+				return nil
+			}
 			return os.Remove(path)
 		})
 	})
diff --git a/pkg/codesearch/codesearch.go b/pkg/codesearch/codesearch.go
index 6eab749db..746984369 100644
--- a/pkg/codesearch/codesearch.go
+++ b/pkg/codesearch/codesearch.go
@@ -82,7 +82,17 @@ var Commands = []Command{
 	}},
 }
 
-var SourceExtensions = map[string]bool{".c": true, ".h": true, ".S": true, ".rs": true}
+func IsSourceFile(file string) bool {
+	return sourceFiles[file] || sourceExtensions[filepath.Ext(file)]
+}
+
+var (
+	// Files and extensions we want to keep in the build dir and make available to LLM agents.
+	sourceExtensions = map[string]bool{".c": true, ".h": true, ".S": true, ".rs": true}
+	sourceFiles      = map[string]bool{
+		".config": true,
+	}
+)
 
 func NewIndex(databaseFile string, srcDirs []string) (*Index, error) {
 	db, err := osutil.ReadJSON[*Database](databaseFile)
@@ -275,6 +285,7 @@ func escaping(path string) error {
 }
 
 func dirIndex(root, subdir string) (bool, []string, []string, error) {
+	subdir = filepath.Clean(subdir)
 	dir := filepath.Join(root, subdir)
 	entries, err := os.ReadDir(dir)
 	if err != nil {
@@ -293,7 +304,7 @@ func dirIndex(root, subdir string) (bool, []string, []string, error) {
 			// These are internal things like .git, etc.
 		} else if entry.IsDir() {
 			subdirs = append(subdirs, entry.Name())
-		} else if SourceExtensions[filepath.Ext(entry.Name())] {
+		} else if IsSourceFile(filepath.Join(subdir, entry.Name())) {
 			files = append(files, entry.Name())
 		}
 	}