Skip to content

Commit

Permalink
fix the shit
Browse files Browse the repository at this point in the history
  • Loading branch information
ahrav committed Jan 31, 2024
1 parent eef15ac commit 1e9479b
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 4 deletions.
20 changes: 16 additions & 4 deletions pkg/engine/ahocorasick/ahocorasickcore.go
Original file line number Diff line number Diff line change
Expand Up @@ -67,12 +67,24 @@ func NewAhoCorasickCore(allDetectors []detectors.Detector) *AhoCorasickCore {
// provided input. This method populates an existing map rather than allocating a new one because
// it will be called once per chunk and that many allocations has a noticeable performance cost.
func (ac *AhoCorasickCore) PopulateMatchingDetectors(chunkData string, dts map[DetectorKey]detectors.Detector) []detectors.Detector {
matches := ac.prefilter.MatchString(strings.ToLower(chunkData))
lowerChunkData := strings.ToLower(chunkData)
matches := ac.prefilter.MatchString(lowerChunkData)

// Use a map to avoid adding duplicate detectors to the slice.
addedDetectors := make(map[DetectorKey]struct{})
d := make([]detectors.Detector, 0, len(matches))
for _, m := range ac.prefilter.MatchString(strings.ToLower(chunkData)) {

for _, m := range matches {
for _, k := range ac.keywordsToDetectors[m.MatchString()] {
dts[k] = ac.detectorsByKey[k]
d = append(d, ac.detectorsByKey[k])
if _, exists := addedDetectors[k]; !exists {
// Add to the map to track already added detectors
addedDetectors[k] = struct{}{}

// Add the detector to the map and slice
detector := ac.detectorsByKey[k]
dts[k] = detector
d = append(d, detector)
}
}
}

Expand Down
6 changes: 6 additions & 0 deletions pkg/engine/engine.go
Original file line number Diff line number Diff line change
Expand Up @@ -484,6 +484,8 @@ func (e *Engine) Finish(ctx context.Context) error {
close(e.results) // Detector workers are done, close the results channel and call it a day.
e.WgNotifier.Wait() // Wait for the notifier workers to finish notifying results.

fmt.Printf("Counter: %d\n", counter)

if err := cleantemp.CleanTempArtifacts(ctx); err != nil {
ctx.Logger().Error(err, "error cleaning temp artifacts")
}
Expand Down Expand Up @@ -601,6 +603,8 @@ func likelyDuplicate(ctx context.Context, val []byte, dupes map[string]struct{})
return false
}

var counter uint64

func (e *Engine) reverifierWorker(ctx context.Context) {
var wgDetect sync.WaitGroup

Expand All @@ -614,6 +618,7 @@ nextChunk:
for _, detector := range chunk.detectors {
// DO NOT VERIFY at this stage of the pipeline.
results, err := detector.FromData(ctx, false, chunk.chunk.Data)
atomic.AddUint64(&counter, 1)
if err != nil {
ctx.Logger().Error(err, "error verifying chunk")
}
Expand Down Expand Up @@ -707,6 +712,7 @@ func (e *Engine) detectChunk(ctx context.Context, data detectableChunk) {
defer cancel()

results, err := data.detector.FromData(ctx, data.chunk.Verify, data.chunk.Data)
atomic.AddUint64(&counter, 1)
if err != nil {
ctx.Logger().Error(err, "error scanning chunk")
}
Expand Down

0 comments on commit 1e9479b

Please sign in to comment.