elf: support kernel module kfunc calls

This commit adds support for kernel module function call support.

The fdArray parameter is used during BPF program load to pass module BTFs referenced by the program.
ins.Offset is set to index into this array and ins.Constant is assigned a BTFID in the BTF (of module or vmlinux).

cf: https://lore.kernel.org/bpf/[email protected]/

Signed-off-by: shun159 <[email protected]>

Author: shun159

Makefile

@@ -42,6 +42,7 @@ TARGETS := \
 	testdata/fwd_decl \
 	testdata/kconfig \
 	testdata/kfunc \
+	testdata/kfunc-kmod \
 	btf/testdata/relocs \
 	btf/testdata/relocs_read \
 	btf/testdata/relocs_read_tgt \

elf_reader_test.go

@@ -656,7 +656,7 @@ func TestKconfigSyscallWrapper(t *testing.T) {
 
 func TestKfunc(t *testing.T) {
 	testutils.SkipOnOldKernel(t, "5.18", "bpf_kfunc_call_test_mem_len_pass1")
-	testutils.Files(t, testutils.Glob(t, "testdata/kfunc-*.elf"), func(t *testing.T, file string) {
+	testutils.Files(t, testutils.Glob(t, "testdata/kfunc-e*.elf"), func(t *testing.T, file string) {
 		spec, err := LoadCollectionSpec(file)
 		if err != nil {
 			t.Fatal(err)
@@ -686,7 +686,46 @@ func TestKfunc(t *testing.T) {
 		if ret != 1 {
 			t.Fatalf("Expected kfunc to return value 1, got %d", ret)
 		}
+	})
+}
+
+func TestKfuncKmod(t *testing.T) {
+	testutils.SkipOnOldKernel(t, "5.18", "Kernel module function calls")
+
+	if !haveTestmod(t) {
+		t.Skip("bpf_testmod not loaded")
+	}
+
+	testutils.Files(t, testutils.Glob(t, "testdata/kfunc-kmod-*.elf"), func(t *testing.T, file string) {
+		spec, err := LoadCollectionSpec(file)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if spec.ByteOrder != internal.NativeEndian {
+			return
+		}
+
+		var obj struct {
+			Main *Program `ebpf:"call_kfunc"`
+		}
 
+		err = spec.LoadAndAssign(&obj, nil)
+		testutils.SkipIfNotSupported(t, err)
+		if err != nil {
+			t.Fatalf("%v+", err)
+		}
+		defer obj.Main.Close()
+
+		ret, _, err := obj.Main.Test(internal.EmptyBPFContext)
+		testutils.SkipIfNotSupported(t, err)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if ret != 1 {
+			t.Fatalf("Expected kfunc to return value 1, got %d", ret)
+		}
 	})
 }
 

helpers_test.go

@@ -0,0 +1,26 @@
+package ebpf
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/cilium/ebpf/btf"
+	"github.com/cilium/ebpf/internal/testutils"
+)
+
+func haveTestmod(tb testing.TB) bool {
+	haveTestmod := false
+	if !testutils.IsKernelLessThan(tb, "5.11") {
+		// See https://github.com/torvalds/linux/commit/290248a5b7d829871b3ea3c62578613a580a1744
+		testmod, err := btf.FindHandle(func(info *btf.HandleInfo) bool {
+			return info.IsModule() && info.Name == "bpf_testmod"
+		})
+		if err != nil && !errors.Is(err, btf.ErrNotFound) {
+			tb.Fatal(err)
+		}
+		haveTestmod = testmod != nil
+		testmod.Close()
+	}
+
+	return haveTestmod
+}

linker.go

@@ -5,12 +5,47 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"math"
 
 	"github.com/cilium/ebpf/asm"
 	"github.com/cilium/ebpf/btf"
 	"github.com/cilium/ebpf/internal"
 )
 
+// handles stores handle objects to avoid gc cleanup
+type handles []*btf.Handle
+
+func (hs *handles) add(h *btf.Handle) (int, error) {
+	if h == nil {
+		return 0, nil
+	}
+
+	if len(*hs) == math.MaxInt16 {
+		return 0, fmt.Errorf("can't add more than %d module FDs to fdArray", math.MaxInt16)
+	}
+
+	*hs = append(*hs, h)
+
+	// return length of slice so that indexes start at 1
+	return len(*hs), nil
+}
+
+func (hs handles) fdArray() []int32 {
+	// first element of fda is reserved as no module can be indexed with 0
+	fda := []int32{0}
+	for _, h := range hs {
+		fda = append(fda, int32(h.FD()))
+	}
+
+	return fda
+}
+
+func (hs handles) close() {
+	for _, h := range hs {
+		h.Close()
+	}
+}
+
 // splitSymbols splits insns into subsections delimited by Symbol Instructions.
 // insns cannot be empty and must start with a Symbol Instruction.
 //
@@ -190,17 +225,13 @@ func fixupAndValidate(insns asm.Instructions) error {
 		fixupProbeReadKernel(ins)
 	}
 
-	if err := fixupKfuncs(insns); err != nil {
-		return fmt.Errorf("fixing up kfuncs: %w", err)
-	}
-
 	return nil
 }
 
 // fixupKfuncs loops over all instructions in search for kfunc calls.
-// If at least one is found, the current kernels BTF is loaded to set Instruction.Constant
-// to the running kernels btf id of the btf.Func in the instructions Metadata for all kfunc call instructions.
-func fixupKfuncs(insns asm.Instructions) error {
+// If at least one is found, the current kernels BTF and module BTFis are searched to set Instruction.Constant
+// and Instruction.Offset to the correct values.
+func fixupKfuncs(insns asm.Instructions) (handles, error) {
 	iter := insns.Iterate()
 	for iter.Next() {
 		ins := iter.Ins
@@ -209,15 +240,16 @@ func fixupKfuncs(insns asm.Instructions) error {
 		}
 	}
 
-	return nil
+	return nil, nil
 
 fixups:
 	// only load the kernel spec if we found at least one kfunc call
-	s, err := btf.LoadKernelSpec()
+	kernelSpec, err := btf.LoadKernelSpec()
 	if err != nil {
-		return err
+		return nil, err
 	}
 
+	fdArray := make(handles, 0)
 	for {
 		ins := iter.Ins
 
@@ -232,32 +264,45 @@ fixups:
 		// check meta, if no meta return err
 		kfm, _ := ins.Metadata.Get(kfuncMeta{}).(*btf.Func)
 		if kfm == nil {
-			return fmt.Errorf("kfunc call has no kfuncMeta")
+			return nil, fmt.Errorf("kfunc call has no kfuncMeta")
 		}
 
-		var fn *btf.Func
-		if err := s.TypeByName(kfm.Name, &fn); err != nil {
-			return fmt.Errorf("couldn't resolve %s in kernel spec: %v: %w", kfm.Name, err, ErrNotSupported)
+		kfuncHandle, id, err := findKfuncInKernel(kernelSpec, kfm.Name)
+		if err != nil {
+			return nil, err
 		}
 
-		if err := btf.CheckTypeCompatibility(kfm.Type, fn.Type); err != nil {
-			return err
+		// to avoid doing compatibility checks in findKfuncInKernel
+		// we need to retrieve the btf.Type again by its ID.
+		spec := kernelSpec
+		if kfuncHandle != nil {
+			spec, err = kfuncHandle.Spec(kernelSpec)
+			if err != nil {
+				return nil, err
+			}
+		}
+		typ, err := spec.TypeByID(id)
+		if err != nil {
+			return nil, err
+		}
+		if err := btf.CheckTypeCompatibility(kfm.Type, typ.(*btf.Func).Type); err != nil {
+			return nil, err
 		}
 
-		id, err := s.TypeID(fn)
+		idx, err := fdArray.add(kfuncHandle)
 		if err != nil {
-			return err
+			return nil, err
 		}
 
 		ins.Constant = int64(id)
-		ins.Offset = int16(0) // currently always 0, no support for kmods
+		ins.Offset = int16(idx)
 
 		if !iter.Next() {
 			break
 		}
 	}
 
-	return nil
+	return fdArray, nil
 }
 
 // fixupProbeReadKernel replaces calls to bpf_probe_read_{kernel,user}(_str)

prog.go

@@ -10,6 +10,7 @@ import (
 	"runtime"
 	"strings"
 	"time"
+	"unsafe"
 
 	"github.com/cilium/ebpf/asm"
 	"github.com/cilium/ebpf/btf"
@@ -268,6 +269,17 @@ func newProgramWithOptions(spec *ProgramSpec, opts ProgramOptions) (*Program, er
 		return nil, err
 	}
 
+	handles, err := fixupKfuncs(insns)
+	if err != nil {
+		return nil, fmt.Errorf("fixing up kfuncs: %w", err)
+	}
+
+	if len(handles) > 0 {
+		fdArray := handles.fdArray()
+		attr.FdArray = sys.NewPointer(unsafe.Pointer(&fdArray[0]))
+		defer handles.close()
+	}
+
 	buf := bytes.NewBuffer(make([]byte, 0, insns.Size()))
 	err = insns.Marshal(buf, internal.NativeEndian)
 	if err != nil {
@@ -1000,3 +1012,31 @@ func findTargetInProgram(prog *Program, name string, progType ProgramType, attac
 
 	return spec.TypeID(targetFunc)
 }
+
+// find a kfunc in a kernel module or vmlinux.
+func findKfuncInKernel(kernelSpec *btf.Spec, name string) (*btf.Handle, btf.TypeID, error) {
+	var fn *btf.Func
+	err := kernelSpec.TypeByName(name, &fn)
+	if errors.Is(err, btf.ErrNotFound) {
+		// couldn't find kfunc in vmlinux, searching modules.
+		module, id, err := findTargetInModule(kernelSpec, name, fn)
+		if errors.Is(err, btf.ErrNotFound) {
+			return nil, 0, fmt.Errorf("find kfunc '%s' in vmlinux or modules: %v: %w", name, err, ErrNotSupported)
+		}
+		if err != nil {
+			return nil, 0, fmt.Errorf("find kfunc '%s' in modules: %v", name, err)
+		}
+
+		return module, id, nil
+	}
+	if err != nil {
+		return nil, 0, fmt.Errorf("find kfunc '%s' in kernel: %v", name, err)
+	}
+
+	id, err := kernelSpec.TypeID(fn)
+	if err != nil {
+		return nil, 0, err
+	}
+
+	return nil, id, nil
+}

prog_test.go

@@ -758,18 +758,7 @@ func TestProgramAttachToKernel(t *testing.T) {
 	// See https://github.com/torvalds/linux/commit/290248a5b7d829871b3ea3c62578613a580a1744
 	testutils.SkipOnOldKernel(t, "5.5", "attach_btf_id")
 
-	haveTestmod := false
-	if !testutils.IsKernelLessThan(t, "5.11") {
-		// See https://github.com/torvalds/linux/commit/290248a5b7d829871b3ea3c62578613a580a1744
-		testmod, err := btf.FindHandle(func(info *btf.HandleInfo) bool {
-			return info.IsModule() && info.Name == "bpf_testmod"
-		})
-		if err != nil && !errors.Is(err, btf.ErrNotFound) {
-			t.Fatal(err)
-		}
-		haveTestmod = testmod != nil
-		testmod.Close()
-	}
+	haveTestmod := haveTestmod(t)
 
 	tests := []struct {
 		attachTo    string

testdata/kfunc-kmod-eb.elf

@@ -0,0 +1,10 @@
+#include "common.h"
+
+char __license[] __section("license") = "Dual MIT/GPL";
+
+extern void bpf_testmod_test_mod_kfunc(int) __ksym;
+
+__section("tc") int call_kfunc() {
+	bpf_testmod_test_mod_kfunc(0);
+	return 1;
+}