bpf2go: enable ebpf code reuse across go packages

cmd/bpf2go/README.md

@@ -34,6 +34,20 @@ 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
+
+The tool can pull header files from other Go packages. To enable, create
+`bpf2go.hfiles.go` in the output dir. Add packages you wish to pull headers
+from as imports, e.g.:
+
+```
+// bpf2go.hfiles.go
+package awesome
+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/main.go

@@ -4,13 +4,17 @@ import (
 	"errors"
 	"flag"
 	"fmt"
+	"go/build"
+	"go/parser"
+	"go/token"
 	"io"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"regexp"
 	"slices"
 	"sort"
+	"strconv"
 	"strings"
 
 	"github.com/cilium/ebpf"
@@ -273,6 +277,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 +290,79 @@ func (b2g *bpf2go) convertAll() (err error) {
 	return nil
 }
 
+// addHeaders exposes header files from packages listed in
+// $OUTPUT_DIR/bpf2go.hfiles.go to clang. 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 {
+	f, err := os.Open(filepath.Join(b2g.outputDir, "bpf2go.hfiles.go"))
+	if os.IsNotExist(err) {
+		return nil
+	}
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	fmt.Fprintf(b2g.stdout, "Processing packages listed in %s\n", f.Name())
+
+	fset := token.NewFileSet()
+	ast, err := parser.ParseFile(fset, f.Name(), f, parser.ImportsOnly)
+	if err != nil {
+		return err
+	}
+
+	var pkgs []*build.Package
+	buildCtx := build.Default
+	buildCtx.Dir = b2g.outputDir
+	for _, s := range ast.Imports {
+		path, _ := strconv.Unquote(s.Path.Value)
+		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 {

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,70 @@ 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/bpf2go.hfiles.go", `
+package bar
+import (
+	_ "example.org/bar/nested"
+	_ "example.org/foo"
+)`)
+	mustWriteFile(t, dir, "bar/bar.c", `
+//go:build ignore
+
+// include from current module, package listed in bpf2go.hfiles.go
+#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 bpf2go.hfiles.go
+#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),
+		"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 +435,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
+}