support for kmod CO-RE

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

btf/btf.go

@@ -9,10 +9,12 @@ import (
 	"io"
 	"math"
 	"os"
+	"path/filepath"
 	"reflect"
 	"sync"
 
 	"github.com/cilium/ebpf/internal"
+	"github.com/cilium/ebpf/internal/kallsyms"
 	"github.com/cilium/ebpf/internal/sys"
 )
 
@@ -171,6 +173,10 @@ type Spec struct {
 	strings *stringTable
 }
 
+func init() {
+	kernelModuleBTF.spec = make(map[string]*Spec)
+}
+
 // LoadSpec opens file and calls LoadSpecFromReader on it.
 func LoadSpec(file string) (*Spec, error) {
 	fh, err := os.Open(file)
@@ -388,19 +394,45 @@ func LoadKernelSpec() (*Spec, error) {
 	return spec.Copy(), nil
 }
 
+// LoadKernelModuleSpec returns the BTF information for the named kernel module.
+//
+// Defaults to /sys/kernel/btf/<module>.
+// Returns an error wrapping ErrNotSupported if BTF is not enabled.
+func LoadKernelModuleSpec(module string) (*Spec, error) {
+	dir, file := filepath.Split(module)
+	if dir != "" || filepath.Ext(file) != "" {
+		return nil, fmt.Errorf("invalid module name %q", module)
+	}
+	spec, err := kernelModuleSpec(module)
+	if err != nil {
+		return nil, err
+	}
+	return spec.Copy(), nil
+}
+
 var kernelBTF struct {
 	sync.RWMutex
 	spec *Spec
 	// True if the spec was read from an ELF instead of raw BTF in /sys.
 	fallback bool
 }
 
+var kernelModuleBTF struct {
+	sync.RWMutex
+	spec map[string]*Spec
+}
+
 // FlushKernelSpec removes any cached kernel type information.
 func FlushKernelSpec() {
+	kernelModuleBTF.Lock()
+	defer kernelModuleBTF.Unlock()
 	kernelBTF.Lock()
 	defer kernelBTF.Unlock()
 
 	kernelBTF.spec, kernelBTF.fallback = nil, false
+	kernelModuleBTF.spec = make(map[string]*Spec)
+
+	kallsyms.FlushKernelModuleCache()
 }
 
 func kernelSpec() (*Spec, bool, error) {
@@ -428,6 +460,31 @@ func kernelSpec() (*Spec, bool, error) {
 	return spec, fallback, nil
 }
 
+func kernelModuleSpec(module string) (*Spec, error) {
+	kernelModuleBTF.RLock()
+	spec := kernelModuleBTF.spec[module]
+	kernelModuleBTF.RUnlock()
+
+	if spec == nil {
+		kernelModuleBTF.Lock()
+		defer kernelModuleBTF.Unlock()
+
+		spec = kernelModuleBTF.spec[module]
+	}
+
+	if spec != nil {
+		return spec, nil
+	}
+
+	spec, err := loadKernelModuleSpec(module)
+	if err != nil {
+		return nil, err
+	}
+
+	kernelModuleBTF.spec[module] = spec
+	return spec, nil
+}
+
 func loadKernelSpec() (_ *Spec, fallback bool, _ error) {
 	fh, err := os.Open("/sys/kernel/btf/vmlinux")
 	if err == nil {
@@ -447,6 +504,21 @@ func loadKernelSpec() (_ *Spec, fallback bool, _ error) {
 	return spec, true, err
 }
 
+func loadKernelModuleSpec(module string) (*Spec, error) {
+	base, _, err := kernelSpec()
+	if err != nil {
+		return nil, err
+	}
+
+	fh, err := os.Open(filepath.Join("/sys/kernel/btf", module))
+	if err != nil {
+		return nil, err
+	}
+	defer fh.Close()
+
+	return loadRawSpec(fh, internal.NativeEndian, base)
+}
+
 // findVMLinux scans multiple well-known paths for vmlinux kernel images.
 func findVMLinux() (*os.File, error) {
 	release, err := internal.KernelRelease()

btf/core.go

@@ -156,6 +156,33 @@ func (k coreKind) String() string {
 	}
 }
 
+type mergedSpec []*Spec
+
+func (s mergedSpec) TypeByID(id TypeID) (Type, error) {
+	for _, sp := range s {
+		t, err := sp.TypeByID(id)
+		if err != nil {
+			if errors.Is(err, ErrNotFound) {
+				continue
+			}
+			return nil, err
+		}
+		return t, nil
+	}
+	return nil, fmt.Errorf("look up type with ID %d (first ID is %d): %w", id, s[0].imm.firstTypeID, ErrNotFound)
+}
+
+func (s mergedSpec) NamedTypes(name essentialName) []TypeID {
+	var typeIDs []TypeID
+	for _, sp := range s {
+		namedTypes := sp.imm.namedTypes[name]
+		if len(namedTypes) > 0 {
+			typeIDs = append(typeIDs, namedTypes...)
+		}
+	}
+	return typeIDs
+}
+
 // CORERelocate calculates changes needed to adjust eBPF instructions for differences
 // in types.
 //
@@ -169,17 +196,26 @@ 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, 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)
 		}
-	}
+		targets = append(targets, kernelTarget)
 
-	if bo != target.imm.byteOrder {
-		return nil, fmt.Errorf("can't relocate %s against %s", bo, target.imm.byteOrder)
+		if kmodName != "" {
+			kmodTarget, err := kernelModuleSpec(kmodName)
+			if err != nil {
+				return nil, fmt.Errorf("load kernel module spec: %w", err)
+			}
+			targets = append(targets, kmodTarget)
+		}
+	}
+	for _, target := range targets {
+		if bo != target.imm.byteOrder {
+			return nil, fmt.Errorf("can't relocate %s against %s", bo, target.imm.byteOrder)
+		}
 	}
 
 	type reloGroup struct {
@@ -221,14 +257,15 @@ func CORERelocate(relos []*CORERelocation, target *Spec, bo binary.ByteOrder, re
 		group.indices = append(group.indices, i)
 	}
 
+	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)
 		}
 
-		targets := target.imm.namedTypes[newEssentialName(localTypeName)]
-		fixups, err := coreCalculateFixups(group.relos, target, targets, bo)
+		targets := mergeTarget.NamedTypes(newEssentialName(localTypeName))
+		fixups, err := coreCalculateFixups(group.relos, &mergeTarget, targets, bo)
 		if err != nil {
 			return nil, fmt.Errorf("relocate %s: %w", localType, err)
 		}
@@ -251,7 +288,7 @@ var errIncompatibleTypes = errors.New("incompatible types")
 //
 // The best target is determined by scoring: the less poisoning we have to do
 // the better the target is.
-func coreCalculateFixups(relos []*CORERelocation, targetSpec *Spec, targets []TypeID, bo binary.ByteOrder) ([]COREFixup, error) {
+func coreCalculateFixups(relos []*CORERelocation, targetSpec *mergedSpec, targets []TypeID, bo binary.ByteOrder) ([]COREFixup, error) {
 	bestScore := len(relos)
 	var bestFixups []COREFixup
 	for _, targetID := range targets {

btf/core_test.go

@@ -592,7 +592,7 @@ func TestCORERelocation(t *testing.T) {
 					relos = append(relos, reloInfo.relo)
 				}
 
-				fixups, err := CORERelocate(relos, spec, spec.imm.byteOrder, spec.TypeID)
+				fixups, err := CORERelocate(relos, []*Spec{spec}, "", spec.imm.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, spec.imm.byteOrder, spec.TypeID)
+				_, err = CORERelocate([]*CORERelocation{relo}, []*Spec{spec}, "", spec.imm.byteOrder, spec.TypeID)
 				if err != nil {
 					b.Fatal(err)
 				}

internal/kallsyms/kallsyms.go

@@ -0,0 +1,74 @@
+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
+		}
+		switch string(fields[1]) {
+		case "t", "T":
+			mods[string(fields[2])] = string(bytes.Trim(fields[3], "[]"))
+		default:
+			continue
+		}
+	}
+	if scanner.Err() != nil {
+		return nil, scanner.Err()
+	}
+	return mods, nil
+}

internal/kallsyms/kallsyms_test.go

@@ -0,0 +1,43 @@
+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
+0000000000000000 B empty_zero_page
+0000000000000000 D kimage_vaddr
+0000000000000000 R __start_pci_fixups_early
+0000000000000000 V hv_root_partition
+0000000000000000 W calibrate_delay_is_known
+0000000000000000 a nft_counter_seq	[nft_counter]
+0000000000000000 b bootconfig_found
+0000000000000000 d __func__.10
+0000000000000000 r __ksymtab_LZ4_decompress_fast`)
+	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)
+	}
+}

kallsyms.go

@@ -0,0 +1,8 @@
+package ebpf
+
+import "github.com/cilium/ebpf/internal/kallsyms"
+
+// FlushKernelModuleCache removes any cached information about function to kernel module mapping.
+func FlushKernelModuleCache() {
+	kallsyms.FlushKernelModuleCache()
+}

linker.go

@@ -119,7 +119,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, 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()
@@ -138,7 +138,7 @@ func applyRelocations(insns asm.Instructions, target *btf.Spec, bo binary.ByteOr
 		bo = internal.NativeEndian
 	}
 
-	fixups, err := btf.CORERelocate(relos, target, bo, b.Add)
+	fixups, err := btf.CORERelocate(relos, targets, kmodName, bo, b.Add)
 	if err != nil {
 		return err
 	}