bpf2go: enable ebpf code reuse across go packages

cmd/bpf2go/README.md

@@ -34,6 +34,15 @@ up-to-date list.
 disable this behaviour using `-no-global-types`. You can add to the set of
 types by specifying `-type foo` for each type you'd like to generate.
 
+## eBPF packages
+
+`bpf2go` can pull header files from other Go packages. List packages you
+wish to pull headers from using `-import` command line option, e.g.:
+`-import example.org/foo`.
+```
+
+Write `#include "example.org/foo/foo.h"` to include `foo.h` from `example.org/foo`.
+
 ## Examples
 
 See [examples/kprobe](../../examples/kprobe/main.go) for a fully worked out example.

cmd/bpf2go/gen/output.go

@@ -98,6 +98,8 @@ type GenerateArgs struct {
 	Types []btf.Type
 	// Filename of the object to embed.
 	ObjectFile string
+	// Packages used during build, record in generated files to prevents "go mod tidy" from removing
+	BuildDeps []string
 	// Output to write template to.
 	Output io.Writer
 	// Function which transforms the input into a valid go identifier. Uses the default behaviour if nil
@@ -162,6 +164,7 @@ func Generate(args GenerateArgs) error {
 	ctx := struct {
 		*btf.GoFormatter
 		Module      string
+		BuildDeps   []string
 		Package     string
 		Constraints constraint.Expr
 		Name        templateName
@@ -174,6 +177,7 @@ func Generate(args GenerateArgs) error {
 	}{
 		gf,
 		b2gInt.CurrentModule,
+		args.BuildDeps,
 		args.Package,
 		args.Constraints,
 		templateName(args.Stem),

cmd/bpf2go/gen/output.tpl

@@ -10,6 +10,13 @@ import (
 	"io"
 
 	"{{ .Module }}"
+{{- if .BuildDeps }}
+
+	// Build dependencies:
+{{- range $dep := .BuildDeps }}
+	_ "{{ $dep }}"
+{{- end }}
+{{- end }}
 )
 
 {{- if .Types }}

cmd/bpf2go/main.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"flag"
 	"fmt"
+	"go/build"
 	"io"
 	"os"
 	"os/exec"
@@ -83,6 +84,8 @@ type bpf2go struct {
 	// Base directory of the Makefile. Enables outputting make-style dependencies
 	// in .d files.
 	makeBase string
+	// Packages contributing ebpf C code.
+	imports []string
 }
 
 func newB2G(stdout io.Writer, args []string) (*bpf2go, error) {
@@ -107,6 +110,10 @@ func newB2G(stdout io.Writer, args []string) (*bpf2go, error) {
 	fs.StringVar(&b2g.outputStem, "output-stem", "", "alternative stem for names of generated files (defaults to ident)")
 	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.Func("import", "pull ebpf C code from specified go package(s)", func(pkg string) error {
+		b2g.imports = append(b2g.imports, pkg)
+		return nil
+	})
 	fs.SetOutput(b2g.stdout)
 	fs.Usage = func() {
 		fmt.Fprintf(fs.Output(), helpText, fs.Name())
@@ -273,6 +280,10 @@ func (b2g *bpf2go) convertAll() (err error) {
 		}
 	}
 
+	if err := b2g.addHeaders(); err != nil {
+		return fmt.Errorf("adding headers: %w", err)
+	}
+
 	for target, arches := range b2g.targetArches {
 		if err := b2g.convert(target, arches); err != nil {
 			return err
@@ -282,6 +293,61 @@ func (b2g *bpf2go) convertAll() (err error) {
 	return nil
 }
 
+// addHeaders exposes header files from packages listed in
+// b2g.imports. C consumes them by giving a golang package path in
+// include, e.g.
+// #include "github.com/cilium/ebpf/foo/bar.h".
+func (b2g *bpf2go) addHeaders() error {
+	var pkgs []*build.Package
+	buildCtx := build.Default
+	buildCtx.Dir = b2g.outputDir
+	for _, path := range b2g.imports {
+		if build.IsLocalImport(path) {
+			return fmt.Errorf("local imports are not supported: %s", path)
+		}
+		pkg, err := buildCtx.Import(path, b2g.outputDir, 0)
+		if err != nil {
+			return err
+		}
+		if pkg.Dir == "" {
+			return fmt.Errorf("%s is missing locally: consider 'go mod download'", path)
+		}
+		if len(hfiles(pkg)) == 0 {
+			fmt.Fprintf(b2g.stdout, "Package doesn't contain .h files: %s\n", path)
+			continue
+		}
+		pkgs = append(pkgs, pkg)
+	}
+
+	if len(pkgs) == 0 {
+		return nil
+	}
+
+	vfs, err := createVfs(pkgs)
+	if err != nil {
+		return err
+	}
+
+	path, err := persistVfs(vfs)
+	if err != nil {
+		return err
+	}
+
+	b2g.cFlags = append([]string{"-ivfsoverlay", path, "-iquote", vfsRootDir}, b2g.cFlags...)
+	return nil
+}
+
+// hfiles lists .h files in a package
+func hfiles(pkg *build.Package) []string {
+	var res []string
+	for _, h := range pkg.HFiles { // includes .hpp, etc
+		if strings.HasSuffix(h, ".h") {
+			res = append(res, h)
+		}
+	}
+	return res
+}
+
 func (b2g *bpf2go) convert(tgt gen.Target, goarches gen.GoArches) (err error) {
 	removeOnError := func(f *os.File) {
 		if err != nil {
@@ -406,6 +472,7 @@ func (b2g *bpf2go) convert(tgt gen.Target, goarches gen.GoArches) (err error) {
 		Programs:    programs,
 		Types:       types,
 		ObjectFile:  filepath.Base(objFileName),
+		BuildDeps:   b2g.imports,
 		Output:      goFile,
 	})
 	if err != nil {

cmd/bpf2go/main_test.go

@@ -32,24 +32,12 @@ func TestRun(t *testing.T) {
 	}
 
 	modDir := t.TempDir()
-	execInModule := func(name string, args ...string) {
-		t.Helper()
-
-		cmd := exec.Command(name, args...)
-		cmd.Dir = modDir
-		if out, err := cmd.CombinedOutput(); err != nil {
-			if out := string(out); out != "" {
-				t.Log(out)
-			}
-			t.Fatalf("Can't execute %s: %v", name, args)
-		}
-	}
 
 	module := internal.CurrentModule
 
-	execInModule("go", "mod", "init", "bpf2go-test")
+	execInDir(t, modDir, "go", "mod", "init", "bpf2go-test")
 
-	execInModule("go", "mod", "edit",
+	execInDir(t, modDir, "go", "mod", "edit",
 		// Require the module. The version doesn't matter due to the replace
 		// below.
 		fmt.Sprintf("-require=%[email protected]", module),
@@ -106,6 +94,66 @@ func main() {
 	}
 }
 
+func execInDir(t *testing.T, dir, name string, args ...string) {
+	t.Helper()
+
+	cmd := exec.Command(name, args...)
+	cmd.Dir = dir
+	if out, err := cmd.CombinedOutput(); err != nil {
+		if out := string(out); out != "" {
+			t.Log(out)
+		}
+		t.Fatalf("Can't execute %s: %v", name, args)
+	}
+}
+
+func TestImports(t *testing.T) {
+	dir := t.TempDir()
+	mustWriteFile(t, dir, "foo/foo.go", "package foo")
+	mustWriteFile(t, dir, "foo/foo.h", "#define EXAMPLE_ORG__FOO__FOO_H 1")
+	mustWriteFile(t, dir, "bar/nested/nested.go", "package nested")
+	mustWriteFile(t, dir, "bar/nested/nested.h", "#define EXAMPLE_ORG__BAR__NESTED__NESTED_H 1")
+	mustWriteFile(t, dir, "bar/bar.c", `
+//go:build ignore
+
+// include from current module, package listed in -import
+#include "example.org/bar/nested/nested.h"
+#ifndef EXAMPLE_ORG__BAR__NESTED__NESTED_H
+#error "example.org/bar/nested/nested.h: unexpected file contents"
+#endif
+
+// include from external module, package listed in -import
+#include "example.org/foo/foo.h"
+#ifndef EXAMPLE_ORG__FOO__FOO_H
+#error "example.org/foo/foo.h: unexpected file contents"
+#endif`)
+
+	fooModDir := filepath.Join(dir, "foo")
+	execInDir(t, fooModDir, "go", "mod", "init", "example.org/foo")
+
+	barModDir := filepath.Join(dir, "bar")
+	execInDir(t, barModDir, "go", "mod", "init", "example.org/bar")
+	execInDir(t, barModDir, "go", "mod", "edit", "-require=example.org/[email protected]")
+
+	execInDir(t, dir, "go", "work", "init")
+	execInDir(t, dir, "go", "work", "use", fooModDir)
+	execInDir(t, dir, "go", "work", "use", barModDir)
+
+	err := run(io.Discard, []string{
+		"-go-package", "bar",
+		"-output-dir", barModDir,
+		"-cc", testutils.ClangBin(t),
+		"-import", "example.org/bar/nested",
+		"-import", "example.org/foo",
+		"bar",
+		filepath.Join(barModDir, "bar.c"),
+	})
+
+	if err != nil {
+		t.Fatal("Can't run:", err)
+	}
+}
+
 func TestHelp(t *testing.T) {
 	var stdout bytes.Buffer
 	err := run(&stdout, []string{"-help"})
@@ -383,6 +431,9 @@ func TestParseArgs(t *testing.T) {
 func mustWriteFile(tb testing.TB, dir, name, contents string) {
 	tb.Helper()
 	tmpFile := filepath.Join(dir, name)
+	if err := os.MkdirAll(filepath.Dir(tmpFile), 0770); err != nil {
+		tb.Fatal(err)
+	}
 	if err := os.WriteFile(tmpFile, []byte(contents), 0660); err != nil {
 		tb.Fatal(err)
 	}

cmd/bpf2go/vfs.go

@@ -0,0 +1,107 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"go/build"
+	"os"
+	"path/filepath"
+	"slices"
+	"strings"
+)
+
+// vfs is LLVM virtual file system parsed from a file
+//
+// In a nutshell, it is a tree of "directory" nodes with leafs being
+// either "file" (a reference to file) or "directory-remap" (a reference
+// to directory).
+//
+// https://github.com/llvm/llvm-project/blob/llvmorg-18.1.0/llvm/include/llvm/Support/VirtualFileSystem.h#L637
+type vfs struct {
+	Version       int       `json:"version"`
+	CaseSensitive bool      `json:"case-sensitive"`
+	Roots         []vfsItem `json:"roots"`
+}
+
+type vfsItem struct {
+	Name             string      `json:"name"`
+	Type             vfsItemType `json:"type"`
+	Contents         []vfsItem   `json:"contents,omitempty"`
+	ExternalContents string      `json:"external-contents,omitempty"`
+}
+
+type vfsItemType string
+
+const (
+	vfsFile      vfsItemType = "file"
+	vfsDirectory vfsItemType = "directory"
+)
+
+func (vi *vfsItem) addDir(path string) (*vfsItem, error) {
+	for _, name := range strings.Split(path, "/") {
+		idx := vi.index(name)
+		if idx == -1 {
+			idx = len(vi.Contents)
+			vi.Contents = append(vi.Contents, vfsItem{Name: name, Type: vfsDirectory})
+		}
+		vi = &vi.Contents[idx]
+		if vi.Type != vfsDirectory {
+			return nil, fmt.Errorf("adding %q: non-directory object already exists", path)
+		}
+	}
+	return vi, nil
+}
+
+func (vi *vfsItem) index(name string) int {
+	return slices.IndexFunc(vi.Contents, func(item vfsItem) bool {
+		return item.Name == name
+	})
+}
+
+func persistVfs(vfs *vfs) (_ string, retErr error) {
+	temp, err := os.CreateTemp("", "")
+	if err != nil {
+		return "", err
+	}
+	defer func() {
+		temp.Close()
+		if retErr != nil {
+			os.Remove(temp.Name())
+		}
+	}()
+
+	if err = json.NewEncoder(temp).Encode(vfs); err != nil {
+		return "", err
+	}
+
+	return temp.Name(), nil
+}
+
+// vfsRootDir is the (virtual) directory where we mount go module sources
+// for the C includes to pick them, e.g. "<vfsRootDir>/github.com/cilium/ebpf".
+const vfsRootDir = "/.vfsoverlay.d"
+
+// createVfs produces a vfs from a list of packages. It creates a
+// (virtual) directory tree reflecting package import paths and adds
+// links to header files. E.g. for github.com/foo/bar containing awesome.h:
+//
+//	github.com/
+//	  foo/
+//	    bar/
+//	      awesome.h -> $HOME/go/pkg/mod/github.com/foo/bar@version/awesome.h
+func createVfs(pkgs []*build.Package) (*vfs, error) {
+	roots := [1]vfsItem{{Name: vfsRootDir, Type: vfsDirectory}}
+	for _, pkg := range pkgs {
+		var headers []vfsItem
+		for _, h := range hfiles(pkg) {
+			headers = append(headers, vfsItem{Name: h, Type: vfsFile,
+				ExternalContents: filepath.Join(pkg.Dir, h)})
+		}
+		dir, err := roots[0].addDir(pkg.ImportPath)
+		if err != nil {
+			return nil, err
+		}
+		dir.Contents = headers // NB don't append inplace, same package could be imported twice
+	}
+	return &vfs{CaseSensitive: true, Roots: roots[:]}, nil
+}