[refactor] - Add DataOrErr (#3520)

* adjust error handling to make more explicit * Add DataOrErr * update * fix apk handler to use DataOrErr * fix * fix tests * remove timeout as it wont ever be used
trufflesecurity · Nov 16, 2024 · aeaebdd · aeaebdd
1 parent af3e682
commit aeaebdd
Show file tree

Hide file tree

Showing 14 changed files with 310 additions and 150 deletions.
diff --git a/pkg/engine/engine_test.go b/pkg/engine/engine_test.go
@@ -13,17 +13,16 @@ import (
 
 	"github.com/stretchr/testify/assert"
 
-	"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/gitlab/v2"
-	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
-
 	"github.com/trufflesecurity/trufflehog/v3/pkg/config"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/custom_detectors"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/decoders"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/detectors/gitlab/v2"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/engine/ahocorasick"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/engine/defaults"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/custom_detectorspb"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/source_metadatapb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
@@ -317,8 +316,8 @@ aws_secret_access_key = 5dkLVuqpZhD6V3Zym1hivdSHOzh6FGPjwplXD+5f`,
 		{
 			name: "secret with mixed whitespace before",
 			content: `first line
-   
-		
+
+
 AKIA2OGYBAH6STMMNXNN
 aws_secret_access_key = 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`,
 			expectedLine: 4,
@@ -1248,7 +1247,7 @@ def test_something():
 			conf := Config{
 				Concurrency:   1,
 				Decoders:      decoders.DefaultDecoders(),
-				Detectors:     DefaultDetectors(),
+				Detectors:     defaults.DefaultDetectors(),
 				Verify:        false,
 				SourceManager: sourceManager,
 				Dispatcher:    NewPrinterDispatcher(new(discardPrinter)),

diff --git a/pkg/handlers/apk.go b/pkg/handlers/apk.go
@@ -12,9 +12,9 @@ import (
 	"strings"
 	"time"
 
+	"github.com/avast/apkparser"
 	dextk "github.com/csnewman/dextk"
 
-	"github.com/avast/apkparser"
 	logContext "github.com/trufflesecurity/trufflehog/v3/pkg/context"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/engine/defaults"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/iobuf"
@@ -66,45 +66,40 @@ func newAPKHandler() *apkHandler {
 }
 
 // HandleFile processes apk formatted files.
-func (h *apkHandler) HandleFile(ctx logContext.Context, input fileReader) (chan []byte, error) {
-	apkChan := make(chan []byte, defaultBufferSize)
+func (h *apkHandler) HandleFile(ctx logContext.Context, input fileReader) chan DataOrErr {
+	apkChan := make(chan DataOrErr, defaultBufferSize)
 
 	go func() {
-		ctx, cancel := logContext.WithTimeout(ctx, maxTimeout)
-		defer cancel()
 		defer close(apkChan)
 
-		// Update the metrics for the file processing.
-		start := time.Now()
-		var err error
-		defer func() {
-			h.measureLatencyAndHandleErrors(start, err)
-			h.metrics.incFilesProcessed()
-		}()
-
 		// Defer a panic recovery to handle any panics that occur during the APK processing.
 		defer func() {
 			if r := recover(); r != nil {
 				// Return the panic as an error.
+				var panicErr error
 				if e, ok := r.(error); ok {
-					err = e
+					panicErr = e
 				} else {
-					err = fmt.Errorf("panic occurred: %v", r)
+					panicErr = fmt.Errorf("panic occurred: %v", r)
 				}
-				ctx.Logger().Error(err, "Panic occurred when reading apk archive")
+				ctx.Logger().Error(panicErr, "Panic occurred when reading apk archive")
 			}
 		}()
 
-		if err = h.processAPK(ctx, input, apkChan); err != nil {
-			ctx.Logger().Error(err, "error processing apk content")
+		start := time.Now()
+		err := h.processAPK(ctx, input, apkChan)
+		if err == nil {
+			h.metrics.incFilesProcessed()
 		}
+
+		h.measureLatencyAndHandleErrors(ctx, start, err, apkChan)
 	}()
-	return apkChan, nil
+
+	return apkChan
 }
 
 // processAPK processes the apk file and sends the extracted data to the provided channel.
-func (h *apkHandler) processAPK(ctx logContext.Context, input fileReader, apkChan chan []byte) error {
-
+func (h *apkHandler) processAPK(ctx logContext.Context, input fileReader, apkChan chan DataOrErr) error {
 	// Create a ZIP reader from the input fileReader
 	zipReader, err := createZipReader(input)
 	if err != nil {
@@ -132,7 +127,7 @@ func (h *apkHandler) processAPK(ctx logContext.Context, input fileReader, apkCha
 }
 
 // processResources processes the resources.arsc file and sends the extracted data to the provided channel.
-func (h *apkHandler) processResources(ctx logContext.Context, resTable *apkparser.ResourceTable, apkChan chan []byte) error {
+func (h *apkHandler) processResources(ctx logContext.Context, resTable *apkparser.ResourceTable, apkChan chan DataOrErr) error {
 	if resTable == nil {
 		return errors.New("ResourceTable is nil")
 	}
@@ -144,7 +139,7 @@ func (h *apkHandler) processResources(ctx logContext.Context, resTable *apkparse
 }
 
 // processFile processes the file and sends the extracted data to the provided channel.
-func (h *apkHandler) processFile(ctx logContext.Context, file *zip.File, resTable *apkparser.ResourceTable, apkChan chan []byte) error {
+func (h *apkHandler) processFile(ctx logContext.Context, file *zip.File, resTable *apkparser.ResourceTable, apkChan chan DataOrErr) error {
 	// check if the file is empty
 	if file.UncompressedSize64 == 0 {
 		return nil
@@ -177,7 +172,7 @@ func (h *apkHandler) processFile(ctx logContext.Context, file *zip.File, resTabl
 }
 
 // handleAPKFileContent sends the extracted data to the provided channel via the handleNonArchiveContent function.
-func (h *apkHandler) handleAPKFileContent(ctx logContext.Context, rdr io.Reader, fileName string, apkChan chan []byte) error {
+func (h *apkHandler) handleAPKFileContent(ctx logContext.Context, rdr io.Reader, fileName string, apkChan chan DataOrErr) error {
 	mimeReader, err := newMimeTypeReader(rdr)
 	if err != nil {
 		return fmt.Errorf("failed to create mimeTypeReader for file %s: %w", fileName, err)

diff --git a/pkg/handlers/apk_test.go b/pkg/handlers/apk_test.go
@@ -19,7 +19,6 @@ func TestAPKHandler(t *testing.T) {
 		expectedChunks  int
 		expectedSecrets int
 		matchString     string
-		expectErr       bool
 	}{
 		"apk_with_3_leaked_keys": {
 			"https://github.com/joeleonjr/leakyAPK/raw/refs/heads/main/aws_leak.apk",
@@ -28,7 +27,6 @@ func TestAPKHandler(t *testing.T) {
 			// we're just looking for a string match. There is one extra string match in the APK (but only 3 detected secrets).
 			4,
 			"AKIA2UC3BSXMLSCLTUUS",
-			false,
 		},
 	}
 
@@ -47,19 +45,15 @@ func TestAPKHandler(t *testing.T) {
 			}
 			defer newReader.Close()
 
-			archiveChan, err := handler.HandleFile(logContext.Background(), newReader)
-			if testCase.expectErr {
-				assert.NoError(t, err)
-				return
-			}
+			archiveChan := handler.HandleFile(logContext.Background(), newReader)
 
 			chunkCount := 0
 			secretCount := 0
 			re := regexp.MustCompile(testCase.matchString)
 			matched := false
 			for chunk := range archiveChan {
 				chunkCount++
-				if re.Match(chunk) {
+				if re.Match(chunk.Data) {
 					secretCount++
 					matched = true
 				}
@@ -82,7 +76,7 @@ func TestOpenInvalidAPK(t *testing.T) {
 	assert.NoError(t, err)
 	defer rdr.Close()
 
-	archiveChan := make(chan []byte)
+	archiveChan := make(chan DataOrErr)
 
 	err = handler.processAPK(ctx, rdr, archiveChan)
 	assert.Contains(t, err.Error(), "zip: not a valid zip file")
@@ -106,7 +100,7 @@ func TestOpenValidZipInvalidAPK(t *testing.T) {
 	assert.NoError(t, err)
 	defer newReader.Close()
 
-	archiveChan := make(chan []byte)
+	archiveChan := make(chan DataOrErr)
 	ctx := logContext.AddLogger(context.Background())
 
 	err = handler.processAPK(ctx, newReader, archiveChan)

diff --git a/pkg/handlers/ar.go b/pkg/handlers/ar.go
@@ -22,16 +22,16 @@ func newARHandler() *arHandler {
 
 // HandleFile processes AR formatted files. This function needs to be implemented to extract or
 // manage data from AR files according to specific requirements.
-func (h *arHandler) HandleFile(ctx logContext.Context, input fileReader) (chan []byte, error) {
-	archiveChan := make(chan []byte, defaultBufferSize)
+func (h *arHandler) HandleFile(ctx logContext.Context, input fileReader) chan DataOrErr {
+	dataOrErrChan := make(chan DataOrErr, defaultBufferSize)
 
 	if feature.ForceSkipArchives.Load() {
-		close(archiveChan)
-		return archiveChan, nil
+		close(dataOrErrChan)
+		return dataOrErrChan
 	}
 
 	go func() {
-		defer close(archiveChan)
+		defer close(dataOrErrChan)
 
 		// Defer a panic recovery to handle any panics that occur during the AR processing.
 		defer func() {
@@ -53,19 +53,19 @@ func (h *arHandler) HandleFile(ctx logContext.Context, input fileReader) (chan [
 			return
 		}
 
-		err = h.processARFiles(ctx, arReader, archiveChan)
+		err = h.processARFiles(ctx, arReader, dataOrErrChan)
 		if err == nil {
 			h.metrics.incFilesProcessed()
 		}
 
 		// Update the metrics for the file processing and handle any errors.
-		h.measureLatencyAndHandleErrors(start, err)
+		h.measureLatencyAndHandleErrors(ctx, start, err, dataOrErrChan)
 	}()
 
-	return archiveChan, nil
+	return dataOrErrChan
 }
 
-func (h *arHandler) processARFiles(ctx logContext.Context, reader *deb.Ar, archiveChan chan []byte) error {
+func (h *arHandler) processARFiles(ctx logContext.Context, reader *deb.Ar, dataOrErrChan chan DataOrErr) error {
 	for {
 		select {
 		case <-ctx.Done():
@@ -88,7 +88,7 @@ func (h *arHandler) processARFiles(ctx logContext.Context, reader *deb.Ar, archi
 				return fmt.Errorf("error creating mime-type reader: %w", err)
 			}
 
-			if err := h.handleNonArchiveContent(fileCtx, rdr, archiveChan); err != nil {
+			if err := h.handleNonArchiveContent(fileCtx, rdr, dataOrErrChan); err != nil {
 				fileCtx.Logger().Error(err, "error handling archive content in AR")
 				h.metrics.incErrors()
 			}

diff --git a/pkg/handlers/ar_test.go b/pkg/handlers/ar_test.go
@@ -23,12 +23,12 @@ func TestHandleARFile(t *testing.T) {
 	defer rdr.Close()
 
 	handler := newARHandler()
-	archiveChan, err := handler.HandleFile(context.AddLogger(ctx), rdr)
+	dataOrErrChan := handler.HandleFile(context.AddLogger(ctx), rdr)
 	assert.NoError(t, err)
 
 	wantChunkCount := 102
 	count := 0
-	for range archiveChan {
+	for range dataOrErrChan {
 		count++
 	}
 

diff --git a/pkg/handlers/archive.go b/pkg/handlers/archive.go
@@ -44,16 +44,16 @@ func newArchiveHandler() *archiveHandler {
 // utilizing a single output channel. It first tries to identify the input as an archive. If it is an archive,
 // it processes it accordingly; otherwise, it handles the input as non-archive content.
 // The function returns a channel that will receive the extracted data bytes and an error if the initial setup fails.
-func (h *archiveHandler) HandleFile(ctx logContext.Context, input fileReader) (chan []byte, error) {
-	dataChan := make(chan []byte, defaultBufferSize)
+func (h *archiveHandler) HandleFile(ctx logContext.Context, input fileReader) chan DataOrErr {
+	dataOrErrChan := make(chan DataOrErr, defaultBufferSize)
 
 	if feature.ForceSkipArchives.Load() {
-		close(dataChan)
-		return dataChan, nil
+		close(dataOrErrChan)
+		return dataOrErrChan
 	}
 
 	go func() {
-		defer close(dataChan)
+		defer close(dataOrErrChan)
 
 		// The underlying 7zip library may panic when attempting to open an archive.
 		// This is due to an Index Out Of Range (IOOR) error when reading the archive header.
@@ -71,16 +71,16 @@ func (h *archiveHandler) HandleFile(ctx logContext.Context, input fileReader) (c
 		}()
 
 		start := time.Now()
-		err := h.openArchive(ctx, 0, input, dataChan)
+		err := h.openArchive(ctx, 0, input, dataOrErrChan)
 		if err == nil {
 			h.metrics.incFilesProcessed()
 		}
 
 		// Update the metrics for the file processing and handle any errors.
-		h.measureLatencyAndHandleErrors(start, err)
+		h.measureLatencyAndHandleErrors(ctx, start, err, dataOrErrChan)
 	}()
 
-	return dataChan, nil
+	return dataOrErrChan
 }
 
 var ErrMaxDepthReached = errors.New("max archive depth reached")
@@ -89,7 +89,12 @@ var ErrMaxDepthReached = errors.New("max archive depth reached")
 // It takes a reader from which it attempts to identify and process the archive format. Depending on the archive type,
 // it either decompresses or extracts the contents directly, sending data to the provided channel.
 // Returns an error if the archive cannot be processed due to issues like exceeding maximum depth or unsupported formats.
-func (h *archiveHandler) openArchive(ctx logContext.Context, depth int, reader fileReader, archiveChan chan []byte) error {
+func (h *archiveHandler) openArchive(
+	ctx logContext.Context,
+	depth int,
+	reader fileReader,
+	dataOrErrChan chan DataOrErr,
+) error {
 	ctx.Logger().V(4).Info("Starting archive processing", "depth", depth)
 	defer ctx.Logger().V(4).Info("Finished archive processing", "depth", depth)
 
@@ -104,7 +109,7 @@ func (h *archiveHandler) openArchive(ctx logContext.Context, depth int, reader f
 
 	if reader.format == nil {
 		if depth > 0 {
-			return h.handleNonArchiveContent(ctx, newMimeTypeReaderFromFileReader(reader), archiveChan)
+			return h.handleNonArchiveContent(ctx, newMimeTypeReaderFromFileReader(reader), dataOrErrChan)
 		}
 		return fmt.Errorf("unknown archive format")
 	}
@@ -132,9 +137,9 @@ func (h *archiveHandler) openArchive(ctx logContext.Context, depth int, reader f
 		}
 		defer rdr.Close()
 
-		return h.openArchive(ctx, depth+1, rdr, archiveChan)
+		return h.openArchive(ctx, depth+1, rdr, dataOrErrChan)
 	case archiver.Extractor:
-		err := archive.Extract(logContext.WithValue(ctx, depthKey, depth+1), reader, nil, h.extractorHandler(archiveChan))
+		err := archive.Extract(logContext.WithValue(ctx, depthKey, depth+1), reader, nil, h.extractorHandler(dataOrErrChan))
 		if err != nil {
 			return fmt.Errorf("error extracting archive with format: %s: %w", reader.format.Name(), err)
 		}
@@ -148,7 +153,7 @@ func (h *archiveHandler) openArchive(ctx logContext.Context, depth int, reader f
 // It logs the extraction, checks for cancellation, and decides whether to skip the file based on its name or type,
 // particularly for binary files if configured to skip. If the file is not skipped, it recursively calls openArchive
 // to handle nested archives or to continue processing based on the file's content and depth in the archive structure.
-func (h *archiveHandler) extractorHandler(archiveChan chan []byte) func(context.Context, archiver.File) error {
+func (h *archiveHandler) extractorHandler(dataOrErrChan chan DataOrErr) func(context.Context, archiver.File) error {
 	return func(ctx context.Context, file archiver.File) error {
 		lCtx := logContext.WithValues(
 			logContext.AddLogger(ctx),
@@ -220,6 +225,6 @@ func (h *archiveHandler) extractorHandler(archiveChan chan []byte) func(context.
 		h.metrics.observeFileSize(fileSize)
 
 		lCtx.Logger().V(4).Info("Processed file successfully", "filename", file.Name(), "size", file.Size())
-		return h.openArchive(lCtx, depth, rdr, archiveChan)
+		return h.openArchive(lCtx, depth, rdr, dataOrErrChan)
 	}
 }
diff --git a/pkg/handlers/archive_test.go b/pkg/handlers/archive_test.go
@@ -91,7 +91,7 @@ func TestArchiveHandler(t *testing.T) {
 			}
 			defer newReader.Close()
 
-			archiveChan, err := handler.HandleFile(logContext.Background(), newReader)
+			dataOrErrChan := handler.HandleFile(logContext.Background(), newReader)
 			if testCase.expectErr {
 				assert.NoError(t, err)
 				return
@@ -100,9 +100,9 @@ func TestArchiveHandler(t *testing.T) {
 			count := 0
 			re := regexp.MustCompile(testCase.matchString)
 			matched := false
-			for chunk := range archiveChan {
+			for chunk := range dataOrErrChan {
 				count++
-				if re.Match(chunk) {
+				if re.Match(chunk.Data) {
 					matched = true
 				}
 			}
@@ -123,8 +123,8 @@ func TestOpenInvalidArchive(t *testing.T) {
 	assert.NoError(t, err)
 	defer rdr.Close()
 
-	archiveChan := make(chan []byte)
+	dataOrErrChan := make(chan DataOrErr)
 
-	err = handler.openArchive(ctx, 0, rdr, archiveChan)
+	err = handler.openArchive(ctx, 0, rdr, dataOrErrChan)
 	assert.Error(t, err)
 }