tetragon: Add throttle logic to fork program

Signed-off-by: Jiri Olsa <[email protected]>
cilium · Mar 4, 2024 · a2fce80 · a2fce80
1 parent 891e246
commit a2fce80
Show file tree

Hide file tree

Showing 6 changed files with 73 additions and 7 deletions.
diff --git a/bpf/process/bpf_fork.c b/bpf/process/bpf_fork.c
@@ -10,6 +10,14 @@
 #include "bpf_task.h"
 #include "environ_conf.h"
 #include "bpf_process_event.h"
+#include "bpf_rate.h"
+
+struct {
+	__uint(type, BPF_MAP_TYPE_PROG_ARRAY);
+	__uint(max_entries, 1);
+	__type(key, __u32);
+	__type(value, __u32);
+} fork_calls SEC(".maps");
 
 char _license[] __attribute__((section("license"), used)) = "Dual BSD/GPL";
 #ifdef VMLINUX_KERNEL_VERSION
@@ -19,6 +27,14 @@ int _version __attribute__((section(("version")), used)) =
 
 __attribute__((section("kprobe/wake_up_new_task"), used)) int
 BPF_KPROBE(event_wake_up_new_task, struct task_struct *task)
+{
+	if (fork_cgroup_rate(ctx, task))
+		tail_call(ctx, &fork_calls, 0);
+	return 0;
+}
+
+__attribute__((section("kprobe/0"), used)) int
+BPF_KPROBE(fork_send, struct task_struct *task)
 {
 	struct execve_map_value *curr, *parent;
 	struct msg_clone_event msg;

diff --git a/bpf/process/bpf_process_event.h b/bpf/process/bpf_process_event.h
@@ -571,6 +571,29 @@ __event_get_current_cgroup_name(struct cgroup *cgrp,
 		process->flags |= EVENT_ERROR_CGROUP_NAME;
 }
 
+static inline __attribute__((always_inline)) __u64
+get_task_cgroupid(struct task_struct *task)
+{
+	int zero = 0, subsys_idx = 0;
+	struct tetragon_conf *conf;
+	__u64 cgrpfs_magic = 0;
+	struct cgroup *cgrp;
+	__u32 flags = 0;
+
+	conf = map_lookup_elem(&tg_conf_map, &zero);
+	if (conf) {
+		/* Select which cgroup version */
+		cgrpfs_magic = conf->cgrp_fs_magic;
+		subsys_idx = conf->tg_cgrp_subsys_idx;
+	}
+
+	cgrp = get_task_cgroup(task, subsys_idx, &flags);
+	if (!cgrp)
+		return 0;
+
+	return __tg_get_current_cgroup_id(cgrp, cgrpfs_magic);
+}
+
 /**
  * __event_get_cgroup_info() Collect cgroup info from current task.
  * @task: must be current task.

diff --git a/bpf/process/bpf_rate.h b/bpf/process/bpf_rate.h
@@ -147,4 +147,27 @@ execve_cgroup_rate(struct sched_execve_args *ctx)
 	return send;
 }
 
+static inline __attribute__((always_inline)) bool
+fork_cgroup_rate(void *ctx, struct task_struct *task)
+{
+	struct cgroup_rate_key key = { .op = MSG_OP_CLONE };
+	struct cgroup_rate_settings settings = {
+		.tokens      = 1000,
+		.interval_ns = 1 * NSEC_PER_SEC,
+		.throttle_ns = 5 * NSEC_PER_SEC,
+	};
+	struct msg_execve_event *msg;
+	bool throttle = false, send;
+
+	msg = map_lookup_elem(&execve_msg_heap_map, &(__u32){ 0 });
+	if (!msg)
+		return false;
+
+	key.cgroupid = get_task_cgroupid(task);
+	send = cgroup_rate(&key, msg->common.ktime, &settings, &throttle);
+	if (throttle)
+		send_throttle(ctx, MSG_OP_CLONE);
+	return send;
+}
+
 #endif /* __RATE_H__ */
diff --git a/pkg/sensors/base/base.go b/pkg/sensors/base/base.go
@@ -55,6 +55,8 @@ var (
 
 	ExecveJoinMap = program.MapBuilder("tg_execve_joined_info_map", ExecveBprmCommit)
 
+	ForkTailCallsMap = program.MapBuilderPin("fork_calls", "fork_calls", Fork)
+
 	/* Tetragon runtime configuration */
 	TetragonConfMap = program.MapBuilder("tg_conf_map", Execve)
 
@@ -118,6 +120,7 @@ func GetDefaultMaps() []*program.Map {
 		ExecveStats,
 		ExecveJoinMapStats,
 		ExecveTailCallsMap,
+		ForkTailCallsMap,
 		TCPMonMap,
 		TetragonConfMap,
 		StatsMap,

diff --git a/pkg/sensors/program/loader.go b/pkg/sensors/program/loader.go
@@ -455,7 +455,7 @@ func LoadRawTracepointProgram(bpfDir string, load *Program, verbose int) error {
 func LoadKprobeProgram(bpfDir string, load *Program, verbose int) error {
 	var ci *customInstall
 	for mName, mPath := range load.PinMap {
-		if mName == "kprobe_calls" || mName == "retkprobe_calls" {
+		if mName == "kprobe_calls" || mName == "retkprobe_calls" || mName == "fork_calls" {
 			ci = &customInstall{mPath, "kprobe"}
 			break
 		}

diff --git a/pkg/testutils/sensors/load.go b/pkg/testutils/sensors/load.go
@@ -136,21 +136,22 @@ func mergeInBaseSensorMaps(t *testing.T, sensorMaps []SensorMap, sensorProgs []S
 		3: SensorProg{Name: "execve_send", Type: ebpf.TracePoint},
 		4: SensorProg{Name: "tg_kp_bprm_committing_creds", Type: ebpf.Kprobe},
 		5: SensorProg{Name: "execve_rate", Type: ebpf.TracePoint},
+		6: SensorProg{Name: "fork_send", Type: ebpf.Kprobe},
 	}
 
 	var baseMaps = []SensorMap{
 		// all programs
-		SensorMap{Name: "execve_map", Progs: []uint{0, 1, 2, 3, 4, 5}},
-		SensorMap{Name: "tcpmon_map", Progs: []uint{0, 1, 2, 3, 5}},
+		SensorMap{Name: "execve_map", Progs: []uint{0, 1, 2, 3, 4, 5, 6}},
+		SensorMap{Name: "tcpmon_map", Progs: []uint{0, 1, 2, 3, 5, 6}},
 
 		// all but event_execve
-		SensorMap{Name: "execve_map_stats", Progs: []uint{1, 2}},
+		SensorMap{Name: "execve_map_stats", Progs: []uint{1, 6}},
 
 		// event_execve
-		SensorMap{Name: "tg_conf_map", Progs: []uint{0}},
+		SensorMap{Name: "tg_conf_map", Progs: []uint{0, 2}},
 
-		// event_wake_up_new_task
-		SensorMap{Name: "execve_val", Progs: []uint{2}},
+		// fork_send
+		SensorMap{Name: "execve_val", Progs: []uint{6}},
 
 		// event_execve and tg_kp_bprm_committing_creds
 		SensorMap{Name: "tg_execve_joined_info_map", Progs: []uint{0, 4}},