Bug fixes

* Use a `nil` `WriteCloser` when no flag is passed * Call the finish func in `EndUnitChunking` * Ignore eviction logs
trufflesecurity · Jan 13, 2024 · 473d651 · 473d651
1 parent 552578c
commit 473d651
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 24 deletions.
diff --git a/main.go b/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"fmt"
+	"io"
 	"net/http"
 	_ "net/http/pprof"
 	"os"
@@ -394,6 +395,10 @@ func run(state overseer.State) {
 		fmt.Fprintf(os.Stderr, "🐷🔑🐷  TruffleHog. Unearth your secrets. 🐷🔑🐷\n\n")
 	}
 
+	var jobReportWriter io.WriteCloser
+	if *jobReportFile != nil {
+		jobReportWriter = *jobReportFile
+	}
 	e, err := engine.Start(ctx,
 		engine.WithConcurrency(uint8(*concurrency)),
 		engine.WithDecoders(decoders.DefaultDecoders()...),
@@ -408,7 +413,7 @@ func run(state overseer.State) {
 		engine.WithPrintAvgDetectorTime(*printAvgDetectorTime),
 		engine.WithPrinter(printer),
 		engine.WithFilterEntropy(*filterEntropy),
-		engine.WithJobReportWriter(*jobReportFile),
+		engine.WithJobReportWriter(jobReportWriter),
 	)
 	if err != nil {
 		logFatal(err, "error initializing engine")

diff --git a/pkg/engine/engine.go b/pkg/engine/engine.go
@@ -10,6 +10,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/go-logr/logr"
 	lru "github.com/hashicorp/golang-lru"
 	"google.golang.org/protobuf/proto"
 
@@ -339,26 +340,28 @@ func (e *Engine) initSourceManager(ctx context.Context) {
 	}
 	if e.jobReportWriter != nil {
 		var mu sync.Mutex
-		unitHook := sources.NewUnitHook(ctx, sources.OnUnitHookFinish(func(metrics sources.UnitMetrics) {
-			metrics.Errors = common.ExportErrors(metrics.Errors...)
-			details, err := json.Marshal(map[string]any{
-				"version": 1,
-				"data":    metrics,
-			})
-			if err != nil {
-				ctx.Logger().Error(err, "error marshalling job details")
-				return
-			}
-			mu.Lock()
-			defer mu.Unlock()
-			if _, err := e.jobReportWriter.Write(details); err != nil {
-				ctx.Logger().Error(err, "error writing to file")
-				return
-			}
-			if _, err := e.jobReportWriter.Write([]byte{'\n'}); err != nil {
-				ctx.Logger().Error(err, "error writing to file")
-			}
-		}))
+		// Pass a discard logger for the unit hook to ignore the LRU cache eviction logs.
+		unitHook := sources.NewUnitHook(context.WithLogger(ctx, logr.Discard()),
+			sources.OnUnitHookFinish(func(metrics sources.UnitMetrics) {
+				metrics.Errors = common.ExportErrors(metrics.Errors...)
+				details, err := json.Marshal(map[string]any{
+					"version": 1,
+					"data":    metrics,
+				})
+				if err != nil {
+					ctx.Logger().Error(err, "error marshalling job details")
+					return
+				}
+				mu.Lock()
+				defer mu.Unlock()
+				if _, err := e.jobReportWriter.Write(details); err != nil {
+					ctx.Logger().Error(err, "error writing to file")
+					return
+				}
+				if _, err := e.jobReportWriter.Write([]byte{'\n'}); err != nil {
+					ctx.Logger().Error(err, "error writing to file")
+				}
+			}))
 		opts = append(opts, sources.WithReportHook(unitHook))
 	}
 	e.sourceManager = sources.NewManager(opts...)

diff --git a/pkg/sources/job_progress_hook.go b/pkg/sources/job_progress_hook.go
@@ -79,6 +79,9 @@ func (u *UnitHook) EndUnitChunking(ref JobProgressRef, unit SourceUnit, end time
 		return
 	}
 	metrics.EndTime = &end
+	if u.finishFunc != nil {
+		go u.finishFunc(*metrics)
+	}
 }
 
 func (u *UnitHook) ReportChunk(ref JobProgressRef, unit SourceUnit, chunk *Chunk) {
@@ -147,9 +150,9 @@ func (u *UnitHook) Finish(ref JobProgressRef) {
 			metric.StartTime = snap.StartTime
 			metric.EndTime = snap.EndTime
 			metric.Errors = snap.Errors
-		}
-		if u.finishFunc != nil {
-			go u.finishFunc(*metric)
+			if u.finishFunc != nil {
+				go u.finishFunc(*metric)
+			}
 		}
 	}
 }