btf: fix NewHandleFromID for kernel modules

btf/btf.go

@@ -27,7 +27,7 @@ var (
 )
 
 // ID represents the unique ID of a BTF object.
-type ID uint32
+type ID = sys.BTFID
 
 // Spec represents decoded BTF.
 type Spec struct {
@@ -685,8 +685,10 @@ func (iter *TypesIterator) Next() bool {
 
 // Handle is a reference to BTF loaded into the kernel.
 type Handle struct {
-	spec *Spec
-	fd   *sys.FD
+	fd *sys.FD
+
+	// Size of the raw BTF in bytes.
+	size uint32
 }
 
 // NewHandle loads BTF into the kernel.
@@ -729,7 +731,7 @@ func NewHandle(spec *Spec) (*Handle, error) {
 		return nil, internal.ErrorWithLog(err, logBuf)
 	}
 
-	return &Handle{spec.Copy(), fd}, nil
+	return &Handle{fd, attr.BtfSize}, nil
 }
 
 // NewHandleFromID returns the BTF handle for a given id.
@@ -745,24 +747,45 @@ func NewHandleFromID(id ID) (*Handle, error) {
 		return nil, fmt.Errorf("get FD for ID %d: %w", id, err)
 	}
 
-	info, err := newInfoFromFd(fd)
+	info, err := newHandleInfoFromFD(fd)
 	if err != nil {
 		_ = fd.Close()
-		return nil, fmt.Errorf("get BTF spec for handle: %w", err)
+		return nil, err
 	}
 
-	return &Handle{info.BTF, fd}, nil
+	return &Handle{fd, info.size}, nil
 }
 
-// Spec returns the Spec that defined the BTF loaded into the kernel.
-func (h *Handle) Spec() *Spec {
-	return h.spec
+// Spec parses the kernel BTF into Go types.
+//
+// base is used to decode split BTF and may be nil.
+func (h *Handle) Spec(base *Spec) (*Spec, error) {
+	var btfInfo sys.BtfInfo
+	btfBuffer := make([]byte, h.size)
+	btfInfo.Btf, btfInfo.BtfSize = sys.NewSlicePointerLen(btfBuffer)
+
+	if err := sys.ObjInfo(h.fd, &btfInfo); err != nil {
+		return nil, err
+	}
+
+	var baseTypes types
+	var baseStrings *stringTable
+	if base != nil {
+		baseTypes = base.types
+		baseStrings = base.strings
+	}
+
+	return loadRawSpec(bytes.NewReader(btfBuffer), internal.NativeEndian, baseTypes, baseStrings)
 }
 
 // Close destroys the handle.
 //
 // Subsequent calls to FD will return an invalid value.
 func (h *Handle) Close() error {
+	if h == nil {
+		return nil
+	}
+
 	return h.fd.Close()
 }
 
@@ -771,6 +794,11 @@ func (h *Handle) FD() int {
 	return h.fd.Int()
 }
 
+// Info returns metadata about the handle.
+func (h *Handle) Info() (*HandleInfo, error) {
+	return newHandleInfoFromFD(h.fd)
+}
+
 func marshalBTF(types interface{}, strings []byte, bo binary.ByteOrder) []byte {
 	const minHeaderLength = 24
 

btf/handle.go

@@ -0,0 +1,121 @@
+package btf
+
+import (
+	"errors"
+	"fmt"
+	"os"
+
+	"github.com/cilium/ebpf/internal/sys"
+	"github.com/cilium/ebpf/internal/unix"
+)
+
+// HandleInfo describes a Handle.
+type HandleInfo struct {
+	// ID of this handle in the kernel. The ID is only valid as long as the
+	// associated handle is kept alive.
+	ID ID
+
+	// Name is an identifying name for the BTF, currently only used by the
+	// kernel.
+	Name string
+
+	// IsKernel is true if the BTF originated with the kernel and not
+	// userspace.
+	IsKernel bool
+
+	// Size of the raw BTF in bytes.
+	size uint32
+}
+
+func newHandleInfoFromFD(fd *sys.FD) (*HandleInfo, error) {
+	// We invoke the syscall once with a empty BTF and name buffers to get size
+	// information to allocate buffers. Then we invoke it a second time with
+	// buffers to receive the data.
+	var btfInfo sys.BtfInfo
+	if err := sys.ObjInfo(fd, &btfInfo); err != nil {
+		return nil, fmt.Errorf("get BTF info for fd %s: %w", fd, err)
+	}
+
+	if btfInfo.NameLen > 0 {
+		// NameLen doesn't account for the terminating NUL.
+		btfInfo.NameLen++
+	}
+
+	// Don't pull raw BTF by default, since it may be quite large.
+	btfSize := btfInfo.BtfSize
+	btfInfo.BtfSize = 0
+
+	nameBuffer := make([]byte, btfInfo.NameLen)
+	btfInfo.Name, btfInfo.NameLen = sys.NewSlicePointerLen(nameBuffer)
+	if err := sys.ObjInfo(fd, &btfInfo); err != nil {
+		return nil, err
+	}
+
+	return &HandleInfo{
+		ID:       ID(btfInfo.Id),
+		Name:     unix.ByteSliceToString(nameBuffer),
+		IsKernel: btfInfo.KernelBtf != 0,
+		size:     btfSize,
+	}, nil
+}
+
+// IsModule returns true if the BTF is for the kernel itself.
+func (i *HandleInfo) IsVmlinux() bool {
+	return i.IsKernel && i.Name == "vmlinux"
+}
+
+// IsModule returns true if the BTF is for a kernel module.
+func (i *HandleInfo) IsModule() bool {
+	return i.IsKernel && i.Name != "vmlinux"
+}
+
+// HandleIterator allows enumerating BTF blobs loaded into the kernel.
+type HandleIterator struct {
+	// The ID of the last retrieved handle. Only valid after a call to Next.
+	ID  ID
+	err error
+}
+
+// Next retrieves a handle for the next BTF blob.
+//
+// [Handle.Close] is called if *handle is non-nil to avoid leaking fds.
+//
+// Returns true if another BTF blob was found. Call [HandleIterator.Err] after
+// the function returns false.
+func (it *HandleIterator) Next(handle **Handle) bool {
+	if *handle != nil {
+		(*handle).Close()
+		*handle = nil
+	}
+
+	id := it.ID
+	for {
+		attr := &sys.BtfGetNextIdAttr{Id: id}
+		err := sys.BtfGetNextId(attr)
+		if errors.Is(err, os.ErrNotExist) {
+			// There are no more BTF objects.
+			return false
+		} else if err != nil {
+			it.err = fmt.Errorf("get next BTF ID: %w", err)
+			return false
+		}
+
+		id = attr.NextId
+		*handle, err = NewHandleFromID(id)
+		if errors.Is(err, os.ErrNotExist) {
+			// Try again with the next ID.
+			continue
+		} else if err != nil {
+			it.err = fmt.Errorf("retrieve handle for ID %d: %w", id, err)
+			return false
+		}
+
+		it.ID = id
+		return true
+	}
+}
+
+// Err returns an error if iteration failed for some reason.
+func (it *HandleIterator) Err() error {
+	return it.err
+}

btf/handle_test.go

@@ -1,23 +1,131 @@
 package btf_test
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/cilium/ebpf/btf"
 	"github.com/cilium/ebpf/internal/testutils"
 )
 
-func TestNewHandleFromID(t *testing.T) {
-	// vmlinux is not guaranteed to be at ID 1, but it's highly likely, since
-	// module loading causes vmlinux to be parsed.
-	const vmlinux = btf.ID(1)
-
+func TestHandleIterator(t *testing.T) {
+	// There is no guarantee that there is a BTF ID allocated, but loading a module
+	// triggers loading vmlinux.
 	// See https://github.com/torvalds/linux/commit/5329722057d41aebc31e391907a501feaa42f7d9
 	testutils.SkipOnOldKernel(t, "5.11", "vmlinux BTF ID")
 
-	h, err := btf.NewHandleFromID(vmlinux)
-	if err != nil {
+	var h *btf.Handle
+	defer h.Close()
+
+	it := new(btf.HandleIterator)
+	if !it.Next(&h) {
+		t.Fatalf("No BTF loaded")
+	}
+	if h == nil {
+		t.Fatal("Next doesn't assign handle")
+	}
+	prev := it.ID
+	for it.Next(&h) {
+		// Iterate all loaded BTF.
+		if h == nil {
+			t.Fatal("Next doesn't assign handle")
+		}
+		if it.ID == prev {
+			t.Fatal("Iterator doesn't advance ID")
+		}
+		prev = it.ID
+	}
+	if err := it.Err(); err != nil {
+		t.Fatal("Iteration returned an error:", err)
+	}
+
+	if h != nil {
+		t.Fatal("Next doesn't clean up handle on last iteration")
+	}
+	if prev != it.ID {
+		t.Fatal("Next changes ID on last iteration")
+	}
+}
+
+func TestParseModuleSplitSpec(t *testing.T) {
+	// See TestNewHandleFromID for reasoning.
+	testutils.SkipOnOldKernel(t, "5.11", "vmlinux BTF ID")
+
+	var module *btf.Handle
+	defer module.Close()
+
+	it := new(btf.HandleIterator)
+	for it.Next(&module) {
+		info, err := module.Info()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if !info.IsModule() {
+			continue
+		}
+
+		t.Log("Using module", info.Name)
+		break
+	}
+	if err := it.Err(); err != nil {
+		t.Fatal(err)
+	}
+
+	if module == nil {
+		t.Fatal("No BTF for kernel module found")
+	}
+
+	var vmlinux *btf.Handle
+	defer vmlinux.Close()
+
+	it = new(btf.HandleIterator)
+	for it.Next(&vmlinux) {
+		info, err := vmlinux.Info()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if !info.IsVmlinux() {
+			continue
+		}
+
+		break
+	}
+	if err := it.Err(); err != nil {
 		t.Fatal(err)
 	}
-	h.Close()
+
+	if vmlinux == nil {
+		t.Fatal("No BTF for kernel found")
+	}
+
+	vmlinuxSpec, err := vmlinux.Spec(nil)
+	if err != nil {
+		t.Fatal("Parse vmlinux BTF:", err)
+	}
+
+	_, err = module.Spec(vmlinuxSpec)
+	if err != nil {
+		t.Fatal("Parse module BTF:", err)
+	}
+
+	_, err = module.Spec(nil)
+	if err == nil {
+		t.Fatal("Parsing module BTF without vmlinux base didn't fail")
+	}
+}
+
+func ExampleHandleIterator() {
+	var handle *btf.Handle
+	// Ensure that handle is cleaned up. This is valid for nil handles as well.
+	defer handle.Close()
+
+	it := new(btf.HandleIterator)
+	for it.Next(&handle) {
+		fmt.Printf("Found handle with ID %d\n", it.ID)
+	}
+	if err := it.Err(); err != nil {
+		panic(err)
+	}
 }