mirror of
https://github.com/trufflesecurity/trufflehog.git
synced 2024-11-10 07:04:24 +00:00
[bug] - copy chunk before sending on chunksChan (#1633)
* Redclare chunk before sending on chunksChan. * add integration test. * update test.
This commit is contained in:
parent
fae54c7ffa
commit
b8bb94f2b1
2 changed files with 49 additions and 2 deletions
|
@ -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
|
||||
}
|
||||
|
|
47
pkg/sources/s3/s3_integration_test.go
Normal file
47
pkg/sources/s3/s3_integration_test.go
Normal 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)
|
||||
}
|
Loading…
Reference in a new issue