[not-fixup] - Reduce memory consumption for Buffered File Writer

pkg/engine/ahocorasick/ahocorasickcore.go

@@ -23,6 +23,9 @@ type DetectorKey struct {
 	customDetectorName string
 }
 
+// Type returns the detector type of the key.
+func (k DetectorKey) Type() detectorspb.DetectorType { return k.detectorType }
+
 // AhoCorasickCore encapsulates the operations and data structures used for keyword matching via the
 // Aho-Corasick algorithm. It is responsible for constructing and managing the trie for efficient
 // substring searches, as well as mapping keywords to their associated detectors for rapid lookups.

pkg/engine/engine.go

@@ -598,8 +598,8 @@ func (e *Engine) detectorWorker(ctx context.Context) {
 // by the same detector in the chunk. Exact matches on lookup indicate a duplicate secret for a detector
 // in that chunk - which is expected and not malicious. Those intra-detector dupes are still verified.
 type chunkSecretKey struct {
-	secret       string
-	detectorInfo ahocorasick.DetectorInfo
+	secret      string
+	detectorKey ahocorasick.DetectorKey
 }
 
 func likelyDuplicate(ctx context.Context, val chunkSecretKey, dupes map[chunkSecretKey]struct{}) bool {
@@ -615,7 +615,7 @@ func likelyDuplicate(ctx context.Context, val chunkSecretKey, dupes map[chunkSec
 
 		// If the detector type is the same, we don't need to compare the strings.
 		// These are not duplicates, and should be verified.
-		if val.detectorInfo.Type() == dupeKey.detectorInfo.Type() {
+		if val.detectorKey.Type() == dupeKey.detectorKey.Type() {
 			continue
 		}
 
@@ -674,7 +674,7 @@ func (e *Engine) verificationOverlapWorker(ctx context.Context) {
 				// Ex:
 				// - postman api key: PMAK-qnwfsLyRSyfCwfpHaQP1UzDhrgpWvHjbYzjpRCMshjt417zWcrzyHUArs7r
 				// - malicious detector "api key": qnwfsLyRSyfCwfpHaQP1UzDhrgpWvHjbYzjpRCMshjt417zWcrzyHUArs7r
-				key := chunkSecretKey{secret: string(val), detectorInfo: detector}
+				key := chunkSecretKey{secret: string(val), detectorKey: detector.Key}
 				if _, ok := chunkSecrets[key]; ok {
 					continue
 				}

pkg/engine/engine_test.go

@@ -562,47 +562,47 @@ func TestLikelyDuplicate(t *testing.T) {
 	}{
 		{
 			name: "exact duplicate different detector",
-			val:  chunkSecretKey{"PMAK-qnwfsLyRSyfCwfpHaQP1UzDhrgpWvHjbYzjpRCMshjt417zWcrzyHUArs7r", detectorA},
+			val:  chunkSecretKey{"PMAK-qnwfsLyRSyfCwfpHaQP1UzDhrgpWvHjbYzjpRCMshjt417zWcrzyHUArs7r", detectorA.Key},
 			dupes: map[chunkSecretKey]struct{}{
-				{"PMAK-qnwfsLyRSyfCwfpHaQP1UzDhrgpWvHjbYzjpRCMshjt417zWcrzyHUArs7r", detectorB}: {},
+				{"PMAK-qnwfsLyRSyfCwfpHaQP1UzDhrgpWvHjbYzjpRCMshjt417zWcrzyHUArs7r", detectorB.Key}: {},
 			},
 			expected: true,
 		},
 		{
 			name: "non-duplicate length outside range",
-			val:  chunkSecretKey{"short", detectorA},
+			val:  chunkSecretKey{"short", detectorA.Key},
 			dupes: map[chunkSecretKey]struct{}{
-				{"muchlongerthanthevalstring", detectorB}: {},
+				{"muchlongerthanthevalstring", detectorB.Key}: {},
 			},
 			expected: false,
 		},
 		{
 			name: "similar within threshold",
-			val:  chunkSecretKey{"PMAK-qnwfsLyRSyfCwfpHaQP1UzDhrgpWvHjbYzjpRCMshjt417zWcrzyHUArs7r", detectorA},
+			val:  chunkSecretKey{"PMAK-qnwfsLyRSyfCwfpHaQP1UzDhrgpWvHjbYzjpRCMshjt417zWcrzyHUArs7r", detectorA.Key},
 			dupes: map[chunkSecretKey]struct{}{
-				{"qnwfsLyRSyfCwfpHaQP1UzDhrgpWvHjbYzjpRCMshjt417zWcrzyHUArs7r", detectorB}: {},
+				{"qnwfsLyRSyfCwfpHaQP1UzDhrgpWvHjbYzjpRCMshjt417zWcrzyHUArs7r", detectorB.Key}: {},
 			},
 			expected: true,
 		},
 		{
 			name: "similar outside threshold",
-			val:  chunkSecretKey{"anotherkey", detectorA},
+			val:  chunkSecretKey{"anotherkey", detectorA.Key},
 			dupes: map[chunkSecretKey]struct{}{
-				{"completelydifferent", detectorB}: {},
+				{"completelydifferent", detectorB.Key}: {},
 			},
 			expected: false,
 		},
 		{
 			name:     "empty strings",
-			val:      chunkSecretKey{"", detectorA},
-			dupes:    map[chunkSecretKey]struct{}{{"", detectorB}: {}},
+			val:      chunkSecretKey{"", detectorA.Key},
+			dupes:    map[chunkSecretKey]struct{}{{"", detectorB.Key}: {}},
 			expected: true,
 		},
 		{
 			name: "similar within threshold same detector",
-			val:  chunkSecretKey{"PMAK-qnwfsLyRSyfCwfpHaQP1UzDhrgpWvHjbYzjpRCMshjt417zWcrzyHUArs7r", detectorA},
+			val:  chunkSecretKey{"PMAK-qnwfsLyRSyfCwfpHaQP1UzDhrgpWvHjbYzjpRCMshjt417zWcrzyHUArs7r", detectorA.Key},
 			dupes: map[chunkSecretKey]struct{}{
-				{"qnwfsLyRSyfCwfpHaQP1UzDhrgpWvHjbYzjpRCMshjt417zWcrzyHUArs7r", detectorA}: {},
+				{"qnwfsLyRSyfCwfpHaQP1UzDhrgpWvHjbYzjpRCMshjt417zWcrzyHUArs7r", detectorA.Key}: {},
 			},
 			expected: false,
 		},

pkg/gitparse/gitparse.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
+	bufferedfilewriter "github.com/trufflesecurity/trufflehog/v3/pkg/writers/buffered_file_writer"
 )
 
 const (
@@ -100,21 +101,26 @@ func (b *buffer) String() (string, error) { return b.Buffer.String(), nil }
 // The use of contentWriter enables the management of diff data either in memory or on disk,
 // based on its size, optimizing resource usage and performance.
 type Diff struct {
-	PathB         string
-	LineStart     int
+	PathB     string
+	LineStart int
+	IsBinary  bool
+
 	contentWriter contentWriter
-	IsBinary      bool
 }
 
 type diffOption func(*Diff)
 
 // withPathB sets the PathB option.
 func withPathB(pathB string) diffOption { return func(d *Diff) { d.PathB = pathB } }
 
+// withCustomContentWriter sets the useCustomContentWriter option.
+func withCustomContentWriter(cr contentWriter) diffOption {
+	return func(d *Diff) { d.contentWriter = cr }
+}
+
 // NewDiff creates a new Diff with a threshold.
 func NewDiff(opts ...diffOption) *Diff {
 	diff := new(Diff)
-	diff.contentWriter = newBuffer()
 	for _, opt := range opts {
 		opt(diff)
 	}
@@ -203,7 +209,8 @@ type Parser struct {
 	maxDiffSize   int
 	maxCommitSize int
 	dateFormat    string
-	contentWriter contentWriter
+
+	useCustomContentWriter bool
 }
 
 type ParseState int
@@ -250,11 +257,9 @@ func (state ParseState) String() string {
 	}[state]
 }
 
-// WithContentWriter sets the ContentWriter for the Parser.
-func WithContentWriter(writer contentWriter) Option {
-	return func(parser *Parser) {
-		parser.contentWriter = writer
-	}
+// UseCustomContentWriter sets useCustomContentWriter option.
+func UseCustomContentWriter() Option {
+	return func(parser *Parser) { parser.useCustomContentWriter = true }
 }
 
 // WithMaxDiffSize sets maxDiffSize option. Diffs larger than maxDiffSize will
@@ -283,7 +288,6 @@ func NewParser(options ...Option) *Parser {
 		dateFormat:    defaultDateFormat,
 		maxDiffSize:   defaultMaxDiffSize,
 		maxCommitSize: defaultMaxCommitSize,
-		contentWriter: newBuffer(),
 	}
 	for _, option := range options {
 		option(parser)
@@ -387,7 +391,9 @@ func (c *Parser) FromReader(ctx context.Context, stdOut io.Reader, commitChan ch
 		totalLogSize int
 	)
 	var latestState = Initial
-	currentDiff := NewDiff()
+
+	writer := c.contentWriter()
+	currentDiff := NewDiff(withCustomContentWriter(writer))
 
 	defer common.RecoverWithExit(ctx)
 	defer close(commitChan)
@@ -430,7 +436,7 @@ func (c *Parser) FromReader(ctx context.Context, stdOut io.Reader, commitChan ch
 				totalLogSize += currentCommit.Size
 			}
 			// Create a new currentDiff and currentCommit
-			currentDiff = NewDiff()
+			currentDiff = NewDiff(withCustomContentWriter(c.contentWriter()))
 			currentCommit = &Commit{Message: strings.Builder{}}
 			// Check that the commit line contains a hash and set it.
 			if len(line) >= 47 {
@@ -498,7 +504,7 @@ func (c *Parser) FromReader(ctx context.Context, stdOut io.Reader, commitChan ch
 					currentCommit.Message.WriteString(oldCommit.Message.String())
 				}
 			}
-			currentDiff = NewDiff()
+			currentDiff = NewDiff(withCustomContentWriter(c.contentWriter()))
 		case isModeLine(isStaged, latestState, line):
 			latestState = ModeLine
 			// NoOp
@@ -538,7 +544,7 @@ func (c *Parser) FromReader(ctx context.Context, stdOut io.Reader, commitChan ch
 				}
 				currentCommit.Diffs = append(currentCommit.Diffs, *currentDiff)
 			}
-			currentDiff = NewDiff(withPathB(currentDiff.PathB))
+			currentDiff = NewDiff(withCustomContentWriter(c.contentWriter()), withPathB(currentDiff.PathB))
 
 			words := bytes.Split(line, []byte(" "))
 			if len(words) >= 3 {
@@ -606,6 +612,13 @@ func (c *Parser) FromReader(ctx context.Context, stdOut io.Reader, commitChan ch
 	ctx.Logger().V(2).Info("finished parsing git log.", "total_log_size", totalLogSize)
 }
 
+func (c *Parser) contentWriter() contentWriter {
+	if c.useCustomContentWriter {
+		return bufferedfilewriter.New()
+	}
+	return newBuffer()
+}
+
 func isMergeLine(isStaged bool, latestState ParseState, line []byte) bool {
 	if isStaged || latestState != CommitLine {
 		return false

pkg/sources/git/git.go

@@ -34,7 +34,6 @@ import (
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sanitizer"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
-	bufferedfilewriter "github.com/trufflesecurity/trufflehog/v3/pkg/writers/buffered_file_writer"
 )
 
 const SourceType = sourcespb.SourceType_SOURCE_TYPE_GIT
@@ -99,7 +98,7 @@ type Config struct {
 func NewGit(config *Config) *Git {
 	var parser *gitparse.Parser
 	if config.UseCustomContentWriter {
-		parser = gitparse.NewParser(gitparse.WithContentWriter(bufferedfilewriter.New()))
+		parser = gitparse.NewParser(gitparse.UseCustomContentWriter())
 	} else {
 		parser = gitparse.NewParser()
 	}
@@ -522,7 +521,7 @@ func (s *Git) ScanCommits(ctx context.Context, repo *git.Repository, path string
 		repoCtx = context.WithValue(ctx, "repo", path)
 	}
 
-	commitChan, err := gitparse.NewParser().RepoPath(repoCtx, path, scanOptions.HeadHash, scanOptions.BaseHash == "", scanOptions.ExcludeGlobs, scanOptions.Bare)
+	commitChan, err := s.parser.RepoPath(repoCtx, path, scanOptions.HeadHash, scanOptions.BaseHash == "", scanOptions.ExcludeGlobs, scanOptions.Bare)
 	if err != nil {
 		return err
 	}