Add context to handler (#877)

* Add context to handler * Return rather than break out of select
trufflesecurity · Oct 28, 2022 · ab71b93 · ab71b93
1 parent 034ca4f
commit ab71b93
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 11 deletions.
diff --git a/pkg/handlers/handlers.go b/pkg/handlers/handlers.go
@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"context"
 	"io"
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
@@ -18,7 +19,7 @@ type Handler interface {
 	New()
 }
 
-func HandleFile(file io.Reader, chunkSkel *sources.Chunk, chunksChan chan (*sources.Chunk)) bool {
+func HandleFile(ctx context.Context, file io.Reader, chunkSkel *sources.Chunk, chunksChan chan (*sources.Chunk)) bool {
 	for _, handler := range DefaultHandlers() {
 		handler.New()
 		var isType bool
@@ -27,10 +28,18 @@ func HandleFile(file io.Reader, chunkSkel *sources.Chunk, chunksChan chan (*sour
 			continue
 		}
 		handlerChan := handler.FromFile(file)
-		for data := range handlerChan {
-			chunk := *chunkSkel
-			chunk.Data = data
-			chunksChan <- &chunk
+		for {
+			select {
+			case data := <-handlerChan:
+				chunk := *chunkSkel
+				chunk.Data = data
+				chunksChan <- &chunk
+			case <-ctx.Done():
+				return false
+			}
+			if handlerChan == nil {
+				break
+			}
 		}
 		return true
 	}

diff --git a/pkg/sources/filesystem/filesystem.go b/pkg/sources/filesystem/filesystem.go
@@ -131,7 +131,7 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
 				},
 				Verify: s.verify,
 			}
-			if handlers.HandleFile(reReader, chunkSkel, chunksChan) {
+			if handlers.HandleFile(ctx, reReader, chunkSkel, chunksChan) {
 				return nil
 			}
 

diff --git a/pkg/sources/git/git.go b/pkg/sources/git/git.go
@@ -366,7 +366,7 @@ func (s *Git) ScanCommits(ctx context.Context, repo *git.Repository, path string
 					SourceMetadata: metadata,
 					Verify:         s.verify,
 				}
-				if err := handleBinary(repo, chunksChan, chunkSkel, commitHash, fileName); err != nil {
+				if err := handleBinary(ctx, repo, chunksChan, chunkSkel, commitHash, fileName); err != nil {
 					log.WithError(err).WithField("file", fileName).Debug("Error handling binary file")
 				}
 				continue
@@ -504,7 +504,7 @@ func (s *Git) ScanUnstaged(ctx context.Context, repo *git.Repository, path strin
 					SourceMetadata: metadata,
 					Verify:         s.verify,
 				}
-				if err := handleBinary(repo, chunksChan, chunkSkel, commitHash, fileName); err != nil {
+				if err := handleBinary(ctx, repo, chunksChan, chunkSkel, commitHash, fileName); err != nil {
 					log.WithError(err).WithField("file", fileName).Debug("Error handling binary file")
 				}
 				continue
@@ -782,7 +782,7 @@ func getSafeRemoteURL(repo *git.Repository, preferred string) string {
 	return safeURL
 }
 
-func handleBinary(repo *git.Repository, chunksChan chan *sources.Chunk, chunkSkel *sources.Chunk, commitHash plumbing.Hash, path string) error {
+func handleBinary(ctx context.Context, repo *git.Repository, chunksChan chan *sources.Chunk, chunkSkel *sources.Chunk, commitHash plumbing.Hash, path string) error {
 	log.WithField("path", path).Trace("Binary file found in repository.")
 	commit, err := repo.CommitObject(commitHash)
 	if err != nil {
@@ -805,7 +805,7 @@ func handleBinary(repo *git.Repository, chunksChan chan *sources.Chunk, chunkSke
 		return err
 	}
 
-	if handlers.HandleFile(reader, chunkSkel, chunksChan) {
+	if handlers.HandleFile(ctx, reader, chunkSkel, chunksChan) {
 		return nil
 	}
 

diff --git a/pkg/sources/s3/s3.go b/pkg/sources/s3/s3.go
@@ -287,7 +287,7 @@ func (s *Source) pageChunker(ctx context.Context, client *s3.S3, chunksChan chan
 				},
 				Verify: s.verify,
 			}
-			if handlers.HandleFile(reader, chunkSkel, chunksChan) {
+			if handlers.HandleFile(ctx, reader, chunkSkel, chunksChan) {
 				return
 			}