From e6eab5810ca7039a825d07e84ec5b281a1bc25cc Mon Sep 17 00:00:00 2001 From: Miccah Castorina Date: Wed, 2 Nov 2022 14:55:51 -0500 Subject: [PATCH 1/3] Check for closed channel in HandleFile --- pkg/handlers/archive_test.go | 27 +++++++++++++++++++++++++++ pkg/handlers/handlers.go | 18 ++++++++++++------ 2 files changed, 39 insertions(+), 6 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..5e57d61918e3 100644 --- a/pkg/handlers/handlers.go +++ b/pkg/handlers/handlers.go @@ -28,18 +28,24 @@ func HandleFile(ctx context.Context, file io.Reader, chunkSkel *sources.Chunk, c continue } handlerChan := handler.FromFile(file) - for { + var closed bool + for !closed { select { - case data := <-handlerChan: + case data, open := <-handlerChan: + if !open { + closed = true + break + } chunk := *chunkSkel chunk.Data = data - chunksChan <- &chunk + select { + case chunksChan <- &chunk: + case <-ctx.Done(): + return false + } case <-ctx.Done(): return false } - if handlerChan == nil { - break - } } return true } From 314f1075b34790a5c5a8e94875905cab7b1d94da Mon Sep 17 00:00:00 2001 From: Miccah Castorina Date: Wed, 2 Nov 2022 15:28:42 -0500 Subject: [PATCH 2/3] Refactor to be more readable --- pkg/handlers/handlers.go | 46 ++++++++++++++++++++++------------------ 1 file changed, 25 insertions(+), 21 deletions(-) diff --git a/pkg/handlers/handlers.go b/pkg/handlers/handlers.go index 5e57d61918e3..7ba572a1962b 100644 --- a/pkg/handlers/handlers.go +++ b/pkg/handlers/handlers.go @@ -20,34 +20,38 @@ type Handler interface { } func HandleFile(ctx context.Context, file io.Reader, chunkSkel *sources.Chunk, chunksChan chan (*sources.Chunk)) bool { - for _, handler := range DefaultHandlers() { + // Find a handler for this file. + var handler Handler + for _, handler = range DefaultHandlers() { handler.New() var isType bool - file, isType = handler.IsFiletype(file) - if !isType { - continue + if file, isType = handler.IsFiletype(file); isType { + break } - handlerChan := handler.FromFile(file) - var closed bool - for !closed { + } + 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, open := <-handlerChan: - if !open { - closed = true - break - } - chunk := *chunkSkel - chunk.Data = data - select { - case chunksChan <- &chunk: - case <-ctx.Done(): - return false - } + case chunksChan <- &chunk: case <-ctx.Done(): return false } + case <-ctx.Done(): + return false } - return true } - return false } From 492372d9010e2c9828862c0dae8666c73fdf7eb9 Mon Sep 17 00:00:00 2001 From: Miccah Castorina Date: Wed, 2 Nov 2022 15:47:27 -0500 Subject: [PATCH 3/3] Fix handler search --- pkg/handlers/handlers.go | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/pkg/handlers/handlers.go b/pkg/handlers/handlers.go index 7ba572a1962b..d542ecb638f5 100644 --- a/pkg/handlers/handlers.go +++ b/pkg/handlers/handlers.go @@ -22,10 +22,11 @@ type Handler interface { func HandleFile(ctx context.Context, file io.Reader, chunkSkel *sources.Chunk, chunksChan chan (*sources.Chunk)) bool { // Find a handler for this file. var handler Handler - for _, handler = range DefaultHandlers() { - handler.New() + for _, h := range DefaultHandlers() { + h.New() var isType bool - if file, isType = handler.IsFiletype(file); isType { + if file, isType = h.IsFiletype(file); isType { + handler = h break } }