From b8058284a8cd30ddcaadabf299be24901f73565a Mon Sep 17 00:00:00 2001
From: Nick Zavaritsky <mejedi@gmail.com>
Date: Thu, 12 Dec 2024 08:20:04 +0000
Subject: [PATCH] bpf2go: optionally add suffix in output file names

Testing complex ebpf code sometimes requires writing C code specifically
for the test. Exclude testing code from the build unless running tests.

Add "_test" suffix in generated file names if invoking "go:generate"
directive was in a test file (exposed via $GOFILE).

Provide -output-suffix command line option to override.

Signed-off-by: Nick Zavaritsky <mejedi@gmail.com>
---
 cmd/bpf2go/main.go      | 14 +++++++++++++-
 cmd/bpf2go/main_test.go | 24 ++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/cmd/bpf2go/main.go b/cmd/bpf2go/main.go
index 367246ba9..91e51e349 100644
--- a/cmd/bpf2go/main.go
+++ b/cmd/bpf2go/main.go
@@ -62,6 +62,8 @@ type bpf2go struct {
 	outputDir string
 	// Alternative output stem. If empty, identStem is used.
 	outputStem string
+	// Suffix in generated file names such as _test.
+	outputSuffix string
 	// Valid go package name.
 	pkg string
 	// Valid go identifier.
@@ -105,6 +107,12 @@ func newB2G(stdout io.Writer, args []string) (*bpf2go, error) {
 	fs.Var(&b2g.cTypes, "type", "`Name` of a type to generate a Go declaration for, may be repeated")
 	fs.BoolVar(&b2g.skipGlobalTypes, "no-global-types", false, "Skip generating types for map keys and values, etc.")
 	fs.StringVar(&b2g.outputStem, "output-stem", "", "alternative stem for names of generated files (defaults to ident)")
+	outputSuffix := ""
+	if strings.HasSuffix(getEnv("GOFILE", ""), "_test.go") {
+		outputSuffix = "_test"
+	}
+	fs.StringVar(&b2g.outputSuffix, "output-suffix", outputSuffix,
+		"suffix in generated file names such as _test (default based on $GOFILE)")
 	outDir := fs.String("output-dir", "", "target directory of generated files (defaults to current directory)")
 	outPkg := fs.String("go-package", "", "package for output go file (default as ENV GOPACKAGE)")
 	fs.SetOutput(b2g.stdout)
@@ -183,6 +191,10 @@ func newB2G(stdout io.Writer, args []string) (*bpf2go, error) {
 		return nil, fmt.Errorf("-output-stem %q must not contain path separation characters", b2g.outputStem)
 	}
 
+	if strings.ContainsRune(b2g.outputSuffix, filepath.Separator) {
+		return nil, fmt.Errorf("-output-suffix %q must not contain path separation characters", b2g.outputSuffix)
+	}
+
 	targetArches := make(map[gen.Target]gen.GoArches)
 	for _, tgt := range strings.Split(*flagTarget, ",") {
 		target, goarches, err := gen.FindTarget(tgt)
@@ -295,7 +307,7 @@ func (b2g *bpf2go) convert(tgt gen.Target, goarches gen.GoArches) (err error) {
 		outputStem = strings.ToLower(b2g.identStem)
 	}
 
-	stem := fmt.Sprintf("%s_%s", outputStem, tgt.Suffix())
+	stem := fmt.Sprintf("%s_%s%s", outputStem, tgt.Suffix(), b2g.outputSuffix)
 
 	absOutPath, err := filepath.Abs(b2g.outputDir)
 	if err != nil {
diff --git a/cmd/bpf2go/main_test.go b/cmd/bpf2go/main_test.go
index 5cf5cf027..eefdb420e 100644
--- a/cmd/bpf2go/main_test.go
+++ b/cmd/bpf2go/main_test.go
@@ -378,6 +378,30 @@ func TestParseArgs(t *testing.T) {
 		qt.Assert(t, qt.IsNil(err))
 		qt.Assert(t, qt.Equals(b2g.outputDir, outputDir))
 	})
+
+	t.Run("output suffix default", func(t *testing.T) {
+		t.Setenv(gopackageEnv, pkg)
+		b2g, err := newB2G(&bytes.Buffer{}, []string{stem, csource})
+		qt.Assert(t, qt.IsNil(err))
+		qt.Assert(t, qt.Equals(b2g.outputSuffix, ""))
+	})
+
+	t.Run("output suffix GOFILE=_test", func(t *testing.T) {
+		t.Setenv(gopackageEnv, pkg)
+		t.Setenv("GOFILE", "foo_test.go")
+		b2g, err := newB2G(&bytes.Buffer{}, []string{stem, csource})
+		qt.Assert(t, qt.IsNil(err))
+		qt.Assert(t, qt.Equals(b2g.outputSuffix, "_test"))
+	})
+
+	t.Run("output suffix custom", func(t *testing.T) {
+		t.Setenv(gopackageEnv, pkg)
+		t.Setenv("GOFILE", "foo_test.go")
+		args := []string{"-output-suffix", "_custom", stem, csource}
+		b2g, err := newB2G(&bytes.Buffer{}, args)
+		qt.Assert(t, qt.IsNil(err))
+		qt.Assert(t, qt.Equals(b2g.outputSuffix, "_custom"))
+	})
 }
 
 func mustWriteFile(tb testing.TB, dir, name, contents string) {