From 8be03bd93f38ac847f102ea1b106fda5fc8ca380 Mon Sep 17 00:00:00 2001
From: Jiri Olsa <jolsa@kernel.org>
Date: Tue, 12 Mar 2024 22:49:01 +0000
Subject: [PATCH] tetragon: Add throttle exec/fork event test

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 .../observer_test_helper.go                   | 27 +++++--
 pkg/sensors/exec/exec_test.go                 | 72 +++++++++++++++++++
 2 files changed, 94 insertions(+), 5 deletions(-)

diff --git a/pkg/observer/observertesthelper/observer_test_helper.go b/pkg/observer/observertesthelper/observer_test_helper.go
index db0c3526df3..a484e9a390f 100644
--- a/pkg/observer/observertesthelper/observer_test_helper.go
+++ b/pkg/observer/observertesthelper/observer_test_helper.go
@@ -60,9 +60,11 @@ var (
 )
 
 type testObserverOptions struct {
-	crd    bool
-	config string
-	lib    string
+	crd            bool
+	config         string
+	lib            string
+	execCgroupRate string
+	forkCgroupRate string
 }
 
 type testExporterOptions struct {
@@ -124,6 +126,18 @@ func WithLib(lib string) TestOption {
 	}
 }
 
+func WithExecCgroupRate(rate string) TestOption {
+	return func(o *TestOptions) {
+		o.observer.execCgroupRate = rate
+	}
+}
+
+func WithForkCgroupRate(rate string) TestOption {
+	return func(o *TestOptions) {
+		o.observer.forkCgroupRate = rate
+	}
+}
+
 func testDone(tb testing.TB, obs *observer.Observer) {
 	if tb.Failed() {
 		bugtoolFname := "/tmp/tetragon-bugtool.tar.gz"
@@ -212,7 +226,7 @@ func newDefaultObserver(oo *testObserverOptions) *observer.Observer {
 	return observer.NewObserver(oo.config)
 }
 
-func getDefaultObserver(tb testing.TB, ctx context.Context, base *sensors.Sensor, opts ...TestOption) (*observer.Observer, error) {
+func getDefaultObserver(tb testing.TB, ctx context.Context, initial *sensors.Sensor, opts ...TestOption) (*observer.Observer, error) {
 	testutils.CaptureLog(tb, logger.GetLogger().(*logrus.Logger))
 
 	o := newDefaultTestOptions(opts...)
@@ -221,6 +235,7 @@ func getDefaultObserver(tb testing.TB, ctx context.Context, base *sensors.Sensor
 	if option.Config.HubbleLib == "" {
 		option.Config.HubbleLib = o.observer.lib
 	}
+
 	procfs := os.Getenv("TETRAGON_PROCFS")
 	if procfs != "" {
 		option.Config.ProcFS = procfs
@@ -244,10 +259,12 @@ func getDefaultObserver(tb testing.TB, ctx context.Context, base *sensors.Sensor
 		}
 	}
 
-	if err := loadObserver(tb, ctx, base, tp); err != nil {
+	if err := loadObserver(tb, ctx, initial, tp); err != nil {
 		return nil, err
 	}
 
+	base.Config(o.observer.execCgroupRate, o.observer.forkCgroupRate)
+
 	exportFname, err := testutils.GetExportFilename(tb)
 	if err != nil {
 		return nil, err
diff --git a/pkg/sensors/exec/exec_test.go b/pkg/sensors/exec/exec_test.go
index a87e14d22dc..cf9bef35ddc 100644
--- a/pkg/sensors/exec/exec_test.go
+++ b/pkg/sensors/exec/exec_test.go
@@ -1436,3 +1436,75 @@ func TestExecDeletedBinary(t *testing.T) {
 	err = jsonchecker.JsonTestCheck(t, checker)
 	assert.NoError(t, err)
 }
+
+func TestThrottleExec(t *testing.T) {
+	var doneWG, readyWG sync.WaitGroup
+	defer doneWG.Wait()
+
+	testBin := testutils.RepoRootPath("contrib/tester-progs/execbomb")
+
+	processChecker := ec.NewProcessChecker().
+		WithBinary(sm.Full(testBin))
+
+	throttleChecker := ec.NewProcessThrottleChecker("THROTTLE_EXEC").
+		WithProcess(processChecker).
+		WithOp(tetragon.OpType(tetragon.OpType_OP_EXECVE))
+
+	checker := ec.NewUnorderedEventChecker(throttleChecker)
+
+	ctx, cancel := context.WithTimeout(context.Background(), tus.Conf().CmdWaitTime)
+	defer cancel()
+
+	obs, err := observertesthelper.GetDefaultObserver(t, ctx, tus.Conf().TetragonLib,
+		observertesthelper.WithMyPid(),
+		observertesthelper.WithExecCgroupRate("10,1s,1s"))
+	if err != nil {
+		t.Fatalf("Failed to run observer: %s", err)
+	}
+
+	observertesthelper.LoopEvents(ctx, t, &doneWG, &readyWG, obs)
+	readyWG.Wait()
+
+	if err := exec.Command(testBin, "10", "0").Run(); err != nil {
+		t.Fatalf("Failed to execute test binary: %s\n", err)
+	}
+
+	err = jsonchecker.JsonTestCheck(t, checker)
+	assert.NoError(t, err)
+}
+
+func TestThrottleFork(t *testing.T) {
+	var doneWG, readyWG sync.WaitGroup
+	defer doneWG.Wait()
+
+	testBin := testutils.RepoRootPath("contrib/tester-progs/forkbomb")
+
+	processChecker := ec.NewProcessChecker().
+		WithBinary(sm.Full(testBin))
+
+	throttleChecker := ec.NewProcessThrottleChecker("THROTTLE_FORK").
+		WithProcess(processChecker).
+		WithOp(tetragon.OpType(tetragon.OpType_OP_CLONE))
+
+	checker := ec.NewUnorderedEventChecker(throttleChecker)
+
+	ctx, cancel := context.WithTimeout(context.Background(), tus.Conf().CmdWaitTime)
+	defer cancel()
+
+	obs, err := observertesthelper.GetDefaultObserver(t, ctx, tus.Conf().TetragonLib,
+		observertesthelper.WithMyPid(),
+		observertesthelper.WithForkCgroupRate("10,1s,1s"))
+	if err != nil {
+		t.Fatalf("Failed to run observer: %s", err)
+	}
+
+	observertesthelper.LoopEvents(ctx, t, &doneWG, &readyWG, obs)
+	readyWG.Wait()
+
+	if err := exec.Command(testBin, "10", "0").Run(); err != nil {
+		t.Fatalf("Failed to execute test binary: %s\n", err)
+	}
+
+	err = jsonchecker.JsonTestCheck(t, checker)
+	assert.NoError(t, err)
+}