review changes

Signed-off-by: Bryce Kahle <bryce.kahle@datadoghq.com>

btf/btf.go

@@ -13,6 +13,7 @@ import (
 	"sync"
 
 	"github.com/cilium/ebpf/internal"
+	"github.com/cilium/ebpf/internal/kallsyms"
 	"github.com/cilium/ebpf/internal/sys"
 )
 
@@ -299,6 +300,8 @@ func FlushKernelSpec() {
 
 	kernelBTF.spec, kernelBTF.fallback = nil, false
 	kernelModuleBTF.spec = make(map[string]*Spec)
+
+	kallsyms.FlushKernelModuleCache()
 }
 
 func kernelSpec() (*Spec, bool, error) {
@@ -759,46 +762,54 @@ func (iter *TypesIterator) Next() bool {
 	return true
 }
 
-type mergedSpec struct {
-	base *Spec
-	kmod *Spec
-}
+type mergedSpec []*Spec
 
-func (s *mergedSpec) TypeID(typ Type) (TypeID, error) {
-	id, err := s.base.TypeID(typ)
-	if err != nil && errors.Is(err, ErrNotFound) && s.kmod != nil {
-		return s.kmod.TypeID(typ)
+func (s mergedSpec) TypeID(typ Type) (TypeID, error) {
+	for _, sp := range s {
+		id, err := sp.TypeID(typ)
+		if err != nil {
+			if errors.Is(err, ErrNotFound) {
+				continue
+			}
+			return 0, err
+		}
+		return id, nil
 	}
-	return id, err
+	return 0, ErrNotFound
 }
 
-func (s *mergedSpec) NamedTypesIterate(name essentialName) *namedTypesIterator {
-	iter := &namedTypesIterator{baseTypes: s.base.namedTypes[name], index: 0}
-	if s.kmod != nil {
-		iter.kmodTypes = s.kmod.namedTypes[name]
+func (s mergedSpec) NamedTypesIterate(name essentialName) *namedTypesIterator {
+	iter := &namedTypesIterator{specIndex: 0, typeIndex: 0}
+	for _, sp := range s {
+		namedTypes := sp.namedTypes[name]
+		if len(namedTypes) > 0 {
+			iter.types = append(iter.types, namedTypes)
+		}
 	}
 	return iter
 }
 
 type namedTypesIterator struct {
-	baseTypes []Type
-	kmodTypes []Type
-	index     int
+	types     [][]Type
+	specIndex int
+	typeIndex int
 
 	Type Type
 }
 
 func (iter *namedTypesIterator) Next() bool {
-	if len(iter.baseTypes)+len(iter.kmodTypes) <= iter.index {
+	if len(iter.types) == 0 {
 		return false
 	}
-
-	if iter.index < len(iter.baseTypes) {
-		iter.Type = iter.baseTypes[iter.index]
-	} else {
-		iter.Type = iter.kmodTypes[iter.index-len(iter.baseTypes)]
+	for len(iter.types[iter.specIndex]) <= iter.typeIndex {
+		iter.specIndex++
+		iter.typeIndex = 0
+	}
+	if len(iter.types) <= iter.specIndex {
+		return false
 	}
 
-	iter.index++
+	iter.Type = iter.types[iter.specIndex][iter.typeIndex]
+	iter.typeIndex++
 	return true
 }

btf/core.go

@@ -169,28 +169,25 @@ func (k coreKind) String() string {
 //
 // Fixups are returned in the order of relos, e.g. fixup[i] is the solution
 // for relos[i].
-func CORERelocate(relos []*CORERelocation, target *Spec, kmodTarget *Spec, kmodName string, bo binary.ByteOrder, resolveLocalTypeID func(Type) (TypeID, error)) ([]COREFixup, error) {
-	if target == nil {
-		var err error
-		target, _, err = kernelSpec()
+func CORERelocate(relos []*CORERelocation, targets []*Spec, kmodName string, bo binary.ByteOrder, resolveLocalTypeID func(Type) (TypeID, error)) ([]COREFixup, error) {
+	if len(targets) == 0 {
+		kernelTarget, _, err := kernelSpec()
 		if err != nil {
 			return nil, fmt.Errorf("load kernel spec: %w", err)
 		}
-	}
-	if bo != target.byteOrder {
-		return nil, fmt.Errorf("can't relocate %s against %s", bo, target.byteOrder)
-	}
+		targets = append(targets, kernelTarget)
 
-	if kmodTarget == nil && kmodName != "" {
-		var err error
-		kmodTarget, err = kernelModuleSpec(kmodName)
-		if err != nil {
-			return nil, fmt.Errorf("load kernel module spec: %w", err)
+		if kmodName != "" {
+			kmodTarget, err := kernelModuleSpec(kmodName)
+			if err != nil {
+				return nil, fmt.Errorf("load kernel module spec: %w", err)
+			}
+			targets = append(targets, kmodTarget)
 		}
 	}
-	if kmodTarget != nil {
-		if bo != kmodTarget.byteOrder {
-			return nil, fmt.Errorf("can't relocate %s against %s", bo, kmodTarget.byteOrder)
+	for _, target := range targets {
+		if bo != target.byteOrder {
+			return nil, fmt.Errorf("can't relocate %s against %s", bo, target.byteOrder)
 		}
 	}
 
@@ -233,14 +230,14 @@ func CORERelocate(relos []*CORERelocation, target *Spec, kmodTarget *Spec, kmodN
 		group.indices = append(group.indices, i)
 	}
 
-	mergeTarget := &mergedSpec{base: target, kmod: kmodTarget}
+	mergeTarget := mergedSpec(targets)
 	for localType, group := range relosByType {
 		localTypeName := localType.TypeName()
 		if localTypeName == "" {
 			return nil, fmt.Errorf("relocate unnamed or anonymous type %s: %w", localType, ErrNotSupported)
 		}
 
-		fixups, err := coreCalculateFixups(group.relos, mergeTarget, mergeTarget.NamedTypesIterate(newEssentialName(localTypeName)), bo)
+		fixups, err := coreCalculateFixups(group.relos, &mergeTarget, mergeTarget.NamedTypesIterate(newEssentialName(localTypeName)), bo)
 		if err != nil {
 			return nil, fmt.Errorf("relocate %s: %w", localType, err)
 		}

btf/core_reloc_test.go

@@ -37,7 +37,7 @@ func TestCORERelocationLoad(t *testing.T) {
 				}
 
 				prog, err := ebpf.NewProgramWithOptions(progSpec, ebpf.ProgramOptions{
-					KernelTypes: spec.Types,
+					KernelTypes: []*btf.Spec{spec.Types},
 				})
 				testutils.SkipIfNotSupported(t, err)
 
@@ -89,7 +89,7 @@ func TestCORERelocationRead(t *testing.T) {
 		for _, progSpec := range spec.Programs {
 			t.Run(progSpec.Name, func(t *testing.T) {
 				prog, err := ebpf.NewProgramWithOptions(progSpec, ebpf.ProgramOptions{
-					KernelTypes: targetSpec,
+					KernelTypes: []*btf.Spec{targetSpec},
 				})
 				testutils.SkipIfNotSupported(t, err)
 				if err != nil {

btf/core_test.go

@@ -592,7 +592,7 @@ func TestCORERelocation(t *testing.T) {
 					relos = append(relos, reloInfo.relo)
 				}
 
-				fixups, err := CORERelocate(relos, spec, nil, "", spec.byteOrder, spec.TypeID)
+				fixups, err := CORERelocate(relos, []*Spec{spec}, "", spec.byteOrder, spec.TypeID)
 				if want := errs[name]; want != nil {
 					if !errors.Is(err, want) {
 						t.Fatal("Expected", want, "got", err)
@@ -744,7 +744,7 @@ func BenchmarkCORESkBuff(b *testing.B) {
 			b.ReportAllocs()
 
 			for i := 0; i < b.N; i++ {
-				_, err = CORERelocate([]*CORERelocation{relo}, spec, nil, "", spec.byteOrder, spec.TypeID)
+				_, err = CORERelocate([]*CORERelocation{relo}, []*Spec{spec}, "", spec.byteOrder, spec.TypeID)
 				if err != nil {
 					b.Fatal(err)
 				}

elf_reader_test.go

@@ -1134,7 +1134,7 @@ func TestLibBPFCompat(t *testing.T) {
 				}
 
 				opts := opts // copy
-				opts.Programs.KernelTypes = btfSpec
+				opts.Programs.KernelTypes = []*btf.Spec{btfSpec}
 				load(t, spec, opts, valid)
 			})
 		}

internal/kallsyms/kallsyms.go

@@ -0,0 +1,69 @@
+package kallsyms
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+	"os"
+	"sync"
+)
+
+var kernelModules struct {
+	sync.RWMutex
+	// function to kernel module mapping
+	kmods map[string]string
+}
+
+// KernelModule returns the kernel module, if any, a probe-able function is contained in.
+func KernelModule(fn string) (string, error) {
+	kernelModules.RLock()
+	kmods := kernelModules.kmods
+	kernelModules.RUnlock()
+
+	if kmods == nil {
+		kernelModules.Lock()
+		defer kernelModules.Unlock()
+		kmods = kernelModules.kmods
+	}
+
+	if kmods != nil {
+		return kmods[fn], nil
+	}
+
+	f, err := os.Open("/proc/kallsyms")
+	if err != nil {
+		return "", err
+	}
+	defer f.Close()
+	kmods, err = loadKernelModuleMapping(f)
+	if err != nil {
+		return "", err
+	}
+
+	kernelModules.kmods = kmods
+	return kmods[fn], nil
+}
+
+// FlushKernelModuleCache removes any cached information about function to kernel module mapping.
+func FlushKernelModuleCache() {
+	kernelModules.Lock()
+	defer kernelModules.Unlock()
+
+	kernelModules.kmods = nil
+}
+
+func loadKernelModuleMapping(f io.Reader) (map[string]string, error) {
+	mods := make(map[string]string)
+	scanner := bufio.NewScanner(f)
+	for scanner.Scan() {
+		fields := bytes.Fields(scanner.Bytes())
+		if len(fields) < 4 {
+			continue
+		}
+		mods[string(fields[2])] = string(bytes.Trim(fields[3], "[]"))
+	}
+	if scanner.Err() != nil {
+		return nil, scanner.Err()
+	}
+	return mods, nil
+}

internal/kallsyms/kallsyms_test.go

@@ -0,0 +1,34 @@
+package kallsyms
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/go-quicktest/qt"
+)
+
+func TestKernelModule(t *testing.T) {
+	kallsyms := []byte(`0000000000000000 t hid_generic_probe	[hid_generic]
+0000000000000000 T tcp_connect`)
+	krdr := bytes.NewBuffer(kallsyms)
+	kmods, err := loadKernelModuleMapping(krdr)
+	qt.Assert(t, qt.IsNil(err))
+
+	// present and in module
+	kmod := kmods["hid_generic_probe"]
+	if kmod != "hid_generic" {
+		t.Errorf("expected %q got %q", "hid_generic", kmod)
+	}
+
+	// present but not kernel module
+	kmod = kmods["tcp_connect"]
+	if kmod != "" {
+		t.Errorf("expected %q got %q", "", kmod)
+	}
+
+	// not present
+	kmod = kmods["asdfasdf"]
+	if kmod != "" {
+		t.Errorf("expected %q got %q", "", kmod)
+	}
+}

linker.go

@@ -120,7 +120,7 @@ func hasFunctionReferences(insns asm.Instructions) bool {
 //
 // Passing a nil target will relocate against the running kernel. insns are
 // modified in place.
-func applyRelocations(insns asm.Instructions, target *btf.Spec, kmodTarget *btf.Spec, kmodName string, bo binary.ByteOrder, b *btf.Builder) error {
+func applyRelocations(insns asm.Instructions, targets []*btf.Spec, kmodName string, bo binary.ByteOrder, b *btf.Builder) error {
 	var relos []*btf.CORERelocation
 	var reloInsns []*asm.Instruction
 	iter := insns.Iterate()
@@ -139,7 +139,7 @@ func applyRelocations(insns asm.Instructions, target *btf.Spec, kmodTarget *btf.
 		bo = internal.NativeEndian
 	}
 
-	fixups, err := btf.CORERelocate(relos, target, kmodTarget, kmodName, bo, b.Add)
+	fixups, err := btf.CORERelocate(relos, targets, kmodName, bo, b.Add)
 	if err != nil {
 		return err
 	}