From 597025ae285b2c07f2db5eb7c09f0eb85f3d8bf6 Mon Sep 17 00:00:00 2001
From: Jiri Olsa <jolsa@kernel.org>
Date: Mon, 11 Sep 2023 08:04:43 +0000
Subject: [PATCH] tetragon: Add support for fmod_ret override

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 bpf/process/bpf_generic_kprobe.c     | 14 ++++++
 pkg/sensors/program/loader.go        | 67 +++++++++++++++++++++++++++-
 pkg/sensors/program/program.go       |  3 +-
 pkg/sensors/tracing/generickprobe.go |  4 ++
 4 files changed, 85 insertions(+), 3 deletions(-)

diff --git a/bpf/process/bpf_generic_kprobe.c b/bpf/process/bpf_generic_kprobe.c
index 49350f84a9d..26fb09acacf 100644
--- a/bpf/process/bpf_generic_kprobe.c
+++ b/bpf/process/bpf_generic_kprobe.c
@@ -302,3 +302,17 @@ generic_kprobe_override(void *ctx)
 	map_delete_elem(&override_tasks, &id);
 	return 0;
 }
+
+__attribute__((section("fmod_ret/security_task_prctl"), used)) int
+generic_fmodret_override(void *ctx)
+{
+	__u64 id = get_current_pid_tgid();
+	__s32 *error;
+
+	error = map_lookup_elem(&override_tasks, &id);
+	if (!error)
+		return 0;
+
+	map_delete_elem(&override_tasks, &id);
+	return *error;
+}
diff --git a/pkg/sensors/program/loader.go b/pkg/sensors/program/loader.go
index 869aaf385c8..b183b4b3db9 100644
--- a/pkg/sensors/program/loader.go
+++ b/pkg/sensors/program/loader.go
@@ -138,6 +138,18 @@ func KprobeOpen(load *Program) OpenFunc {
 		// loaded and bpftool will show it.
 		if !load.Override {
 			disableProg(coll, "generic_kprobe_override")
+			disableProg(coll, "generic_fmodret_override")
+		} else {
+			if load.OverrideFmodRet {
+				spec, ok := coll.Programs["generic_fmodret_override"]
+				if !ok {
+					return fmt.Errorf("failed to find generic_fmodret_override")
+				}
+				spec.AttachTo = load.Attach
+				disableProg(coll, "generic_kprobe_override")
+			} else {
+				disableProg(coll, "generic_fmodret_override")
+			}
 		}
 		return nil
 	}
@@ -196,13 +208,64 @@ func kprobeAttachOverride(load *Program, bpfDir string,
 	return nil
 }
 
+func fmodretAttachOverride(load *Program, bpfDir string,
+	coll *ebpf.Collection, collSpec *ebpf.CollectionSpec) error {
+
+	spec, ok := collSpec.Programs["generic_fmodret_override"]
+	if !ok {
+		return fmt.Errorf("spec for generic_fmodret_override program not found")
+	}
+
+	prog, ok := coll.Programs["generic_fmodret_override"]
+	if !ok {
+		return fmt.Errorf("program generic_fmodret_override not found")
+	}
+
+	prog, err := prog.Clone()
+	if err != nil {
+		return fmt.Errorf("failed to clone generic_fmodret_override program: %w", err)
+	}
+
+	pinPath := filepath.Join(bpfDir, fmt.Sprint(load.PinPath, "-override"))
+
+	if err := prog.Pin(pinPath); err != nil {
+		return fmt.Errorf("pinning '%s' to '%s' failed: %w", load.Label, pinPath, err)
+	}
+
+	linkFn := func() (link.Link, error) {
+		return link.AttachTracing(link.TracingOptions{
+			Program: prog,
+		})
+	}
+
+	lnk, err := linkFn()
+	if err != nil {
+		return fmt.Errorf("attaching '%s' failed: %w", spec.Name, err)
+	}
+
+	load.unloaderOverride = &unloader.RelinkUnloader{
+		UnloadProg: unloader.PinUnloader{Prog: prog}.Unload,
+		IsLinked:   true,
+		Link:       lnk,
+		RelinkFn:   linkFn,
+	}
+
+	return nil
+}
+
 func KprobeAttach(load *Program, bpfDir string) AttachFunc {
 	return func(coll *ebpf.Collection, collSpec *ebpf.CollectionSpec,
 		prog *ebpf.Program, spec *ebpf.ProgramSpec) (unloader.Unloader, error) {
 
 		if load.Override {
-			if err := kprobeAttachOverride(load, bpfDir, coll, collSpec); err != nil {
-				return nil, err
+			if load.OverrideFmodRet {
+				if err := fmodretAttachOverride(load, bpfDir, coll, collSpec); err != nil {
+					return nil, err
+				}
+			} else {
+				if err := kprobeAttachOverride(load, bpfDir, coll, collSpec); err != nil {
+					return nil, err
+				}
 			}
 		}
 
diff --git a/pkg/sensors/program/program.go b/pkg/sensors/program/program.go
index 4f739a32c74..2d0f786caee 100644
--- a/pkg/sensors/program/program.go
+++ b/pkg/sensors/program/program.go
@@ -74,7 +74,8 @@ type Program struct {
 	ErrorFatal bool
 
 	// Needs override bpf program
-	Override bool
+	Override        bool
+	OverrideFmodRet bool
 
 	// Type is the type of BPF program. For example, tc, skb, tracepoint,
 	// etc.
diff --git a/pkg/sensors/tracing/generickprobe.go b/pkg/sensors/tracing/generickprobe.go
index 2c51b5663bb..bc5635a5ecc 100644
--- a/pkg/sensors/tracing/generickprobe.go
+++ b/pkg/sensors/tracing/generickprobe.go
@@ -727,6 +727,9 @@ func addKprobe(funcName string, f *v1alpha1.KProbeSpec, in *addKprobeIn) (out *a
 		"generic_kprobe").
 		SetLoaderData(kprobeEntry.tableId)
 	load.Override = kprobeEntry.hasOverride
+	if load.Override {
+		load.OverrideFmodRet = bpf.HasModifyReturn()
+	}
 	out.progs = append(out.progs, load)
 
 	fdinstall := program.MapBuilderPin("fdinstall_map", sensors.PathJoin(in.sensorPath, "fdinstall_map"), load)
@@ -919,6 +922,7 @@ func loadMultiKprobeSensor(ids []idtable.EntryID, bpfDir, mapDir string, load *p
 	}
 
 	load.Override = len(data.Overrides) > 0
+	load.OverrideFmodRet = false
 	load.SetAttachData(data)
 
 	if err := program.LoadMultiKprobeProgram(bpfDir, mapDir, load, verbose); err == nil {