From ee8232f7a1ad84bcee537dbd5857724ee83ebdd1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Momar=20TOUR=C3=89?= <momar.toure@datadoghq.com>
Date: Sun, 29 Dec 2024 16:47:39 +0100
Subject: [PATCH] fix mmap - mkdir_error should fail now - fix errors

---
 pkg/security/tests/module_tester.go       | 42 +++++++++++++++++------
 pkg/security/tests/module_tester_linux.go | 10 +-----
 2 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/pkg/security/tests/module_tester.go b/pkg/security/tests/module_tester.go
index a72da8d6e389a..3f10111a3dd5c 100644
--- a/pkg/security/tests/module_tester.go
+++ b/pkg/security/tests/module_tester.go
@@ -19,7 +19,6 @@ import (
 	"os/exec"
 	"path"
 	"reflect"
-	"slices"
 	"strings"
 	"sync"
 	"testing"
@@ -321,7 +320,7 @@ func (tm *testModule) RegisterRuleEventHandler(cb onRuleHandler) {
 	tm.eventHandlers.Unlock()
 }
 
-func (tm *testModule) GetCustomEventSent(tb testing.TB, action func() error, cb func(rule *rules.Rule, event *events.CustomEvent) bool, timeout time.Duration, eventType model.EventType, ruleIDs ...string) error {
+func (tm *testModule) GetCustomEventSent(tb testing.TB, action func() error, cb func(rule *rules.Rule, event *events.CustomEvent) bool, timeout time.Duration, eventType model.EventType, ruleID string) error {
 	tb.Helper()
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -329,7 +328,7 @@ func (tm *testModule) GetCustomEventSent(tb testing.TB, action func() error, cb
 	message := make(chan ActionMessage, 1)
 
 	tm.RegisterCustomSendEventHandler(func(rule *rules.Rule, event *events.CustomEvent) {
-		if event.GetEventType() != eventType || !slices.Contains(ruleIDs, rule.ID) {
+		if event.GetEventType() != eventType || rule.ID != ruleID {
 			return
 		}
 
@@ -351,15 +350,11 @@ func (tm *testModule) GetCustomEventSent(tb testing.TB, action func() error, cb
 	})
 	defer tm.RegisterCustomSendEventHandler(nil)
 
-	if action == nil {
-		message <- Continue
-	} else {
-		if err := action(); err != nil {
-			message <- Skip
-			return err
-		}
-		message <- Continue
+	if err := action(); err != nil {
+		message <- Skip
+		return err
 	}
+	message <- Continue
 
 	select {
 	case <-time.After(timeout):
@@ -369,6 +364,31 @@ func (tm *testModule) GetCustomEventSent(tb testing.TB, action func() error, cb
 	}
 }
 
+// WaitForPotentialAbnormalPath waits for potential abnormal_path errors. It is use to check before closing the test module
+func (tm *testModule) WaitForPotentialAbnormalPath(timeout time.Duration) bool {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	message := make(chan string, 1)
+
+	tm.RegisterCustomSendEventHandler(func(rule *rules.Rule, event *events.CustomEvent) {
+		if rule.ID == events.AbnormalPathRuleID {
+			message <- "FOUND"
+			cancel()
+		}
+	})
+	defer tm.RegisterCustomSendEventHandler(nil)
+
+	select {
+	case <-message:
+		return true
+	case <-time.After(timeout):
+		return false
+	case <-ctx.Done():
+		return false
+	}
+}
+
 func (tm *testModule) GetEventSent(tb testing.TB, action func() error, cb func(rule *rules.Rule, event *model.Event) bool, timeout time.Duration, ruleID eval.RuleID) error {
 	tb.Helper()
 	ctx, cancel := context.WithCancel(context.Background())
diff --git a/pkg/security/tests/module_tester_linux.go b/pkg/security/tests/module_tester_linux.go
index 96c032ceb3319..1e47e6b2c10ca 100644
--- a/pkg/security/tests/module_tester_linux.go
+++ b/pkg/security/tests/module_tester_linux.go
@@ -37,7 +37,6 @@ import (
 	"github.com/DataDog/datadog-agent/pkg/eventmonitor"
 	secconfig "github.com/DataDog/datadog-agent/pkg/security/config"
 	"github.com/DataDog/datadog-agent/pkg/security/ebpf/kernel"
-	"github.com/DataDog/datadog-agent/pkg/security/events"
 	"github.com/DataDog/datadog-agent/pkg/security/module"
 	sprobe "github.com/DataDog/datadog-agent/pkg/security/probe"
 	"github.com/DataDog/datadog-agent/pkg/security/proto/api"
@@ -966,8 +965,7 @@ func (tm *testModule) validateSyscallsInFlight() {
 }
 
 func (tm *testModule) Close() {
-
-	waitForPotentialEventError(tm, nil, 2*time.Second)
+	tm.WaitForPotentialAbnormalPath(2 * time.Second)
 
 	if !tm.opts.staticOpts.disableRuntimeSecurity {
 		// The stats from the rate_limiter should sent, tm.eventMonitor.SendStats() does not do that
@@ -1110,12 +1108,6 @@ func waitForIMDSResponseProbeEvent(test *testModule, action func() error, proces
 	}...)
 }
 
-func waitForPotentialEventError(test *testModule, action func() error, timeout time.Duration) error {
-	return test.GetCustomEventSent(test.t, action, func(rule *rules.Rule, event *events.CustomEvent) bool {
-		return true
-	}, timeout, model.CustomEventType, events.AbnormalPathRuleID)
-}
-
 //nolint:deadcode,unused
 func checkKernelCompatibility(tb testing.TB, why string, skipCheck func(kv *kernel.Version) bool) {
 	tb.Helper()