From 657353aa125bb28e6d90ed8c30cc85a74cd94379 Mon Sep 17 00:00:00 2001
From: Miccah <m.castorina93@gmail.com>
Date: Wed, 2 Nov 2022 16:35:19 -0500
Subject: [PATCH] Check for closed channel in HandleFile (#895)

* Check for closed channel in HandleFile

* Refactor to be more readable

* Fix handler search
---
 pkg/handlers/archive_test.go | 27 ++++++++++++++++++++++
 pkg/handlers/handlers.go     | 43 ++++++++++++++++++++++--------------
 2 files changed, 54 insertions(+), 16 deletions(-)

diff --git a/pkg/handlers/archive_test.go b/pkg/handlers/archive_test.go
index 99fa1289dbf9..71618af06389 100644
--- a/pkg/handlers/archive_test.go
+++ b/pkg/handlers/archive_test.go
@@ -1,11 +1,15 @@
 package handlers
 
 import (
+	"context"
 	"net/http"
 	"regexp"
+	"strings"
 	"testing"
 
 	diskbufferreader "github.com/bill-rich/disk-buffer-reader"
+	"github.com/stretchr/testify/assert"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
 )
 
 func TestArchiveHandler(t *testing.T) {
@@ -94,3 +98,26 @@ func TestArchiveHandler(t *testing.T) {
 		}
 	}
 }
+
+func TestHandleFile(t *testing.T) {
+	ch := make(chan *sources.Chunk, 2)
+
+	// Context cancels the operation.
+	canceledCtx, cancel := context.WithCancel(context.Background())
+	cancel()
+	assert.False(t, HandleFile(canceledCtx, strings.NewReader("file"), &sources.Chunk{}, ch))
+
+	// Only one chunk is sent on the channel.
+	// TODO: Embed a zip without making an HTTP request.
+	resp, err := http.Get("https://raw.githubusercontent.com/bill-rich/bad-secrets/master/aws-canary-creds.zip")
+	assert.NoError(t, err)
+	defer resp.Body.Close()
+	archive := Archive{}
+	archive.New()
+	reader, err := diskbufferreader.New(resp.Body)
+	assert.NoError(t, err)
+
+	assert.Equal(t, 0, len(ch))
+	assert.True(t, HandleFile(context.Background(), reader, &sources.Chunk{}, ch))
+	assert.Equal(t, 1, len(ch))
+}
diff --git a/pkg/handlers/handlers.go b/pkg/handlers/handlers.go
index 5b78aec67368..d542ecb638f5 100644
--- a/pkg/handlers/handlers.go
+++ b/pkg/handlers/handlers.go
@@ -20,28 +20,39 @@ type Handler interface {
 }
 
 func HandleFile(ctx context.Context, file io.Reader, chunkSkel *sources.Chunk, chunksChan chan (*sources.Chunk)) bool {
-	for _, handler := range DefaultHandlers() {
-		handler.New()
+	// Find a handler for this file.
+	var handler Handler
+	for _, h := range DefaultHandlers() {
+		h.New()
 		var isType bool
-		file, isType = handler.IsFiletype(file)
-		if !isType {
-			continue
+		if file, isType = h.IsFiletype(file); isType {
+			handler = h
+			break
 		}
-		handlerChan := handler.FromFile(file)
-		for {
+	}
+	if handler == nil {
+		return false
+	}
+
+	// Process the file and read all []byte chunks from handlerChan.
+	handlerChan := handler.FromFile(file)
+	for {
+		select {
+		case data, open := <-handlerChan:
+			if !open {
+				// We finished reading everything from handlerChan.
+				return true
+			}
+			chunk := *chunkSkel
+			chunk.Data = data
+			// Send data on chunksChan.
 			select {
-			case data := <-handlerChan:
-				chunk := *chunkSkel
-				chunk.Data = data
-				chunksChan <- &chunk
+			case chunksChan <- &chunk:
 			case <-ctx.Done():
 				return false
 			}
-			if handlerChan == nil {
-				break
-			}
+		case <-ctx.Done():
+			return false
 		}
-		return true
 	}
-	return false
 }