link: attach multi-kprobes to entry and return of kernel symbols (ses…

…sion)

Adding support to attach kprobe multi link as session [1] by adding
Session bool to KprobeMultiOptions. When set true it attaches the
link with BPF_TRACE_KPROBE_SESSION attach_type, which means:

  - program is attached to both function entry and return
  - entry program can decided if the return program gets executed
  - entry program can share u64 cookie value with return program

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

Signed-off-by: Jiri Olsa <[email protected]>
Co-authored-by: Timo Beckers <[email protected]>

.github/workflows/ci.yml

@@ -155,7 +155,7 @@ jobs:
     needs: build-and-lint
     timeout-minutes: 15
     env:
-      EBPF_TEST_IGNORE_VERSION: 'TestKprobeMulti,TestKprobeMultiErrors,TestKprobeMultiCookie,TestKprobeMultiProgramCall,TestHaveBPFLinkKprobeMulti,TestHaveProgramType/LircMode2'
+      EBPF_TEST_IGNORE_VERSION: 'TestKprobeMulti,TestKprobeMultiErrors,TestKprobeMultiCookie,TestKprobeMultiProgramCall,TestHaveBPFLinkKprobeMulti,TestKprobeSession,TestHaveBPFLinkKprobeSession,TestHaveProgramType/LircMode2'
     steps:
       - uses: actions/checkout@v4
 

link/kprobe_multi.go

@@ -37,6 +37,14 @@ type KprobeMultiOptions struct {
 	// Each Cookie is assigned to the Symbol or Address specified at the
 	// corresponding slice index.
 	Cookies []uint64
+
+	// Session must be true when attaching Programs with the
+	// [ebpf.AttachTraceKprobeSession] attach type.
+	//
+	// This makes a Kprobe execute on both function entry and return. The entry
+	// program can share a cookie value with the return program and can decide
+	// whether the return program gets executed.
+	Session bool
 }
 
 // KprobeMulti attaches the given eBPF program to the entry point of a given set
@@ -82,9 +90,14 @@ func kprobeMulti(prog *ebpf.Program, opts KprobeMultiOptions, flags uint32) (Lin
 		return nil, fmt.Errorf("Cookies must be exactly Symbols or Addresses in length: %w", errInvalidInput)
 	}
 
+	attachType := sys.BPF_TRACE_KPROBE_MULTI
+	if opts.Session {
+		attachType = sys.BPF_TRACE_KPROBE_SESSION
+	}
+
 	attr := &sys.LinkCreateKprobeMultiAttr{
 		ProgFd:           uint32(prog.FD()),
-		AttachType:       sys.BPF_TRACE_KPROBE_MULTI,
+		AttachType:       attachType,
 		KprobeMultiFlags: flags,
 	}
 
@@ -103,21 +116,31 @@ func kprobeMulti(prog *ebpf.Program, opts KprobeMultiOptions, flags uint32) (Lin
 	}
 
 	fd, err := sys.LinkCreateKprobeMulti(attr)
+	if err == nil {
+		return &kprobeMultiLink{RawLink{fd, ""}}, nil
+	}
+
 	if errors.Is(err, unix.ESRCH) {
 		return nil, fmt.Errorf("couldn't find one or more symbols: %w", os.ErrNotExist)
 	}
-	if errors.Is(err, unix.EINVAL) {
-		return nil, fmt.Errorf("%w (missing kernel symbol or prog's AttachType not AttachTraceKprobeMulti?)", err)
-	}
 
-	if err != nil {
+	if opts.Session {
+		if haveFeatErr := haveBPFLinkKprobeSession(); haveFeatErr != nil {
+			return nil, haveFeatErr
+		}
+	} else {
 		if haveFeatErr := haveBPFLinkKprobeMulti(); haveFeatErr != nil {
 			return nil, haveFeatErr
 		}
-		return nil, err
 	}
 
-	return &kprobeMultiLink{RawLink{fd, ""}}, nil
+	// Check EINVAL after running feature probes, since it's also returned when
+	// the kernel doesn't support the multi/session attach types.
+	if errors.Is(err, unix.EINVAL) {
+		return nil, fmt.Errorf("%w (missing kernel symbol or prog's AttachType not %s?)", err, ebpf.AttachType(attachType))
+	}
+
+	return nil, err
 }
 
 type kprobeMultiLink struct {
@@ -189,3 +212,44 @@ var haveBPFLinkKprobeMulti = internal.NewFeatureTest("bpf_link_kprobe_multi", fu
 
 	return nil
 }, "5.18")
+
+var haveBPFLinkKprobeSession = internal.NewFeatureTest("bpf_link_kprobe_session", func() error {
+	prog, err := ebpf.NewProgram(&ebpf.ProgramSpec{
+		Name: "probe_kps_link",
+		Type: ebpf.Kprobe,
+		Instructions: asm.Instructions{
+			asm.Mov.Imm(asm.R0, 0),
+			asm.Return(),
+		},
+		AttachType: ebpf.AttachTraceKprobeSession,
+		License:    "MIT",
+	})
+	if errors.Is(err, unix.E2BIG) {
+		// Kernel doesn't support AttachType field.
+		return internal.ErrNotSupported
+	}
+	if err != nil {
+		return err
+	}
+	defer prog.Close()
+
+	fd, err := sys.LinkCreateKprobeMulti(&sys.LinkCreateKprobeMultiAttr{
+		ProgFd:     uint32(prog.FD()),
+		AttachType: sys.BPF_TRACE_KPROBE_SESSION,
+		Count:      1,
+		Syms:       sys.NewStringSlicePointer([]string{"vprintk"}),
+	})
+	switch {
+	case errors.Is(err, unix.EINVAL):
+		return internal.ErrNotSupported
+	// If CONFIG_FPROBE isn't set.
+	case errors.Is(err, unix.EOPNOTSUPP):
+		return internal.ErrNotSupported
+	case err != nil:
+		return err
+	}
+
+	fd.Close()
+
+	return nil
+}, "6.10")

link/kprobe_multi_test.go

@@ -5,6 +5,8 @@ import (
 	"os"
 	"testing"
 
+	"github.com/go-quicktest/qt"
+
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/internal/testutils"
 	"github.com/cilium/ebpf/internal/unix"
@@ -134,3 +136,20 @@ func TestKprobeMultiProgramCall(t *testing.T) {
 func TestHaveBPFLinkKprobeMulti(t *testing.T) {
 	testutils.CheckFeatureTest(t, haveBPFLinkKprobeMulti)
 }
+
+func TestKprobeSession(t *testing.T) {
+	testutils.SkipIfNotSupported(t, haveBPFLinkKprobeMulti())
+
+	prog := mustLoadProgram(t, ebpf.Kprobe, ebpf.AttachTraceKprobeSession, "")
+
+	km, err := KprobeMulti(prog, KprobeMultiOptions{Symbols: kprobeMultiSyms, Session: true})
+	testutils.SkipIfNotSupported(t, err)
+	qt.Assert(t, qt.IsNil(err))
+	defer km.Close()
+
+	testLink(t, km, prog)
+}
+
+func TestHaveBPFLinkKprobeSession(t *testing.T) {
+	testutils.CheckFeatureTest(t, haveBPFLinkKprobeSession)
+}

types.go

@@ -236,6 +236,7 @@ const (
 	AttachSkReuseportSelectOrMigrate = AttachType(sys.BPF_SK_REUSEPORT_SELECT_OR_MIGRATE)
 	AttachPerfEvent                  = AttachType(sys.BPF_PERF_EVENT)
 	AttachTraceKprobeMulti           = AttachType(sys.BPF_TRACE_KPROBE_MULTI)
+	AttachTraceKprobeSession         = AttachType(sys.BPF_TRACE_KPROBE_SESSION)
 	AttachLSMCgroup                  = AttachType(sys.BPF_LSM_CGROUP)
 	AttachStructOps                  = AttachType(sys.BPF_STRUCT_OPS)
 	AttachNetfilter                  = AttachType(sys.BPF_NETFILTER)