cmd/bpf2go: generate types from BTF

Use GoFormatter to generate type definitions for types used as map keys
and values. Allow users to emit additional types via a command line flag.

cmd/bpf2go/README.md

@@ -23,6 +23,12 @@ across a project, e.g. to set specific C flags:
 By exporting `$BPF_CFLAGS` from your build system you can then control
 all builds from a single location.
 
+## Generated types
+
+`bpf2go` generates Go types for all map keys and values by default. You can
+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.
+
 ## Examples
 
 See [examples/kprobe](../../examples/kprobe/main.go) for a fully worked out example.

cmd/bpf2go/main.go

@@ -10,6 +10,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"regexp"
 	"runtime"
 	"sort"
 	"strings"
@@ -78,6 +79,8 @@ func run(stdout io.Writer, pkg, outputDir string, args []string) (err error) {
 	fs.StringVar(&b2g.tags, "tags", "", "list of Go build tags to include in generated files")
 	flagTarget := fs.String("target", "bpfel,bpfeb", "clang target to compile for")
 	fs.StringVar(&b2g.makeBase, "makebase", "", "write make compatible depinfo files relative to `directory`")
+	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.SetOutput(stdout)
 	fs.Usage = func() {
@@ -193,6 +196,44 @@ func run(stdout io.Writer, pkg, outputDir string, args []string) (err error) {
 	return nil
 }
 
+// cTypes collects the C type names a user wants to generate Go types for.
+//
+// Names are guaranteed to be unique, and only a subset of names is accepted so
+// that we may extend the flag syntax in the future.
+type cTypes []string
+
+var _ flag.Value = (*cTypes)(nil)
+
+func (ct *cTypes) String() string {
+	if ct == nil {
+		return "[]"
+	}
+	return fmt.Sprint(*ct)
+}
+
+const validCTypeChars = `[a-z0-9_]`
+
+var reValidCType = regexp.MustCompile(`(?i)^` + validCTypeChars + `+$`)
+
+func (ct *cTypes) Set(value string) error {
+	if !reValidCType.MatchString(value) {
+		return fmt.Errorf("%q contains characters outside of %s", value, validCTypeChars)
+	}
+
+	i := sort.SearchStrings(*ct, value)
+	if i >= len(*ct) {
+		*ct = append(*ct, value)
+		return nil
+	}
+
+	if (*ct)[i] == value {
+		return fmt.Errorf("duplicate type %q", value)
+	}
+
+	*ct = append((*ct)[:i], append([]string{value}, (*ct)[i:]...)...)
+	return nil
+}
+
 type bpf2go struct {
 	stdout io.Writer
 	// Absolute path to a .c file.
@@ -209,7 +250,10 @@ type bpf2go struct {
 	strip            string
 	disableStripping bool
 	// C flags passed to the compiler.
-	cFlags []string
+	cFlags          []string
+	skipGlobalTypes bool
+	// C types to include in the generatd output.
+	cTypes cTypes
 	// Go tags included in the .go
 	tags string
 	// Base directory of the Makefile. Enables outputting make-style dependencies
@@ -282,12 +326,14 @@ func (b2g *bpf2go) convert(tgt target, arches []string) (err error) {
 	}
 	defer removeOnError(goFile)
 
-	err = writeCommon(writeArgs{
-		pkg:   b2g.pkg,
-		ident: b2g.ident,
-		tags:  tags,
-		obj:   objFileName,
-		out:   goFile,
+	err = output(outputArgs{
+		pkg:             b2g.pkg,
+		ident:           b2g.ident,
+		cTypes:          b2g.cTypes,
+		skipGlobalTypes: b2g.skipGlobalTypes,
+		tags:            tags,
+		obj:             objFileName,
+		out:             goFile,
 	})
 	if err != nil {
 		return fmt.Errorf("can't write %s: %s", goFileName, err)

cmd/bpf2go/main_test.go

@@ -12,6 +12,7 @@ import (
 	"strings"
 	"testing"
 
+	qt "github.com/frankban/quicktest"
 	"github.com/google/go-cmp/cmp"
 )
 
@@ -243,3 +244,37 @@ func TestConvertGOARCH(t *testing.T) {
 		t.Fatal("Can't target GOARCH:", err)
 	}
 }
+
+func TestCTypes(t *testing.T) {
+	var ct cTypes
+	valid := []string{
+		"abcdefghijklmnopqrstuvqxyABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890_",
+		"y",
+	}
+	for _, value := range valid {
+		if err := ct.Set(value); err != nil {
+			t.Fatalf("Set returned an error for %q: %s", value, err)
+		}
+	}
+	qt.Assert(t, ct, qt.ContentEquals, cTypes(valid))
+
+	for _, value := range []string{
+		"",
+		" ",
+		" frood",
+		"foo\nbar",
+		".",
+		",",
+		"+",
+		"-",
+	} {
+		ct = nil
+		if err := ct.Set(value); err == nil {
+			t.Fatalf("Set did not return an error for %q", value)
+		}
+	}
+
+	ct = nil
+	qt.Assert(t, ct.Set("foo"), qt.IsNil)
+	qt.Assert(t, ct.Set("foo"), qt.IsNotNil)
+}

cmd/bpf2go/output.go

@@ -5,13 +5,15 @@ import (
 	"fmt"
 	"go/token"
 	"io"
-	"os"
+	"io/ioutil"
 	"path/filepath"
+	"sort"
 	"strings"
 	"text/template"
 
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/internal"
+	"github.com/cilium/ebpf/internal/btf"
 )
 
 const ebpfModule = "github.com/cilium/ebpf"
@@ -32,6 +34,13 @@ import (
 	"{{ .Module }}"
 )
 
+{{- if .Types }}
+{{- range $type := .Types }}
+{{ $.TypeDeclaration (index $.TypeNames $type) $type }}
+
+{{ end }}
+{{- end }}
+
 // {{ .Name.Load }} returns the embedded CollectionSpec for {{ .Name }}.
 func {{ .Name.Load }}() (*ebpf.CollectionSpec, error) {
 	reader := bytes.NewReader({{ .Name.Bytes }})
@@ -173,15 +182,15 @@ func (n templateName) Bytes() string {
 }
 
 func (n templateName) Specs() string {
-	return n.maybeExport(string(n) + "Specs")
+	return string(n) + "Specs"
 }
 
 func (n templateName) ProgramSpecs() string {
-	return n.maybeExport(string(n) + "ProgramSpecs")
+	return string(n) + "ProgramSpecs"
 }
 
 func (n templateName) MapSpecs() string {
-	return n.maybeExport(string(n) + "MapSpecs")
+	return string(n) + "MapSpecs"
 }
 
 func (n templateName) Load() string {
@@ -193,36 +202,39 @@ func (n templateName) LoadObjects() string {
 }
 
 func (n templateName) Objects() string {
-	return n.maybeExport(string(n) + "Objects")
+	return string(n) + "Objects"
 }
 
 func (n templateName) Maps() string {
-	return n.maybeExport(string(n) + "Maps")
+	return string(n) + "Maps"
 }
 
 func (n templateName) Programs() string {
-	return n.maybeExport(string(n) + "Programs")
+	return string(n) + "Programs"
 }
 
 func (n templateName) CloseHelper() string {
 	return "_" + toUpperFirst(string(n)) + "Close"
 }
 
-type writeArgs struct {
-	pkg   string
-	ident string
-	tags  []string
-	obj   string
-	out   io.Writer
+type outputArgs struct {
+	pkg             string
+	ident           string
+	tags            []string
+	cTypes          []string
+	skipGlobalTypes bool
+	obj             string
+	out             io.Writer
 }
 
-func writeCommon(args writeArgs) error {
-	obj, err := os.ReadFile(args.obj)
+func output(args outputArgs) error {
+	obj, err := ioutil.ReadFile(args.obj)
 	if err != nil {
 		return fmt.Errorf("read object file contents: %s", err)
 	}
 
-	spec, err := ebpf.LoadCollectionSpecFromReader(bytes.NewReader(obj))
+	rd := bytes.NewReader(obj)
+	spec, err := ebpf.LoadCollectionSpecFromReader(rd)
 	if err != nil {
 		return fmt.Errorf("can't load BPF from ELF: %s", err)
 	}
@@ -242,21 +254,68 @@ func writeCommon(args writeArgs) error {
 		programs[name] = internal.Identifier(name)
 	}
 
+	// Collect any types which we've been asked for explicitly.
+	cTypes, err := collectCTypes(spec.BTF, args.cTypes)
+	if err != nil {
+		return err
+	}
+
+	typeNames := make(map[btf.Type]string)
+	for _, cType := range cTypes {
+		typeNames[cType] = args.ident + internal.Identifier(cType.TypeName())
+	}
+
+	// Collect map key and value types, unless we've been asked not to.
+	if !args.skipGlobalTypes {
+		for _, typ := range collectMapTypes(spec.Maps) {
+			switch btf.UnderlyingType(typ).(type) {
+			case *btf.Datasec:
+				// Avoid emitting .rodata, .bss, etc. for now. We might want to
+				// name these types differently, etc.
+				continue
+
+			case *btf.Int:
+				// Don't emit primitive types by default.
+				continue
+			}
+
+			typeNames[typ] = args.ident + internal.Identifier(typ.TypeName())
+		}
+	}
+
+	// Ensure we don't have conflicting names and generate a sorted list of
+	// named types so that the output is stable.
+	types, err := sortTypes(typeNames)
+	if err != nil {
+		return err
+	}
+
+	gf := &btf.GoFormatter{
+		Names:      typeNames,
+		Identifier: internal.Identifier,
+	}
+
 	ctx := struct {
-		Module   string
-		Package  string
-		Tags     []string
-		Name     templateName
-		Maps     map[string]string
-		Programs map[string]string
-		File     string
+		*btf.GoFormatter
+		Module    string
+		Package   string
+		Tags      []string
+		Name      templateName
+		Maps      map[string]string
+		Programs  map[string]string
+		Types     []btf.Type
+		TypeNames map[btf.Type]string
+		File      string
 	}{
+		gf,
 		ebpfModule,
 		args.pkg,
 		args.tags,
 		templateName(args.ident),
 		maps,
 		programs,
+		types,
+		typeNames,
 		filepath.Base(args.obj),
 	}
 
@@ -268,6 +327,62 @@ func writeCommon(args writeArgs) error {
 	return internal.WriteFormatted(buf.Bytes(), args.out)
 }
 
+func collectCTypes(types *btf.Spec, names []string) ([]btf.Type, error) {
+	var result []btf.Type
+	for _, cType := range names {
+		typ, err := types.AnyTypeByName(cType)
+		if err != nil {
+			return nil, err
+		}
+		result = append(result, typ)
+	}
+	return result, nil
+}
+
+// collectMapTypes returns a list of all types used as map keys or values.
+func collectMapTypes(maps map[string]*ebpf.MapSpec) []btf.Type {
+	var result []btf.Type
+	for _, m := range maps {
+		if m.BTF == nil {
+			continue
+		}
+
+		if m.BTF.Key != nil && m.BTF.Key.TypeName() != "" {
+			result = append(result, m.BTF.Key)
+		}
+
+		if m.BTF.Value != nil && m.BTF.Value.TypeName() != "" {
+			result = append(result, m.BTF.Value)
+		}
+	}
+	return result
+}
+
+// sortTypes returns a list of types sorted by their (generated) Go type name.
+//
+// Duplicate Go type names are rejected.
+func sortTypes(typeNames map[btf.Type]string) ([]btf.Type, error) {
+	var types []btf.Type
+	var names []string
+	for typ, name := range typeNames {
+		i := sort.SearchStrings(names, name)
+		if i >= len(names) {
+			types = append(types, typ)
+			names = append(names, name)
+			continue
+		}
+
+		if names[i] == name {
+			return nil, fmt.Errorf("type name %q is used multiple times", name)
+		}
+
+		types = append(types[:i], append([]btf.Type{typ}, types[i:]...)...)
+		names = append(names[:i], append([]string{name}, names[i:]...)...)
+	}
+
+	return types, nil
+}
+
 func tag(str string) string {
 	return "`ebpf:\"" + str + "\"`"
 }