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 <[email protected]>
cilium · Jan 8, 2025 · b805828 · b805828
1 parent be7a2fa
commit b805828
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 1 deletion.
diff --git 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
@@ -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) {