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.

1 files changed, 13 insertions, 6 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)
 		})
 	})