From 71717b1875bd23acd5b8c4290aa2ae008b87df2f Mon Sep 17 00:00:00 2001
From: Jiri Olsa <jolsa@kernel.org>
Date: Fri, 26 Jul 2024 14:32:14 +0000
Subject: [PATCH] tetragon: Add cgroup rate support for kprobe

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 bpf/process/bpf_generic_kprobe.c     | 11 +++++++++++
 pkg/sensors/base/base.go             |  7 +++++--
 pkg/sensors/tracing/generickprobe.go | 17 +++++++++++++++++
 pkg/sensors/tracing/kprobe_test.go   |  4 ++--
 4 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/bpf/process/bpf_generic_kprobe.c b/bpf/process/bpf_generic_kprobe.c
index d8743d5d97b..2bdc38d1d8b 100644
--- a/bpf/process/bpf_generic_kprobe.c
+++ b/bpf/process/bpf_generic_kprobe.c
@@ -15,6 +15,7 @@
 #include "generic_calls.h"
 #include "pfilter.h"
 #include "policy_filter.h"
+#include "bpf_rate.h"
 
 char _license[] __attribute__((section("license"), used)) = "Dual BSD/GPL";
 
@@ -112,6 +113,16 @@ static struct generic_maps maps = {
 __attribute__((section((MAIN)), used)) int
 generic_kprobe_event(struct pt_regs *ctx)
 {
+	__u64 ktime = ktime_get_ns();
+	struct task_struct *task;
+	struct msg_k8s kube;
+
+	task = (struct task_struct *)get_current_task();
+	__event_get_cgroup_info(task, &kube);
+
+	if (!cgroup_rate(ctx, &kube, ktime))
+		return 0;
+
 	return generic_start_process_filter(ctx, &maps);
 }
 
diff --git a/pkg/sensors/base/base.go b/pkg/sensors/base/base.go
index 5a1dc856b66..818a66b69a3 100644
--- a/pkg/sensors/base/base.go
+++ b/pkg/sensors/base/base.go
@@ -16,7 +16,7 @@ import (
 )
 
 const (
-	cgroupRateMaxEntries = 32768 // this value could be fine tuned
+	CgroupRateMaxEntries = 32768 // this value could be fine tuned
 )
 
 var (
@@ -80,6 +80,8 @@ var (
 	CgroupRateMap        = program.MapBuilder("cgroup_rate_map", Execve, Exit, Fork, CgroupRmdir)
 	CgroupRateOptionsMap = program.MapBuilder("cgroup_rate_options_map", Execve)
 
+	HasCgroupRate bool
+
 	sensor = sensors.Sensor{
 		Name: "__base__",
 	}
@@ -178,5 +180,6 @@ func ConfigCgroupRate(opts *option.CgroupRate) {
 		return
 	}
 
-	CgroupRateMap.SetMaxEntries(cgroupRateMaxEntries)
+	HasCgroupRate = true
+	CgroupRateMap.SetMaxEntries(CgroupRateMaxEntries)
 }
diff --git a/pkg/sensors/tracing/generickprobe.go b/pkg/sensors/tracing/generickprobe.go
index 057a79591d8..6dd519cf415 100644
--- a/pkg/sensors/tracing/generickprobe.go
+++ b/pkg/sensors/tracing/generickprobe.go
@@ -33,6 +33,7 @@ import (
 	"github.com/cilium/tetragon/pkg/policyfilter"
 	"github.com/cilium/tetragon/pkg/selectors"
 	"github.com/cilium/tetragon/pkg/sensors"
+	"github.com/cilium/tetragon/pkg/sensors/base"
 	"github.com/cilium/tetragon/pkg/sensors/program"
 	lru "github.com/hashicorp/golang-lru/v2"
 	"github.com/sirupsen/logrus"
@@ -381,6 +382,14 @@ func createMultiKprobeSensor(sensorPath, policyName string, multiIDs []idtable.E
 	}
 	maps = append(maps, overrideTasksMap)
 
+	if base.HasCgroupRate {
+		cgroupRateMap := program.MapBuilder("cgroup_rate_map", load)
+		cgroupRateOptionsMap := program.MapBuilder("cgroup_rate_options_map", load)
+
+		cgroupRateMap.SetMaxEntries(base.CgroupRateMaxEntries)
+		maps = append(maps, cgroupRateMap, cgroupRateOptionsMap)
+	}
+
 	if len(multiRetIDs) != 0 {
 		loadret := program.Builder(
 			path.Join(option.Config.HubbleLib, loadProgRetName),
@@ -976,6 +985,14 @@ func createKprobeSensorFromEntry(kprobeEntry *genericKprobe, sensorPath string,
 	}
 	maps = append(maps, overrideTasksMap)
 
+	if base.HasCgroupRate {
+		cgroupRateMap := program.MapBuilder("cgroup_rate_map", load)
+		cgroupRateOptionsMap := program.MapBuilder("cgroup_rate_options_map", load)
+
+		cgroupRateMap.SetMaxEntries(base.CgroupRateMaxEntries)
+		maps = append(maps, cgroupRateMap, cgroupRateOptionsMap)
+	}
+
 	if kprobeEntry.loadArgs.retprobe {
 		pinRetProg := sensors.PathJoin(pinPath, fmt.Sprintf("%s_ret_prog", kprobeEntry.funcName))
 		loadret := program.Builder(
diff --git a/pkg/sensors/tracing/kprobe_test.go b/pkg/sensors/tracing/kprobe_test.go
index 5ecc0d2902b..6673094c127 100644
--- a/pkg/sensors/tracing/kprobe_test.go
+++ b/pkg/sensors/tracing/kprobe_test.go
@@ -4281,7 +4281,7 @@ func TestLoadKprobeSensor(t *testing.T) {
 		sensorMaps = append(sensorMaps, tus.SensorMap{Name: "execve_map", Progs: []uint{4, 5, 6, 7, 9}})
 
 		// generic_kprobe_process_event*,generic_kprobe_output,generic_retkprobe_output
-		sensorMaps = append(sensorMaps, tus.SensorMap{Name: "tcpmon_map", Progs: []uint{1, 2, 6, 10}})
+		sensorMaps = append(sensorMaps, tus.SensorMap{Name: "tcpmon_map", Progs: []uint{0, 1, 2, 6, 10}})
 
 		// generic_kprobe_process_event*,generic_kprobe_actions,retkprobe
 		sensorMaps = append(sensorMaps, tus.SensorMap{Name: "socktrack_map", Progs: []uint{1, 2, 5, 7, 9}})
@@ -4290,7 +4290,7 @@ func TestLoadKprobeSensor(t *testing.T) {
 		sensorMaps = append(sensorMaps, tus.SensorMap{Name: "execve_map", Progs: []uint{4, 7}})
 
 		// generic_kprobe_output,generic_retkprobe_output
-		sensorMaps = append(sensorMaps, tus.SensorMap{Name: "tcpmon_map", Progs: []uint{6, 10}})
+		sensorMaps = append(sensorMaps, tus.SensorMap{Name: "tcpmon_map", Progs: []uint{0, 6, 10}})
 	}
 
 	readHook := `