From 5b2efedef67c71d1080c42a2c2d760ae241184ed Mon Sep 17 00:00:00 2001
From: Jiri Olsa <jolsa@kernel.org>
Date: Fri, 26 Jul 2024 17:29:20 +0000
Subject: [PATCH] tetragon: Add cgroup rate throttle test for kprobe

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 pkg/sensors/tracing/kprobe_test.go | 68 ++++++++++++++++++++++++++++++
 1 file changed, 68 insertions(+)

diff --git a/pkg/sensors/tracing/kprobe_test.go b/pkg/sensors/tracing/kprobe_test.go
index 6673094c127..be041666b08 100644
--- a/pkg/sensors/tracing/kprobe_test.go
+++ b/pkg/sensors/tracing/kprobe_test.go
@@ -6808,3 +6808,71 @@ spec:
 	err = jsonchecker.JsonTestCheck(t, checker)
 	assert.NoError(t, err)
 }
+
+func testKprobeThrottle(t *testing.T) {
+	var doneWG, readyWG sync.WaitGroup
+	defer doneWG.Wait()
+
+	tracingPolicy := `
+apiVersion: cilium.io/v1alpha1
+kind: TracingPolicy
+metadata:
+  name: "privileges-raise"
+  annotations:
+    description: "Detects privileges change operations"
+spec:
+  kprobes:
+  - call: "sys_prctl"
+    syscall: true
+    args:
+    - index: 0
+      type: "int"
+`
+
+	createCrdFile(t, tracingPolicy)
+
+	throttleStartChecker := ec.NewProcessThrottleChecker("THROTTLE").
+		WithType(tetragon.ThrottleType_THROTTLE_START)
+
+	throttleStopChecker := ec.NewProcessThrottleChecker("THROTTLE").
+		WithType(tetragon.ThrottleType_THROTTLE_STOP)
+
+	checker := ec.NewUnorderedEventChecker(throttleStartChecker, throttleStopChecker)
+
+	ctx, cancel := context.WithTimeout(context.Background(), tus.Conf().CmdWaitTime)
+	defer cancel()
+
+	option.Config.CgroupRate = option.ParseCgroupRate("10,2s")
+	t.Cleanup(func() {
+		option.Config.CgroupRate = option.CgroupRate{}
+	})
+
+	obs, err := observertesthelper.GetDefaultObserverWithFile(t, ctx, testConfigFile, tus.Conf().TetragonLib)
+	if err != nil {
+		t.Fatalf("GetDefaultObserverWithFile error: %s", err)
+	}
+	observertesthelper.LoopEvents(ctx, t, &doneWG, &readyWG, obs)
+	readyWG.Wait()
+
+	// create the load 40 fork/exec per sec for 4 seconds
+	// to get THROTTLE START
+	for cnt := 0; cnt < 2000; cnt++ {
+		unix.Prctl(99999, 99999, 99999, 99999, 99999)
+		time.Sleep(time.Millisecond)
+	}
+
+	// and calm down to get THROTTLE STOP
+	time.Sleep(8 * time.Second)
+
+	err = jsonchecker.JsonTestCheck(t, checker)
+	assert.NoError(t, err)
+}
+
+func TestKprobeThrottle1(t *testing.T) {
+	testKprobeThrottle(t)
+}
+
+// Run throttle twice to test the CgroupRate setup code
+func TestKprobeThrottle2(t *testing.T) {
+	testKprobeThrottle(t)
+}