From e8945406321f2e27a8cec2a22186d4fcffedacdc Mon Sep 17 00:00:00 2001
From: ahrav <ahravdutta02@gmail.com>
Date: Fri, 11 Aug 2023 13:40:10 -0700
Subject: [PATCH] Use the common chunker for scanning the filesystem source
 (#1619)

* Use the common chunker for scanning the filesystem source.

* remove unused conts.

* add test.
---
 pkg/sources/filesystem/filesystem.go      | 51 ++++++++++-------------
 pkg/sources/filesystem/filesystem_test.go | 42 +++++++++++++++++++
 2 files changed, 63 insertions(+), 30 deletions(-)

diff --git a/pkg/sources/filesystem/filesystem.go b/pkg/sources/filesystem/filesystem.go
index 41673a4e2..89d473486 100644
--- a/pkg/sources/filesystem/filesystem.go
+++ b/pkg/sources/filesystem/filesystem.go
@@ -1,7 +1,6 @@
 package filesystem
 
 import (
-	"bufio"
 	"fmt"
 	"io"
 	"io/fs"
@@ -23,13 +22,6 @@ import (
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
 )
 
-const (
-	// These buffer sizes are mainly driven by our largest credential size, which is GCP @ ~2.25KB.
-	// Having a peek size larger than that ensures that we have complete credential coverage in our chunks.
-	BufferSize = 10 * 1024 // 10KB
-	PeekSize   = 3 * 1024  // 3KB
-)
-
 type Source struct {
 	name     string
 	sourceId int64
@@ -187,34 +179,33 @@ func (s *Source) scanFile(ctx context.Context, path string, chunksChan chan *sou
 	}
 	reReader.Stop()
 
-	for {
-		chunkBytes := make([]byte, BufferSize)
-		reader := bufio.NewReaderSize(reReader, BufferSize)
-		n, err := reader.Read(chunkBytes)
-		if err != nil && !errors.Is(err, io.EOF) {
-			break
+	chunkReader := sources.NewChunkReader()
+	chunkResChan := chunkReader(ctx, reReader)
+	for data := range chunkResChan {
+		if err := data.Error(); err != nil {
+			s.log.Error(err, "error reading chunk.")
+			continue
 		}
-		peekData, _ := reader.Peek(PeekSize)
-		if n > 0 {
-			chunksChan <- &sources.Chunk{
-				SourceType: s.Type(),
-				SourceName: s.name,
-				SourceID:   s.SourceID(),
-				Data:       append(chunkBytes[:n], peekData...),
-				SourceMetadata: &source_metadatapb.MetaData{
-					Data: &source_metadatapb.MetaData_Filesystem{
-						Filesystem: &source_metadatapb.Filesystem{
-							File: sanitizer.UTF8(path),
-						},
+
+		chunk := &sources.Chunk{
+			SourceType: s.Type(),
+			SourceName: s.name,
+			SourceID:   s.SourceID(),
+			Data:       data.Bytes(),
+			SourceMetadata: &source_metadatapb.MetaData{
+				Data: &source_metadatapb.MetaData_Filesystem{
+					Filesystem: &source_metadatapb.Filesystem{
+						File: sanitizer.UTF8(path),
 					},
 				},
-				Verify: s.verify,
-			}
+			},
+			Verify: s.verify,
 		}
-		if errors.Is(err, io.EOF) {
-			break
+		if err := common.CancellableWrite(ctx, chunksChan, chunk); err != nil {
+			return err
 		}
 	}
+
 	return nil
 }
 
diff --git a/pkg/sources/filesystem/filesystem_test.go b/pkg/sources/filesystem/filesystem_test.go
index 7d56bc9d9..ef5dfbdd3 100644
--- a/pkg/sources/filesystem/filesystem_test.go
+++ b/pkg/sources/filesystem/filesystem_test.go
@@ -1,10 +1,14 @@
 package filesystem
 
 import (
+	"os"
+	"strings"
 	"testing"
 	"time"
 
+	"github.com/go-logr/logr"
 	"github.com/kylelemons/godebug/pretty"
+	"github.com/stretchr/testify/assert"
 	"google.golang.org/protobuf/types/known/anypb"
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
@@ -78,3 +82,41 @@ func TestSource_Scan(t *testing.T) {
 		})
 	}
 }
+
+func TestScanFile(t *testing.T) {
+	tmpfile, err := os.CreateTemp("", "example.txt")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.Remove(tmpfile.Name())
+
+	chunkSize := sources.ChunkSize
+	secretPart1 := "SECRET"
+	secretPart2 := "SPLIT"
+	// Split the secret into two parts and pad the rest of the chunk with A's.
+	data := strings.Repeat("A", chunkSize-len(secretPart1)) + secretPart1 + secretPart2 + strings.Repeat("A", chunkSize-len(secretPart2))
+
+	_, err = tmpfile.Write([]byte(data))
+	assert.Nil(t, err)
+
+	err = tmpfile.Close()
+	assert.Nil(t, err)
+
+	source := &Source{}
+	chunksChan := make(chan *sources.Chunk, 2)
+
+	ctx := context.WithLogger(context.Background(), logr.Discard())
+	go func() {
+		defer close(chunksChan)
+		err = source.scanFile(ctx, tmpfile.Name(), chunksChan)
+		assert.Nil(t, err)
+	}()
+
+	// Read from the channel and validate the secrets.
+	foundSecret := ""
+	for chunkCh := range chunksChan {
+		foundSecret += string(chunkCh.Data)
+	}
+
+	assert.Contains(t, foundSecret, secretPart1+secretPart2)
+}