From 6558bdad5d35ade347f111f73670b4ae4e5da096 Mon Sep 17 00:00:00 2001
From: Ahrav Dutta <ahrav.dutta@trufflesec.com>
Date: Wed, 17 Apr 2024 17:58:50 -0700
Subject: [PATCH] wrap compReader with DiskbufferReader

---
 pkg/handlers/defaulthandler.go       |   8 +++++++-
 pkg/handlers/handlers_test.go        |  25 +++++++++++++++++++++++++
 pkg/handlers/testdata/example.zip.gz | Bin 0 -> 185 bytes
 3 files changed, 32 insertions(+), 1 deletion(-)
 create mode 100644 pkg/handlers/testdata/example.zip.gz

diff --git a/pkg/handlers/defaulthandler.go b/pkg/handlers/defaulthandler.go
index c6f0fea2e..dae08cafd 100644
--- a/pkg/handlers/defaulthandler.go
+++ b/pkg/handlers/defaulthandler.go
@@ -156,7 +156,13 @@ func (h *defaultHandler) openArchive(ctx logContext.Context, depth int, reader i
 
 		h.metrics.incFilesProcessed()
 
-		return h.openArchive(ctx, depth+1, compReader, archiveChan)
+		reReader, err := diskbufferreader.New(compReader)
+		if err != nil {
+			return fmt.Errorf("error creating reusable reader: %w", err)
+		}
+		defer reReader.Close()
+
+		return h.openArchive(ctx, depth+1, reReader, archiveChan)
 	case archiver.Extractor:
 		err := archive.Extract(logContext.WithValue(ctx, depthKey, depth+1), arReader, nil, h.extractorHandler(archiveChan))
 		if err != nil {
diff --git a/pkg/handlers/handlers_test.go b/pkg/handlers/handlers_test.go
index 192d9473b..3290bfc44 100644
--- a/pkg/handlers/handlers_test.go
+++ b/pkg/handlers/handlers_test.go
@@ -117,6 +117,31 @@ func TestHandleNestedArchives(t *testing.T) {
 	assert.Equal(t, wantCount, count)
 }
 
+func TestHandleCompressedZip(t *testing.T) {
+	file, err := os.Open("testdata/example.zip.gz")
+	assert.Nil(t, err)
+	defer file.Close()
+
+	reader, err := diskbufferreader.New(file)
+	assert.NoError(t, err)
+
+	ctx := logContext.Background()
+
+	chunkCh := make(chan *sources.Chunk)
+	go func() {
+		defer close(chunkCh)
+		err := HandleFile(ctx, reader, &sources.Chunk{}, sources.ChanReporter{Ch: chunkCh})
+		assert.NoError(t, err)
+	}()
+
+	wantCount := 2
+	count := 0
+	for range chunkCh {
+		count++
+	}
+	assert.Equal(t, wantCount, count)
+}
+
 func TestHandleNestedCompressedArchive(t *testing.T) {
 	file, err := os.Open("testdata/nested-compressed-archive.tar.gz")
 	assert.Nil(t, err)
diff --git a/pkg/handlers/testdata/example.zip.gz b/pkg/handlers/testdata/example.zip.gz
new file mode 100644
index 0000000000000000000000000000000000000000..863318791b71cbe49e1e90f80bc438f47017ab7c
GIT binary patch
literal 185
zcmV;q07m~GiwFqiZXjj=17&z&ZE$R5E_!Kj01NPDX5nH0f!Tc%BW}elHk1akL7108
zhM^?2xFjtzC)G%=q@pA=gp+~!X`VtF2$xoHGcdAzWn^Gr5dkU=$;d2L0HVaa{F02+
zA_b5}1&~fXt^gDp;&wFO5(TnF(QGgz+6Gj!8JXmmak*atEDkh)VM!y1h2|Sph+Sy*
npc{$mU}Phckc>3MZzL-lNE;InE(6m0z=i+-BpGxadI10cP*+C@

literal 0
HcmV?d00001