From 382990a6bd0f1125ce336f407e7980c06e6f7c6d Mon Sep 17 00:00:00 2001
From: ahrav <ahravdutta02@gmail.com>
Date: Fri, 2 Feb 2024 13:43:56 -0800
Subject: [PATCH] [bug] - use DetectorKey as the key in the
 detectorKeysWithResults map (#2366)

* use DetectorKey as the key in the map

* nil check

* update comment
---
 pkg/engine/ahocorasick/ahocorasickcore.go |  6 ++++++
 pkg/engine/engine.go                      | 24 +++++++++++++++--------
 2 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/pkg/engine/ahocorasick/ahocorasickcore.go b/pkg/engine/ahocorasick/ahocorasickcore.go
index f54f723353ea..b083715507db 100644
--- a/pkg/engine/ahocorasick/ahocorasickcore.go
+++ b/pkg/engine/ahocorasick/ahocorasickcore.go
@@ -63,6 +63,12 @@ func NewAhoCorasickCore(allDetectors []detectors.Detector) *AhoCorasickCore {
 	}
 }
 
+// GetDetectorByKey returns the detector associated with the given key. If no detector is found, it
+// returns nil.
+func (ac *AhoCorasickCore) GetDetectorByKey(key DetectorKey) detectors.Detector {
+	return ac.detectorsByKey[key]
+}
+
 // DetectorInfo represents a detected pattern's metadata in a data chunk.
 // It encapsulates the key identifying a specific detector and the detector instance itself.
 type DetectorInfo struct {
diff --git a/pkg/engine/engine.go b/pkg/engine/engine.go
index de3001352233..eda8f3e4c0a4 100644
--- a/pkg/engine/engine.go
+++ b/pkg/engine/engine.go
@@ -644,7 +644,7 @@ func (e *Engine) verificationOverlapWorker(ctx context.Context) {
 
 	// Reuse the same map and slice to avoid allocations.
 	const avgSecretsPerDetector = 8
-	detectorsWithResult := make(map[ahocorasick.DetectorInfo]struct{}, avgSecretsPerDetector)
+	detectorKeysWithResults := make(map[ahocorasick.DetectorKey]struct{}, avgSecretsPerDetector)
 	chunkSecrets := make(map[chunkSecretKey]struct{}, avgSecretsPerDetector)
 
 	for chunk := range e.verificationOverlapChunksChan {
@@ -658,8 +658,8 @@ func (e *Engine) verificationOverlapWorker(ctx context.Context) {
 			if len(results) == 0 {
 				continue
 			}
-			if _, ok := detectorsWithResult[detector]; !ok {
-				detectorsWithResult[detector] = struct{}{}
+			if _, ok := detectorKeysWithResults[detector.Key]; !ok {
+				detectorKeysWithResults[detector.Key] = struct{}{}
 			}
 
 			for _, res := range results {
@@ -693,14 +693,22 @@ func (e *Engine) verificationOverlapWorker(ctx context.Context) {
 						wgDoneFn: wgDetect.Done,
 					}, res)
 
-					// Remove the detector from the list of detectors with results.
-					delete(detectorsWithResult, detector)
+					// Remove the detector key from the list of detector keys with results.
+					// This is to ensure that the chunk is not reprocessed with verification enabled
+					// for this detector.
+					delete(detectorKeysWithResults, detector.Key)
 				}
 				chunkSecrets[key] = struct{}{}
 			}
 		}
 
-		for detector := range detectorsWithResult {
+		for key := range detectorKeysWithResults {
+			detector := e.ahoCorasickCore.GetDetectorByKey(key)
+			if detector == nil {
+				ctx.Logger().Info("detector not found", "key", key)
+				continue
+			}
+
 			wgDetect.Add(1)
 			chunk.chunk.Verify = e.verify
 			e.detectableChunksChan <- detectableChunk{
@@ -715,8 +723,8 @@ func (e *Engine) verificationOverlapWorker(ctx context.Context) {
 		for k := range chunkSecrets {
 			delete(chunkSecrets, k)
 		}
-		for k := range detectorsWithResult {
-			delete(detectorsWithResult, k)
+		for k := range detectorKeysWithResults {
+			delete(detectorKeysWithResults, k)
 		}
 
 		chunk.verificationOverlapWgDoneFn()