test(sources/s3): fix infinite blocking and timeout issue in TestSource_Chunks #4048

Juneezee · 2025-04-14T16:35:44Z

Description:

This PR addresses two issues:

Duplicate test functions in s3_integration_test.go.

The file currently contains two separate TestSourceChunksNoResumption. The only differences between them are the use of t.Parallel() in one and the bucket name. This PR merges both into a single, table-driven test

trufflehog/pkg/sources/s3/s3_integration_test.go

Lines 85 to 114 in a1243a4

    
           func TestSourceChunksNoResumption(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{"trufflesec-ahrav-test-2"}, 
        
           	} 
        
           	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 := 19787 
        
           	got := 0 
        
           	for range chunksCh { 
        
           		got++ 
        
           	} 
        
           	assert.Equal(t, got, wantChunkCount) 
        
           }

trufflehog/pkg/sources/s3/s3_integration_test.go

Lines 251 to 282 in a1243a4

    
           func TestSourceChunksNoResumption(t *testing.T) { 
        
           	t.Parallel() 
        
           	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30) 
        
           	defer cancel() 
        
           	s := Source{} 
        
           	connection := &sourcespb.S3{ 
        
           		Credential: &sourcespb.S3_Unauthenticated{}, 
        
           		Buckets:    []string{"integration-resumption-tests"}, 
        
           	} 
        
           	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 := 19787 
        
           	got := 0 
        
           	for range chunksCh { 
        
           		got++ 
        
           	} 
        
           	assert.Equal(t, wantChunkCount, got) 
        
           }

Possible infinite block in TestSource_Chunks

As reported in feat(sources/s3): migrate to AWS SDK v2 #4069 (comment), the test TestSource_Chunks in s3_test.go may block indefinitely. This is due to the use of an unbuffered channel (chunksCh), combined with only a single receive operation (gotChunk := <-chunksCh). If the test bucket contains more than one chunk, the s.Chunks(ctx, chunksCh) call will block, causing the test to hang.

See test(sources/s3): fix infinite blocking and timeout issue in TestSource_Chunks #4048 (comment) for full explanation.

trufflehog/pkg/sources/s3/s3_test.go

Lines 120 to 131 in a1243a4

    
           chunksCh := make(chan *sources.Chunk) 
        
           var wg sync.WaitGroup 
        
           wg.Add(1) 
        
           go func() { 
        
           	defer wg.Done() 
        
           	err = s.Chunks(ctx, chunksCh) 
        
           	if (err != nil) != tt.wantErr { 
        
           		t.Errorf("Source.Chunks() error = %v, wantErr %v", err, tt.wantErr) 
        
           		os.Exit(1) 
        
           	} 
        
           }() 
        
           gotChunk := <-chunksCh

Checklist:

Tests passing (make test-community)?
Lint passing (make lint this requires golangci-lint)?

Signed-off-by: Eng Zer Jun <engzerjun@gmail.com>

zricethezav · 2025-04-14T19:06:43Z

@ahrav any idea which one is the test we want to keep?

Signed-off-by: Eng Zer Jun <engzerjun@gmail.com>

Juneezee · 2025-04-22T16:26:07Z

pkg/sources/s3/s3_test.go

+			chunksCh := make(chan *sources.Chunk, 1)
 			go func() {
-				defer wg.Done()
+				defer close(chunksCh)
 				err = s.Chunks(ctx, chunksCh)


/cc @rosecodym

Here is a more detailed explanation of why TestSource_Chunks may block indefinitely.

The issue arises after the call to s.Chunks(ctx, chunksCh). The relevant call stack is:

(*Source).Chunks

trufflehog/pkg/sources/s3/s3.go

Lines 390 to 398 in 6d4ccfa

// Chunks emits chunks of bytes over a channel.

func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk, _ ...sources.ChunkingTarget) error {

visitor := func(c context.Context, defaultRegionClient *s3.S3, roleArn string, buckets []string) error {

s.scanBuckets(c, defaultRegionClient, roleArn, buckets, chunksChan)

return nil

}

return s.visitRoles(ctx, visitor)

}

(*Source).scanBuckets

trufflehog/pkg/sources/s3/s3.go

Line 363 in 6d4ccfa

s.pageChunker(ctx, pageMetadata, processingState, chunksChan)

(*Source).pageChunker

trufflehog/pkg/sources/s3/s3.go

Lines 579 to 583 in 6d4ccfa

if err := handlers.HandleFile(ctx, res.Body, chunkSkel, sources.ChanReporter{Ch: chunksChan}); err != nil {

ctx.Logger().Error(err, "error handling file")

s.metricsCollector.RecordObjectError(metadata.bucket)

return nil

}

HandleFile

trufflehog/pkg/handlers/handlers.go

Line 390 in 6d4ccfa

return handleChunksWithError(processingCtx, dataOrErrChan, chunkSkel, reporter)

handleChunksWithError

trufflehog/pkg/handlers/handlers.go

Lines 426 to 428 in 6d4ccfa

if err := reporter.ChunkOk(ctx, chunk); err != nil {

return fmt.Errorf("error reporting chunk: %w", err)

}

(ChanReporter).ChunkOk

trufflehog/pkg/sources/legacy_reporters.go

Lines 15 to 17 in 6d4ccfa

func (c ChanReporter) ChunkOk(ctx context.Context, chunk Chunk) error {

return common.CancellableWrite(ctx, c.Ch, &chunk)

}

And blocks infinitely in CancellableWrite. Because chunksCh is an unbuffered channel, the ch <- item case can never proceed since we only receive from chunksCh once.

trufflehog/pkg/common/context.go

Lines 23 to 35 in 6d4ccfa

func CancellableWrite[T any](ctx context.Context, ch chan<- T, item T) error {

select {

case <-ctx.Done(): // priority to context cancellation

return ctx.Err()

default:

select {

case <-ctx.Done():

return ctx.Err()

case ch <- item:

return nil

}

}

}

After changing chunksCh to a buffered channel, I created a public S3 bucket and can confirm that the test no longer hangs infinitely:

TestSource_Chunk_buffered_channel.mp4

rosecodym

@Juneezee thank you for your detailed investigation!

actually, i'm seeing some new test failures :/

rosecodym

Thanks for this! Unfortunately, it introduces a new test failure:

=== RUN   TestSource_Chunks/gets_chunks_after_assuming_role
2025-04-23T09:40:41-04:00	info-0	context	Creating checkpointer	{"timeout": 30}
2025-04-23T09:40:42-04:00	info-0	context	Starting new scan from beginning	{"timeout": 30, "role": "arn:aws:iam::619888638459:role/s3-test-assume-role"}
    s3_test.go:138: 
        	Error Trace:	/Users/cody.rose/src/trufflehog/pkg/sources/s3/s3_test.go:138
        	Error:      	Not equal: 
        	            	expected: ""
        	            	actual  : "{\"current_bucket\":\"truffletestbucket-s3-role-assumption\",\"start_after\":\"s\"}"
        	            	
        	            	Diff:
        	            	--- Expected
        	            	+++ Actual
        	            	@@ -1 +1 @@
        	            	-
        	            	+{"current_bucket":"truffletestbucket-s3-role-assumption","start_after":"s"}
        	Test:       	TestSource_Chunks/gets_chunks_after_assuming_role
    s3_test.go:139: 
        	Error Trace:	/Users/cody.rose/src/trufflehog/pkg/sources/s3/s3_test.go:139
        	Error:      	Not equal: 
        	            	expected: 100
        	            	actual  : 90
        	Test:       	TestSource_Chunks/gets_chunks_after_assuming_role
--- FAIL: TestSource_Chunks (6.15s)
    --- FAIL: TestSource_Chunks/gets_chunks_after_assuming_role (5.66s)

The problem is that the test no longer waits for Chunks to return before checking the source's resumption info and progress completion, so there's a mismatch. I believe this can be resolved by re-adding the removed waitgroup functionality.

Reference: #4048 (review) Signed-off-by: Eng Zer Jun <engzerjun@gmail.com>

Juneezee · 2025-04-23T14:24:42Z

pkg/sources/s3/s3_test.go

+			waitFn := func() {
+				receivedFirstChunk := false
+				for {
+					select {
+					case <-ctx.Done():
+						t.Errorf("TestSource_Chunks timed out: %v", ctx.Err())
+					case gotChunk, ok := <-chunksCh:
+						if !ok {
+							t.Logf("Source.Chunks() finished, channel closed")
+							assert.Equal(t, "", s.GetProgress().EncodedResumeInfo)
+							assert.Equal(t, int64(100), s.GetProgress().PercentComplete)
+							return
+						}
+						if receivedFirstChunk {
+							// wantChunkData is the first chunk data. After the first chunk has
+							// been received and matched below, we want to drain chunksCh
+							// so Source.Chunks() can finish completely.
+							continue
+						}
+
+						receivedFirstChunk = true
+						wantData, _ := base64.StdEncoding.DecodeString(tt.wantChunkData)
+
+						if diff := pretty.Compare(gotChunk.Data, wantData); diff != "" {
+							t.Logf("%s: Source.Chunks() diff: (-got +want)\n%s", tt.name, diff)
+						}
+					}
+				}
 			}


@rosecodym Could you try running the test again with the latest commit (1ca8a95)?

I believe this can be resolved by re-adding the removed waitgroup functionality.

Actually, a WaitGroup alone is not sufficient here. For s.Chunks(ctx, chunksCh) to finish completely, the chunksCh channel must be fully drained.

If we use a WaitGroup without draining chunksCh, the test will still block indefinitely, unless chunksCh is a buffered channel with a large enough buffer to hold all the chunks.

Signed-off-by: Eng Zer Jun <engzerjun@gmail.com>

rosecodym · 2025-04-24T13:37:27Z

I see passing tests now!

rosecodym

Thanks for cleaning this up!

test(sources/s3): remove duplicate integration test case

1a7abe0

Signed-off-by: Eng Zer Jun <engzerjun@gmail.com>

Juneezee requested review from a team as code owners April 14, 2025 16:35

Juneezee mentioned this pull request Apr 19, 2025

feat(sources/s3): migrate to AWS SDK v2 #4069

Merged

2 tasks

Juneezee added 2 commits April 22, 2025 22:52

Merge two TestSourceChunksNoResumption into table-driven test

c2d0401

Signed-off-by: Eng Zer Jun <engzerjun@gmail.com>

Fix infinite blocking and timeout issue in TestSource_Chunks

d6655fe

Signed-off-by: Eng Zer Jun <engzerjun@gmail.com>

Juneezee changed the title ~~test(sources/s3): remove duplicate integration test case~~ test(sources/s3): fix infinite blocking and timeout issue in TestSource_Chunks Apr 22, 2025

Juneezee commented Apr 22, 2025

View reviewed changes

rosecodym previously approved these changes Apr 23, 2025

View reviewed changes

rosecodym requested changes Apr 23, 2025

View reviewed changes

Drain chunksCh channel

6e020ff

Reference: #4048 (review) Signed-off-by: Eng Zer Jun <engzerjun@gmail.com>

Juneezee commented Apr 23, 2025

View reviewed changes

Add missing return for ctx.Done() case

1ca8a95

Signed-off-by: Eng Zer Jun <engzerjun@gmail.com>

Juneezee requested a review from rosecodym April 23, 2025 18:08

rosecodym approved these changes Apr 24, 2025

View reviewed changes

rosecodym merged commit a4838d4 into trufflesecurity:main Apr 24, 2025
13 checks passed

blsaccess mentioned this pull request Apr 30, 2025

Update trufflehog to 3.88.31 blacklanternsecurity/bbot#2385

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

test(sources/s3): fix infinite blocking and timeout issue in TestSource_Chunks #4048

test(sources/s3): fix infinite blocking and timeout issue in TestSource_Chunks #4048

Uh oh!

Juneezee commented Apr 14, 2025 •

edited

Loading

Uh oh!

zricethezav commented Apr 14, 2025

Uh oh!

Juneezee Apr 22, 2025

Uh oh!

rosecodym left a comment

Uh oh!

rosecodym left a comment

Uh oh!

Juneezee Apr 23, 2025 •

edited

Loading

Uh oh!

rosecodym commented Apr 24, 2025

Uh oh!

rosecodym left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	func TestSourceChunksNoResumption(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{"trufflesec-ahrav-test-2"},
	}
	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 := 19787
	got := 0

	for range chunksCh {
	got++
	}
	assert.Equal(t, got, wantChunkCount)
	}

	func TestSourceChunksNoResumption(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
	defer cancel()

	s := Source{}
	connection := &sourcespb.S3{
	Credential: &sourcespb.S3_Unauthenticated{},
	Buckets: []string{"integration-resumption-tests"},
	}
	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 := 19787
	got := 0

	for range chunksCh {
	got++
	}
	assert.Equal(t, wantChunkCount, got)
	}

	chunksCh := make(chan *sources.Chunk)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
	defer wg.Done()
	err = s.Chunks(ctx, chunksCh)
	if (err != nil) != tt.wantErr {
	t.Errorf("Source.Chunks() error = %v, wantErr %v", err, tt.wantErr)
	os.Exit(1)
	}
	}()
	gotChunk := <-chunksCh

	// Chunks emits chunks of bytes over a channel.
	func (s Source) Chunks(ctx context.Context, chunksChan chan sources.Chunk, _ ...sources.ChunkingTarget) error {
	visitor := func(c context.Context, defaultRegionClient *s3.S3, roleArn string, buckets []string) error {
	s.scanBuckets(c, defaultRegionClient, roleArn, buckets, chunksChan)
	return nil
	}

	return s.visitRoles(ctx, visitor)
	}

	if err := handlers.HandleFile(ctx, res.Body, chunkSkel, sources.ChanReporter{Ch: chunksChan}); err != nil {
	ctx.Logger().Error(err, "error handling file")
	s.metricsCollector.RecordObjectError(metadata.bucket)
	return nil
	}

	if err := reporter.ChunkOk(ctx, chunk); err != nil {
	return fmt.Errorf("error reporting chunk: %w", err)
	}

	func (c ChanReporter) ChunkOk(ctx context.Context, chunk Chunk) error {
	return common.CancellableWrite(ctx, c.Ch, &chunk)
	}

	func CancellableWrite[T any](ctx context.Context, ch chan<- T, item T) error {
	select {
	case <-ctx.Done(): // priority to context cancellation
	return ctx.Err()
	default:
	select {
	case <-ctx.Done():
	return ctx.Err()
	case ch <- item:
	return nil
	}
	}
	}

test(sources/s3): fix infinite blocking and timeout issue in TestSource_Chunks #4048

test(sources/s3): fix infinite blocking and timeout issue in TestSource_Chunks #4048

Uh oh!

Conversation

Juneezee commented Apr 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description:

Checklist:

Uh oh!

zricethezav commented Apr 14, 2025

Uh oh!

Juneezee Apr 22, 2025

Choose a reason for hiding this comment

Uh oh!

rosecodym left a comment

Choose a reason for hiding this comment

Uh oh!

rosecodym left a comment

Choose a reason for hiding this comment

Uh oh!

Juneezee Apr 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rosecodym commented Apr 24, 2025

Uh oh!

rosecodym left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Juneezee commented Apr 14, 2025 •

edited

Loading

Juneezee Apr 23, 2025 •

edited

Loading