Concurrent detection (#1580)

* Run detection on each chunk concurrently.

* Add printer functionality.

* Add logic for dedupe.

* cleanup.

* Moddify number of notifier workers.

* Add comment.

* move consts into fxn.

* buffer resutls chan.

* fix test.

* address comments.

* return an error from Finish.

* fix test.

* fix test.

* linter.

* check err.

* address comments.

go.mod

@@ -40,6 +40,7 @@ require (
 	github.com/google/go-github/v42 v42.0.0
 	github.com/googleapis/gax-go/v2 v2.12.0
 	github.com/hashicorp/go-retryablehttp v0.7.4
+	github.com/hashicorp/golang-lru v0.5.1
 	github.com/jlaffaye/ftp v0.2.0
 	github.com/joho/godotenv v1.5.1
 	github.com/jpillora/overseer v1.1.6
@@ -128,6 +129,7 @@ require (
 	github.com/hashicorp/errwrap v1.0.0 // indirect
 	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
 	github.com/hashicorp/go-multierror v1.1.1 // indirect
+	github.com/hashicorp/golang-lru/v2 v2.0.4 // indirect
 	github.com/imdario/mergo v0.3.15 // indirect
 	github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect

go.sum

@@ -301,7 +301,10 @@ github.com/hashicorp/go-multierror v1.1.1/go.mod h1:iw975J/qwKPdAO1clOe2L8331t/9
 github.com/hashicorp/go-retryablehttp v0.7.4 h1:ZQgVdpTdAL7WpMIwLzCfbalOcSUdkDZnpUv3/+BxzFA=
 github.com/hashicorp/go-retryablehttp v0.7.4/go.mod h1:Jy/gPYAdjqffZ/yFGCFV2doI5wjtH1ewM9u8iYVjtX8=
 github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
+github.com/hashicorp/golang-lru v0.5.1 h1:0hERBMJE1eitiLkihrMvRVBYAkpHzc/J3QdDN+dAcgU=
 github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
+github.com/hashicorp/golang-lru/v2 v2.0.4 h1:7GHuZcgid37q8o5i3QI9KMT4nCWQQ3Kx3Ov6bb9MfK0=
+github.com/hashicorp/golang-lru/v2 v2.0.4/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
 github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
 github.com/ianlancetaylor/demangle v0.0.0-20210905161508-09a460cdf81d/go.mod h1:aYm2/VgdVmcIU8iMfdMvDMsRAQjcfZSKFby6HOFvi/w=

main.go

@@ -9,7 +9,6 @@ import (
 	"strconv"
 	"strings"
 	"syscall"
-	"time"
 
 	"github.com/felixge/fgprof"
 	"github.com/go-logr/logr"
@@ -316,16 +315,35 @@ func run(state overseer.State) {
 		return true
 	}
 
-	e := engine.Start(ctx,
-		engine.WithConcurrency(*concurrency),
+	// Set how the engine will print its results.
+	var printer engine.Printer
+	switch {
+	case *jsonLegacy:
+		printer = new(output.LegacyJSONPrinter)
+	case *jsonOut:
+		printer = new(output.JSONPrinter)
+	case *gitHubActionsFormat:
+		printer = new(output.GitHubActionsPrinter)
+	default:
+		printer = new(output.PlainPrinter)
+	}
+
+	e, err := engine.Start(ctx,
+		engine.WithConcurrency(uint8(*concurrency)),
 		engine.WithDecoders(decoders.DefaultDecoders()...),
 		engine.WithDetectors(!*noVerification, engine.DefaultDetectors()...),
 		engine.WithDetectors(!*noVerification, conf.Detectors...),
 		engine.WithFilterDetectors(includeFilter),
 		engine.WithFilterDetectors(excludeFilter),
 		engine.WithFilterDetectors(endpointCustomizer),
 		engine.WithFilterUnverified(*filterUnverified),
+		engine.WithOnlyVerified(*onlyVerified),
+		engine.WithPrintAvgDetectorTime(*printAvgDetectorTime),
+		engine.WithPrinter(printer),
 	)
+	if err != nil {
+		logFatal(err, "error initializing engine")
+	}
 
 	var repoPath string
 	var remote bool
@@ -475,52 +493,48 @@ func run(state overseer.State) {
 			logFatal(err, "Failed to scan Docker.")
 		}
 	}
-	// asynchronously wait for scanning to finish and cleanup
-	go e.Finish(ctx, logFatal)
 
 	if !*jsonLegacy && !*jsonOut {
 		fmt.Fprintf(os.Stderr, "🐷🔑🐷  TruffleHog. Unearth your secrets. 🐷🔑🐷\n\n")
 	}
 
-	// NOTE: this loop will terminate when the results channel is closed in
-	// e.Finish()
-	foundResults := false
-	for r := range e.ResultsChan() {
-		if *onlyVerified && !r.Verified {
-			continue
-		}
-		foundResults = true
-
-		var err error
-		switch {
-		case *jsonLegacy:
-			err = output.PrintLegacyJSON(ctx, &r)
-		case *jsonOut:
-			err = output.PrintJSON(&r)
-		case *gitHubActionsFormat:
-			err = output.PrintGitHubActionsOutput(&r)
-		default:
-			err = output.PrintPlainOutput(&r)
-		}
-		if err != nil {
-			logFatal(err, "error printing results")
-		}
+	// Wait for all workers to finish.
+	if err = e.Finish(ctx); err != nil {
+		logFatal(err, "engine failed to finish execution")
 	}
-	logger.V(2).Info("finished scanning",
-		"chunks", e.ChunksScanned(),
-		"bytes", e.BytesScanned(),
+
+	metrics := e.GetMetrics()
+	// Print results.
+	logger.Info("finished scanning",
+		"chunks", metrics.ChunksScanned,
+		"bytes", metrics.BytesScanned,
+		"verified_secrets", metrics.VerifiedSecretsFound,
+		"unverified_secrets", metrics.UnverifiedSecretsFound,
 	)
 
 	if *printAvgDetectorTime {
 		printAverageDetectorTime(e)
 	}
 
-	if foundResults && *fail {
+	if e.HasFoundResults() && *fail {
 		logger.V(2).Info("exiting with code 183 because results were found")
 		os.Exit(183)
 	}
 }
 
+// logFatalFunc returns a log.Fatal style function. Calling the returned
+// function will terminate the program without cleanup.
+func logFatalFunc(logger logr.Logger) func(error, string, ...any) {
+	return func(err error, message string, keyAndVals ...any) {
+		logger.Error(err, message, keyAndVals...)
+		if err != nil {
+			os.Exit(1)
+			return
+		}
+		os.Exit(0)
+	}
+}
+
 func commaSeparatedToSlice(s []string) []string {
 	var result []string
 	for _, items := range s {
@@ -537,26 +551,8 @@ func commaSeparatedToSlice(s []string) []string {
 
 func printAverageDetectorTime(e *engine.Engine) {
 	fmt.Fprintln(os.Stderr, "Average detector time is the measurement of average time spent on each detector when results are returned.")
-	for detectorName, durations := range e.DetectorAvgTime() {
-		var total time.Duration
-		for _, d := range durations {
-			total += d
-		}
-		avgDuration := total / time.Duration(len(durations))
-		fmt.Fprintf(os.Stderr, "%s: %s\n", detectorName, avgDuration)
-	}
-}
-
-// logFatalFunc returns a log.Fatal style function. Calling the returned
-// function will terminate the program without cleanup.
-func logFatalFunc(logger logr.Logger) func(error, string, ...any) {
-	return func(err error, message string, keyAndVals ...any) {
-		logger.Error(err, message, keyAndVals...)
-		if err != nil {
-			os.Exit(1)
-			return
-		}
-		os.Exit(0)
+	for detectorName, duration := range e.GetDetectorsMetrics() {
+		fmt.Fprintf(os.Stderr, "%s: %s\n", detectorName, duration)
 	}
 }
 

pkg/common/utils.go

@@ -3,7 +3,9 @@ package common
 import (
 	"bufio"
 	"bytes"
+	"crypto/rand"
 	"io"
+	"math/big"
 	"strings"
 )
 
@@ -49,3 +51,16 @@ func ResponseContainsSubstring(reader io.ReadCloser, target string) (bool, error
 	}
 	return false, nil
 }
+
+var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")
+
+// RandomID returns a random string of the given length.
+func RandomID(length int) string {
+	b := make([]rune, length)
+	for i := range b {
+		randInt, _ := rand.Int(rand.Reader, big.NewInt(int64(len(letters))))
+		b[i] = letters[randInt.Int64()]
+	}
+
+	return string(b)
+}