[bug] - copy chunk before sending on chunksChan (#1633)

* Redclare chunk before sending on chunksChan.

* add integration test.

* update test.
This commit is contained in:
ahrav 2023-08-16 16:36:38 -07:00 committed by GitHub
parent fae54c7ffa
commit b8bb94f2b1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 49 additions and 2 deletions

View file

@ -334,15 +334,15 @@ func (s *Source) pageChunker(ctx context.Context, client *s3.S3, chunksChan chan
}
reader.Stop()
chunk := *chunkSkel
chunkReader := sources.NewChunkReader()
chunkResChan := chunkReader(ctx, reader)
for data := range chunkResChan {
chunk.Data = data.Bytes()
if err := data.Error(); err != nil {
s.log.Error(err, "error reading chunk.")
continue
}
chunk := *chunkSkel
chunk.Data = data.Bytes()
if err := common.CancellableWrite(ctx, chunksChan, &chunk); err != nil {
return err
}

View file

@ -0,0 +1,47 @@
//go:build integration
// +build integration
package s3
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
"google.golang.org/protobuf/types/known/anypb"
"github.com/trufflesecurity/trufflehog/v3/pkg/context"
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
)
func TestSource_ChunksCount(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
defer cancel()
s := Source{}
connection := &sourcespb.S3{
Credential: &sourcespb.S3_Unauthenticated{},
Buckets: []string{"truffletestbucket"},
}
conn, err := anypb.New(connection)
if err != nil {
t.Fatal(err)
}
err = s.Init(ctx, "test name", 0, 0, false, conn, 1)
chunksCh := make(chan *sources.Chunk)
go func() {
defer close(chunksCh)
err = s.Chunks(ctx, chunksCh)
assert.Nil(t, err)
}()
wantChunkCount := 120
got := 0
for range chunksCh {
got++
}
assert.Greater(t, got, wantChunkCount)
}