tetragon: Add support for multi kprobe override

bpf/process/bpf_generic_kprobe.c

@@ -117,9 +117,11 @@ generic_kprobe_start_process_filter(void *ctx)
 }
 
 #ifdef __MULTI_KPROBE
-#define MAIN "kprobe.multi/generic_kprobe"
+#define MAIN	 "kprobe.multi/generic_kprobe"
+#define OVERRIDE "kprobe.multi/generic_kprobe_override"
 #else
-#define MAIN "kprobe/generic_kprobe"
+#define MAIN	 "kprobe/generic_kprobe"
+#define OVERRIDE "kprobe/generic_kprobe_override"
 #endif
 
 /* Generic kprobe pseudocode is the following
@@ -286,7 +288,7 @@ generic_kprobe_output(void *ctx)
 	return generic_output(ctx, (struct bpf_map_def *)&process_call_heap);
 }
 
-__attribute__((section("kprobe/override"), used)) int
+__attribute__((section(OVERRIDE), used)) int
 generic_kprobe_override(void *ctx)
 {
 	__u64 id = get_current_pid_tgid();

pkg/sensors/program/loader.go

@@ -276,38 +276,77 @@ func LSMAttach() AttachFunc {
 	}
 }
 
-func MultiKprobeAttach(load *Program) AttachFunc {
+func multiKprobeAttach(load *Program, prog *ebpf.Program,
+	spec *ebpf.ProgramSpec, opts link.KprobeMultiOptions) (unloader.Unloader, error) {
+
+	var lnk link.Link
+	var err error
+
+	if load.RetProbe {
+		lnk, err = link.KretprobeMulti(prog, opts)
+	} else {
+		lnk, err = link.KprobeMulti(prog, opts)
+	}
+	if err != nil {
+		return nil, fmt.Errorf("attaching '%s' failed: %w", spec.Name, err)
+	}
+	return unloader.ChainUnloader{
+		unloader.PinUnloader{
+			Prog: prog,
+		},
+		unloader.LinkUnloader{
+			Link: lnk,
+		},
+	}, nil
+}
+
+func MultiKprobeAttach(load *Program, bpfDir string) AttachFunc {
 	return func(coll *ebpf.Collection, collSpec *ebpf.CollectionSpec,
 		prog *ebpf.Program, spec *ebpf.ProgramSpec) (unloader.Unloader, error) {
 
 		data, ok := load.AttachData.(*MultiKprobeAttachData)
 		if !ok {
 			return nil, fmt.Errorf("attaching '%s' failed: wrong attach data", spec.Name)
 		}
+
+		if load.Override {
+			progOverrideSpec, ok := collSpec.Programs["generic_kprobe_override"]
+			if ok {
+				progOverrideSpec.Type = ebpf.UnspecifiedProgram
+			}
+
+			progOverride, ok := coll.Programs["generic_kprobe_override"]
+			if !ok {
+				return nil, fmt.Errorf("program for section '%s' not found", load.Label)
+			}
+
+			progOverride, err := progOverride.Clone()
+			if err != nil {
+				return nil, fmt.Errorf("failed to clone program '%s': %w", load.Label, err)
+			}
+
+			pinPath := filepath.Join(bpfDir, fmt.Sprint(load.PinPath, "-override"))
+
+			if err := progOverride.Pin(pinPath); err != nil {
+				return nil, fmt.Errorf("pinning '%s' to '%s' failed: %w", load.Label, pinPath, err)
+			}
+
+			opts := link.KprobeMultiOptions{
+				Symbols: data.Overrides,
+			}
+
+			load.unloaderOverride, err = multiKprobeAttach(load, progOverride, progOverrideSpec, opts)
+			if err != nil {
+				logger.GetLogger().Warnf("Failed to attach override program: %w", err)
+			}
+		}
+
 		opts := link.KprobeMultiOptions{
 			Symbols: data.Symbols,
 			Cookies: data.Cookies,
 		}
 
-		var lnk link.Link
-		var err error
-
-		if load.RetProbe {
-			lnk, err = link.KretprobeMulti(prog, opts)
-		} else {
-			lnk, err = link.KprobeMulti(prog, opts)
-		}
-		if err != nil {
-			return nil, fmt.Errorf("attaching '%s' failed: %w", spec.Name, err)
-		}
-		return unloader.ChainUnloader{
-			unloader.PinUnloader{
-				Prog: prog,
-			},
-			unloader.LinkUnloader{
-				Link: lnk,
-			},
-		}, nil
+		return multiKprobeAttach(load, prog, spec, opts)
 	}
 }
 
@@ -366,7 +405,8 @@ func LoadUprobeProgram(bpfDir, mapDir string, load *Program, verbose int) error
 
 func LoadMultiKprobeProgram(bpfDir, mapDir string, load *Program, verbose int) error {
 	opts := &loadOpts{
-		attach: MultiKprobeAttach(load),
+		attach: MultiKprobeAttach(load, bpfDir),
+		open:   KprobeOpen(load),
 		ci:     &customInstall{fmt.Sprintf("%s-kp_calls", load.PinPath), "kprobe"},
 	}
 	return loadProgram(bpfDir, []string{mapDir}, load, opts, verbose)

pkg/sensors/program/program.go

@@ -43,8 +43,9 @@ type MapLoad struct {
 }
 
 type MultiKprobeAttachData struct {
-	Symbols []string
-	Cookies []uint64
+	Symbols   []string
+	Cookies   []uint64
+	Overrides []string
 }
 
 type UprobeAttachData struct {

pkg/sensors/tracing/generickprobe.go

@@ -118,6 +118,9 @@ type genericKprobe struct {
 
 	// policyName is the name of the policy that this tracepoint belongs to
 	policyName string
+
+	// is there override defined for the kprobe
+	hasOverride bool
 }
 
 // pendingEvent is an event waiting to be merged with another event.
@@ -364,10 +367,8 @@ func createGenericKprobeSensor(
 	// use multi kprobe only if:
 	// - it's not disabled by user
 	// - there's support detected
-	// - multiple kprobes are defined
 	useMulti = !option.Config.DisableKprobeMulti &&
-		bpf.HasKprobeMulti() &&
-		len(kprobes) > 1
+		bpf.HasKprobeMulti()
 
 	for i := range kprobes {
 		f := &kprobes[i]
@@ -497,8 +498,6 @@ func createGenericKprobeSensor(
 			}
 		}
 
-		hasOverride := selectors.HasOverride(f)
-
 		// Write attributes into BTF ptr for use with load
 		if !setRetprobe {
 			setRetprobe = f.Return
@@ -529,6 +528,7 @@ func createGenericKprobeSensor(
 			pendingEvents:     nil,
 			tableId:           idtable.UninitializedEntryID,
 			policyName:        policyName,
+			hasOverride:       selectors.HasOverride(f),
 		}
 
 		// Parse Filters into kernel filter logic
@@ -555,6 +555,7 @@ func createGenericKprobeSensor(
 			logger.GetLogger().
 				WithField("return", setRetprobe).
 				WithField("function", kprobeEntry.funcName).
+				WithField("override", kprobeEntry.hasOverride).
 				Infof("Added multi kprobe")
 			continue
 		}
@@ -569,7 +570,7 @@ func createGenericKprobeSensor(
 			pinProg,
 			"generic_kprobe").
 			SetLoaderData(kprobeEntry.tableId)
-		load.Override = hasOverride
+		load.Override = kprobeEntry.hasOverride
 		progs = append(progs, load)
 
 		fdinstall := program.MapBuilderPin("fdinstall_map", sensors.PathJoin(sensorPath, "fdinstall_map"), load)
@@ -643,6 +644,7 @@ func createGenericKprobeSensor(
 		}
 
 		logger.GetLogger().WithField("flags", flagsString(config.Flags)).
+			WithField("override", kprobeEntry.hasOverride).
 			Infof("Added generic kprobe sensor: %s -> %s", load.Name, load.Attach)
 	}
 
@@ -728,8 +730,13 @@ func loadMultiKprobeSensor(ids []idtable.EntryID, bpfDir, mapDir string, load *p
 
 		data.Symbols = append(data.Symbols, gk.funcName)
 		data.Cookies = append(data.Cookies, uint64(index))
+
+		if gk.hasOverride && !load.RetProbe {
+			data.Overrides = append(data.Overrides, gk.funcName)
+		}
 	}
 
+	load.Override = len(data.Overrides) > 0
 	load.SetAttachData(data)
 
 	if err := program.LoadMultiKprobeProgram(bpfDir, mapDir, load, verbose); err == nil {