link: Add support to attach kprobe multi link as session

Adding support to attach kprobe multi link as session by adding
Session bool to KprobeMultiOptions. When true it attaches the
link with BPF_TRACE_KPROBE_SESSION attach_type.

Signed-off-by: Jiri Olsa <[email protected]>

link/kprobe_multi.go

@@ -37,6 +37,9 @@ type KprobeMultiOptions struct {
 	// Each Cookie is assigned to the Symbol or Address specified at the
 	// corresponding slice index.
 	Cookies []uint64
+
+	// Probe is attach as session with BPF_TRACE_KPROBE_SESSION attach type
+	Session bool
 }
 
 // KprobeMulti attaches the given eBPF program to the entry point of a given set
@@ -82,9 +85,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,
 	}
 
@@ -189,3 +197,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.9")

link/kprobe_multi_test.go

@@ -134,3 +134,17 @@ func TestKprobeMultiProgramCall(t *testing.T) {
 func TestHaveBPFLinkKprobeMulti(t *testing.T) {
 	testutils.CheckFeatureTest(t, haveBPFLinkKprobeMulti)
 }
+
+func TestKprobeSession(t *testing.T) {
+	testutils.SkipIfNotSupported(t, haveBPFLinkKprobeSession())
+
+	prog := mustLoadProgram(t, ebpf.Kprobe, ebpf.AttachTraceKprobeSession, "")
+
+	km, err := KprobeMulti(prog, KprobeMultiOptions{Symbols: kprobeMultiSyms, Session: true})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer km.Close()
+
+	testLink(t, km, prog)
+}

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)