diff --git a/pkg/sources/chunker.go b/pkg/sources/chunker.go index 7a97315302ee..ab1d1e82815e 100644 --- a/pkg/sources/chunker.go +++ b/pkg/sources/chunker.go @@ -3,6 +3,7 @@ package sources import ( "bufio" "errors" + "fmt" "io" "github.com/trufflesecurity/trufflehog/v3/pkg/context" @@ -99,6 +100,22 @@ func readInChunks(ctx context.Context, reader io.Reader, config *chunkReaderConf go func() { defer close(chunkResultChan) + // Defer a panic recovery to handle any panics that occur while reading, which can sometimes unavoidably happen + // due to third-party library bugs. + defer func() { + if r := recover(); r != nil { + var panicErr error + if e, ok := r.(error); ok { + panicErr = e + } else { + panicErr = fmt.Errorf("panic occurred: %v", r) + } + chunkResultChan <- ChunkResult{ + err: fmt.Errorf("panic error: %w", panicErr), + } + } + }() + for { chunkRes := ChunkResult{} chunkBytes := make([]byte, config.totalSize) diff --git a/pkg/sources/chunker_test.go b/pkg/sources/chunker_test.go index 818540cb2abb..1051de13576f 100644 --- a/pkg/sources/chunker_test.go +++ b/pkg/sources/chunker_test.go @@ -2,6 +2,7 @@ package sources import ( "bytes" + "io" "math/rand" "runtime" "strings" @@ -10,6 +11,7 @@ import ( "time" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "github.com/trufflesecurity/trufflehog/v3/pkg/context" ) @@ -121,6 +123,21 @@ func TestNewChunkedReader(t *testing.T) { } } +type panicReader struct{} + +var _ io.Reader = (*panicReader)(nil) + +func (_ panicReader) Read([]byte) (int, error) { + panic("panic for testing") +} + +func TestChunkReader_UnderlyingReaderPanics_DoesNotPanic(t *testing.T) { + require.NotPanics(t, func() { + for range NewChunkReader()(context.Background(), &panicReader{}) { + } + }) +} + func BenchmarkChunkReader(b *testing.B) { var bigChunk = make([]byte, 1<<24) // 16MB