btf: fix NewHandleFromID for kernel modules

Kernel modules only expose split BTF when queried via BPF_OBJ_GET_INFO_BY_FD.
NewHandleFromID therefore doesn't work when trying to obtain a handle for a kernel
module BTF. At the same time, always parsing BTF when creating a handle from an ID
is slow and wasteful.

Change Handle.Spec to accept a base Spec to support split BTF from modules.
This is a breaking change.

Updates #705

btf/btf.go

@@ -696,8 +696,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.
@@ -740,7 +742,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.
@@ -759,15 +761,32 @@ func NewHandleFromID(id ID) (*Handle, error) {
 	info, err := newInfoFromFd(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.
@@ -782,6 +801,11 @@ func (h *Handle) FD() int {
 	return h.fd.Int()
 }
 
+// Info returns metadata about the handle.
+func (h *Handle) Info() (*Info, error) {
+	return newInfoFromFd(h.fd)
+}
+
 func marshalBTF(types interface{}, strings []byte, bo binary.ByteOrder) []byte {
 	const minHeaderLength = 24
 

btf/handle_test.go

@@ -16,27 +16,94 @@ func TestNewHandleFromID(t *testing.T) {
 	testutils.SkipOnOldKernel(t, "5.11", "vmlinux BTF ID")
 
 	var id btf.ID
+	h := nextHandle(t, &id)
+	if h == nil {
+		t.Fatalf("No BTF loaded")
+	}
+}
+
+func TestParseModuleSplitSpec(t *testing.T) {
+	// See TestNewHandleFromID for reasoning.
+	testutils.SkipOnOldKernel(t, "5.11", "vmlinux BTF ID")
+
+	var module *btf.Handle
+	for id := btf.ID(0); ; {
+		module = nextHandle(t, &id)
+		if module == nil {
+			t.Fatalf("Can't find module BTF")
+		}
+
+		info, err := module.Info()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if !info.IsModule() {
+			continue
+		}
+
+		t.Log("Using module", info.Name)
+		break
+	}
+
+	var vmlinux *btf.Handle
+	for id := btf.ID(0); ; {
+		vmlinux = nextHandle(t, &id)
+		if vmlinux == nil {
+			t.Fatalf("Can't find vmlinux BTF")
+		}
+
+		info, err := vmlinux.Info()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if !info.IsVmlinux() {
+			continue
+		}
+
+		break
+	}
+
+	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 nextHandle(t *testing.T, prevID *btf.ID) *btf.Handle {
+	t.Helper()
+
 	for {
 		var err error
-		id, err = btf.GetNextID(id)
+		*prevID, err = btf.GetNextID(*prevID)
 		if errors.Is(err, os.ErrNotExist) {
-			t.Fatalf("No BTF loaded")
+			return nil
 		}
-
 		if err != nil {
 			t.Fatal(err)
 		}
 
-		h, err := btf.NewHandleFromID(id)
+		h, err := btf.NewHandleFromID(*prevID)
 		if errors.Is(err, os.ErrNotExist) {
 			// Likely a race where BTF was unloaded before we could retrieve an fd.
 			continue
 		}
 		if err != nil {
 			t.Fatal(err)
 		}
-		h.Close()
 
-		return
+		t.Cleanup(func() { h.Close() })
+		return h
 	}
 }

btf/info.go

@@ -1,56 +1,65 @@
 package btf
 
 import (
-	"bytes"
+	"fmt"
 
-	"github.com/cilium/ebpf/internal"
 	"github.com/cilium/ebpf/internal/sys"
 	"github.com/cilium/ebpf/internal/unix"
 )
 
-// info describes a BTF object.
-type info struct {
-	BTF *Spec
-	ID  ID
+// Info describes a Handle.
+type Info struct {
+	ID ID
 	// Name is an identifying name for the BTF, currently only used by the
 	// kernel.
 	Name string
-	// KernelBTF is true if the BTf originated with the kernel and not
+
+	// KernelBTF is true if the BTF originated with the kernel and not
 	// userspace.
 	KernelBTF bool
+
+	// Size of the raw BTF in bytes.
+	size uint32
 }
 
-func newInfoFromFd(fd *sys.FD) (*info, error) {
+func newInfoFromFd(fd *sys.FD) (*Info, 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, err
+		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++
 	}
 
-	btfBuffer := make([]byte, btfInfo.BtfSize)
+	// 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.Btf, btfInfo.BtfSize = sys.NewSlicePointerLen(btfBuffer)
 	btfInfo.Name, btfInfo.NameLen = sys.NewSlicePointerLen(nameBuffer)
 	if err := sys.ObjInfo(fd, &btfInfo); err != nil {
 		return nil, err
 	}
 
-	spec, err := loadRawSpec(bytes.NewReader(btfBuffer), internal.NativeEndian, nil, nil)
-	if err != nil {
-		return nil, err
-	}
-
-	return &info{
-		BTF:       spec,
+	return &Info{
 		ID:        btfInfo.Id,
 		Name:      unix.ByteSliceToString(nameBuffer),
 		KernelBTF: btfInfo.KernelBtf != 0,
+		size:      btfSize,
 	}, nil
 }
+
+// IsModule returns true if the BTF is for the kernel itself.
+func (i *Info) IsVmlinux() bool {
+	return i.KernelBTF && i.Name == "vmlinux"
+}
+
+// IsModule returns true if the BTF is for a kernel module.
+func (i *Info) IsModule() bool {
+	return i.KernelBTF && i.Name != "vmlinux"
+}

link/tracing.go

@@ -55,13 +55,18 @@ func AttachFreplace(targetProg *ebpf.Program, name string, prog *ebpf.Program) (
 		}
 		defer btfHandle.Close()
 
+		spec, err := btfHandle.Spec(nil)
+		if err != nil {
+			return nil, err
+		}
+
 		var function *btf.Func
-		if err := btfHandle.Spec().TypeByName(name, &function); err != nil {
+		if err := spec.TypeByName(name, &function); err != nil {
 			return nil, err
 		}
 
 		target = targetProg.FD()
-		typeID, err = btfHandle.Spec().TypeID(function)
+		typeID, err = spec.TypeID(function)
 		if err != nil {
 			return nil, err
 		}

prog.go

@@ -852,7 +852,10 @@ func findTargetInProgram(prog *Program, name string, progType ProgramType, attac
 	}
 	defer btfHandle.Close()
 
-	spec := btfHandle.Spec()
+	spec, err := btfHandle.Spec(nil)
+	if err != nil {
+		return 0, err
+	}
 
 	var targetFunc *btf.Func
 	err = spec.TypeByName(typeName, &targetFunc)