cmd/bpf2go: enable building custom tools (wip)

cmd/bpf2go/compile.go

@@ -28,7 +28,7 @@ type compileArgs struct {
 	dep io.Writer
 }
 
-func compile(args compileArgs) error {
+func compile(args compileArgs, tc TargetCustomisations) error {
 	// Default cflags that can be overridden by args.cFlags
 	overrideFlags := []string{
 		// Code needs to be optimized, otherwise the verifier will often fail
@@ -90,7 +90,7 @@ func compile(args compileArgs) error {
 		)
 	}
 
-	if err := cmd.Run(); err != nil {
+	if err := tc.Compile(cmd); err != nil {
 		return fmt.Errorf("can't execute %s: %s", args.cc, err)
 	}
 
@@ -200,10 +200,10 @@ func parseDependencies(baseDir string, in io.Reader) ([]dependency, error) {
 }
 
 // strip DWARF debug info from file by executing exe.
-func strip(exe, file string) error {
+func strip(exe, file string, tc TargetCustomisations) error {
 	cmd := exec.Command(exe, "-g", file)
 	cmd.Stderr = os.Stderr
-	if err := cmd.Run(); err != nil {
+	if err := tc.Strip(cmd); err != nil {
 		return fmt.Errorf("%s: %s", exe, err)
 	}
 	return nil

cmd/bpf2go/compile_test.go

@@ -22,7 +22,7 @@ func TestCompile(t *testing.T) {
 		source: filepath.Join(dir, "test.c"),
 		dest:   filepath.Join(dir, "test.o"),
 		dep:    &dep,
-	})
+	}, &CustomisationsBase{})
 	if err != nil {
 		t.Fatal("Can't compile:", err)
 	}
@@ -55,7 +55,7 @@ func TestReproducibleCompile(t *testing.T) {
 		dir:    dir,
 		source: filepath.Join(dir, "test.c"),
 		dest:   filepath.Join(dir, "a.o"),
-	})
+	}, &CustomisationsBase{})
 	if err != nil {
 		t.Fatal("Can't compile:", err)
 	}
@@ -65,7 +65,7 @@ func TestReproducibleCompile(t *testing.T) {
 		dir:    dir,
 		source: filepath.Join(dir, "test.c"),
 		dest:   filepath.Join(dir, "b.o"),
-	})
+	}, &CustomisationsBase{})
 	if err != nil {
 		t.Fatal("Can't compile:", err)
 	}
@@ -94,7 +94,7 @@ func TestTriggerMissingTarget(t *testing.T) {
 		dir:    dir,
 		source: filepath.Join(dir, "test.c"),
 		dest:   filepath.Join(dir, "a.o"),
-	})
+	}, &CustomisationsBase{})
 
 	if err == nil {
 		t.Fatal("No error when compiling __BPF_TARGET_MISSING")

cmd/bpf2go/customisations.go

@@ -0,0 +1,67 @@
+package main
+
+import (
+	"flag"
+	"io"
+	"os/exec"
+
+	"github.com/cilium/ebpf"
+)
+
+// Run implements bpf2go tool; leverage bpf2go and build custom tools on
+// top by providing Customisations
+func Run(stdout io.Writer, pkg, outputDir string, args []string, c Customisations) (err error) {
+	return run(stdout, pkg, outputDir, args, c)
+}
+
+// Customisations enables building custom tools on top of bpf2go
+type Customisations interface {
+	// ModifyFlags adds custom tool's own flags
+	ModifyFlags(*flag.FlagSet)
+
+	// NewTarget provides decdicated TargetCustomisations per target; a
+	// custom tool implements further customisations via
+	// TargetCustomisations (can contain target-specific state)
+	NewTarget(target) TargetCustomisations
+}
+
+// TargetCustomisations enables building custom tools on top of bpf2go
+type TargetCustomisations interface {
+	// Compile compiles C source code; instead of running the provided
+	// command as is, a custom tool could do something more exciting,
+	// such as transparently running compiler in a Docker container or
+	// implement a cache, or even run a custom preprocessor to have
+	// e.g. golang package-aware includes
+	Compile(*exec.Cmd) error
+
+	// Strip strips the object file; instead of running the provided
+	// command as is, a custom tool could do something more exciting
+	// such as transparently running strip in a Docker container or
+	// implement a cache
+	Strip(*exec.Cmd) error
+
+	// AugmentSpec optionally adds to the spec; e.g. a custom tool could
+	// add types from DWARF debug info so that the tool could generate
+	// golang type even when C type is not included in BTF
+	AugmentSpec(*ebpf.CollectionSpec) error
+}
+
+// CustomisationsBase is a baseline implementation of Cusomisations,
+// TargetCustomisations
+type CustomisationsBase struct{}
+
+func (cb *CustomisationsBase) ModifyFlags(fs *flag.FlagSet) {}
+
+func (cb *CustomisationsBase) NewTarget(target) TargetCustomisations { return cb }
+
+func (cb *CustomisationsBase) Compile(cmd *exec.Cmd) error {
+	return cmd.Run()
+}
+
+func (cb *CustomisationsBase) Strip(cmd *exec.Cmd) error {
+	return cmd.Run()
+}
+
+func (cb *CustomisationsBase) AugmentSpec(*ebpf.CollectionSpec) error {
+	return nil
+}

cmd/bpf2go/main.go

@@ -68,8 +68,8 @@ var targetByGoArch = map[goarch]target{
 	"s390x":    {"bpfeb", "s390"},
 }
 
-func run(stdout io.Writer, pkg, outputDir string, args []string) (err error) {
-	b2g, err := newB2G(stdout, pkg, outputDir, args)
+func run(stdout io.Writer, pkg, outputDir string, args []string, c Customisations) (err error) {
+	b2g, err := newB2G(stdout, pkg, outputDir, args, c)
 	switch {
 	case err == nil:
 		return b2g.convertAll()
@@ -109,13 +109,17 @@ type bpf2go struct {
 	// Base directory of the Makefile. Enables outputting make-style dependencies
 	// in .d files.
 	makeBase string
+
+	customisations Customisations
 }
 
-func newB2G(stdout io.Writer, pkg, outputDir string, args []string) (*bpf2go, error) {
+func newB2G(stdout io.Writer, pkg, outputDir string, args []string, c Customisations) (*bpf2go, error) {
 	b2g := &bpf2go{
 		stdout:    stdout,
 		pkg:       pkg,
 		outputDir: outputDir,
+
+		customisations: c,
 	}
 
 	fs := flag.NewFlagSet("bpf2go", flag.ContinueOnError)
@@ -141,6 +145,7 @@ func newB2G(stdout io.Writer, pkg, outputDir string, args []string) (*bpf2go, er
 		fmt.Fprintln(fs.Output())
 		printTargets(fs.Output())
 	}
+	c.ModifyFlags(fs)
 	if err := fs.Parse(args); err != nil {
 		return nil, err
 	}
@@ -345,6 +350,8 @@ func (b2g *bpf2go) convert(tgt target, goarches []goarch) (err error) {
 		return fmt.Errorf("remove obsolete output: %w", err)
 	}
 
+	tc := b2g.customisations.NewTarget(tgt)
+
 	var dep bytes.Buffer
 	err = compile(compileArgs{
 		cc:     b2g.cc,
@@ -354,15 +361,15 @@ func (b2g *bpf2go) convert(tgt target, goarches []goarch) (err error) {
 		source: b2g.sourceFile,
 		dest:   objFileName,
 		dep:    &dep,
-	})
+	}, tc)
 	if err != nil {
 		return err
 	}
 
 	fmt.Fprintln(b2g.stdout, "Compiled", objFileName)
 
 	if !b2g.disableStripping {
-		if err := strip(b2g.strip, objFileName); err != nil {
+		if err := strip(b2g.strip, objFileName, tc); err != nil {
 			return err
 		}
 		fmt.Fprintln(b2g.stdout, "Stripped", objFileName)
@@ -373,6 +380,11 @@ func (b2g *bpf2go) convert(tgt target, goarches []goarch) (err error) {
 		return fmt.Errorf("can't load BPF from ELF: %s", err)
 	}
 
+	err = tc.AugmentSpec(spec)
+	if err != nil {
+		return err
+	}
+
 	maps, programs, types, err := collectFromSpec(spec, b2g.cTypes, b2g.skipGlobalTypes)
 	if err != nil {
 		return err
@@ -533,7 +545,7 @@ func main() {
 		os.Exit(1)
 	}
 
-	if err := run(os.Stdout, os.Getenv(gopackageEnv), outputDir, os.Args[1:]); err != nil {
+	if err := run(os.Stdout, os.Getenv(gopackageEnv), outputDir, os.Args[1:], &CustomisationsBase{}); err != nil {
 		fmt.Fprintln(os.Stderr, "Error:", err)
 		os.Exit(1)
 	}