implement SourceUnitUnmarshaller on source with the new S3SourceUnit, add test to test resumption on multiple buckets with concurrent ChunkUnit processing

Author: mustansir14

pkg/sources/s3/s3.go

@@ -1,6 +1,7 @@
 package s3
 
 import (
+	"encoding/json"
 	"fmt"
 	"slices"
 	"strings"
@@ -54,8 +55,6 @@ type Source struct {
 	errorCount    *sync.Map
 	jobPool       *errgroup.Group
 	maxObjectSize int64
-
-	sources.CommonSourceUnitUnmarshaller
 }
 
 // Ensure the Source satisfies the interfaces at compile time
@@ -699,22 +698,6 @@ func makeS3Link(bucket, region, key string) string {
 	return fmt.Sprintf("https://%s.s3.%s.amazonaws.com/%s", bucket, region, key)
 }
 
-type S3SourceUnit struct {
-	Bucket string
-	Role   string
-}
-
-func (s S3SourceUnit) SourceUnitID() (string, sources.SourceUnitKind) {
-	// The ID is the bucket name, and the kind is "s3_bucket".
-	return s.Bucket, "s3_bucket"
-}
-
-func (s S3SourceUnit) Display() string {
-	return s.Bucket
-}
-
-var _ sources.SourceUnit = S3SourceUnit{}
-
 // Enumerate implements SourceUnitEnumerator interface. This implementation visits
 // each configured role and passes each s3 bucket as a source unit
 func (s *Source) Enumerate(ctx context.Context, reporter sources.UnitReporter) error {
@@ -771,3 +754,15 @@ func (s *Source) ChunkUnit(ctx context.Context, unit sources.SourceUnit, reporte
 	s.scanBucket(ctx, defaultClient, s3unit.Role, bucket, reporter, startAfterPtr)
 	return nil
 }
+
+func (s *Source) UnmarshalSourceUnit(data []byte) (sources.SourceUnit, error) {
+	var unit S3SourceUnit
+	if err := json.Unmarshal(data, &unit); err != nil {
+		return nil, err
+	}
+	bucket, kind := unit.SourceUnitID()
+	if bucket == "" || kind != SourceUnitKindBucket {
+		return nil, fmt.Errorf("not an S3SourceUnit")
+	}
+	return unit, nil
+}

pkg/sources/s3/s3_integration_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/credentialspb"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/source_metadatapb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sourcestest"
@@ -394,6 +395,8 @@ func TestSourceChunksResumptionMultipleBucketsIgnoredBucket(t *testing.T) {
 }
 
 func TestSource_Enumerate(t *testing.T) {
+	t.Parallel()
+
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
 	defer cancel()
 
@@ -438,6 +441,8 @@ func TestSource_Enumerate(t *testing.T) {
 }
 
 func TestSource_ChunkUnit(t *testing.T) {
+	t.Parallel()
+
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
 	defer cancel()
 
@@ -484,7 +489,7 @@ func TestSource_ChunkUnit(t *testing.T) {
 func TestSource_ChunkUnit_Resumption(t *testing.T) {
 	t.Parallel()
 
-	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
 	defer cancel()
 
 	s := new(Source)
@@ -517,3 +522,65 @@ func TestSource_ChunkUnit_Resumption(t *testing.T) {
 	// Verify that we processed all remaining data on resume.
 	assert.Equal(t, 9638, len(reporter.Chunks), "Should have processed all remaining data on resume")
 }
+
+// TestSource_ChunkUnit_Resumption_MultipleBucketsConcurrent tests resumption across multiple buckets
+// with concurrent ChunkUnit processing.
+func TestSource_ChunkUnit_Resumption_MultipleBucketsConcurrent(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	src := new(Source)
+	connection := &sourcespb.S3{
+		Credential:       &sourcespb.S3_Unauthenticated{},
+		Buckets:          []string{"integration-resumption-tests", "trufflesec-ahrav-test-2"},
+		EnableResumption: true,
+	}
+	conn, err := anypb.New(connection)
+	require.NoError(t, err)
+	err = src.Init(ctx, "test name", 0, 0, false, conn, 2)
+	require.NoError(t, err)
+
+	src.Progress = sources.Progress{
+		Message:           "Buckets: [integration-resumption-tests trufflesec-ahrav-test-2]",
+		EncodedResumeInfo: "{\"integration-resumption-tests\":\"test-dir\", \"trufflesec-ahrav-test-2\":\"test-dir/smmed_random_data.json.zip\"}",
+		SectionsCompleted: 0,
+		SectionsRemaining: 2,
+	}
+
+	reporter := sourcestest.SafeTestReporter{}
+	err = src.Enumerate(ctx, &reporter)
+	require.NoError(t, err)
+	require.Equal(t, 2, len(reporter.Units), "Expected two source units from enumeration")
+
+	var wg sync.WaitGroup
+
+	for _, unit := range reporter.Units {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			err = src.ChunkUnit(ctx, unit, &reporter)
+			assert.NoError(t, err, "Expected no error during ChunkUnit")
+		}()
+	}
+
+	wg.Wait()
+
+	bucketChunkCounts := map[string]int{
+		"integration-resumption-tests": 9638,
+		"trufflesec-ahrav-test-2":      2,
+	}
+	actualBucketChunkCounts := make(map[string]int)
+
+	for _, chunk := range reporter.Chunks {
+		metadata, _ := chunk.SourceMetadata.Data.(*source_metadatapb.MetaData_S3)
+		actualBucketChunkCounts[metadata.S3.Bucket]++
+	}
+
+	for bucket, wantCount := range bucketChunkCounts {
+		gotCount, ok := actualBucketChunkCounts[bucket]
+		require.True(t, ok, "Expected chunks for bucket %s", bucket)
+		assert.Equal(t, wantCount, gotCount, "Chunk count mismatch for bucket %s", bucket)
+	}
+}

pkg/sources/s3/s3_test.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/kylelemons/godebug/pretty"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/types/known/anypb"
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
@@ -163,3 +164,21 @@ func TestSource_Chunks(t *testing.T) {
 		})
 	}
 }
+
+func TestSource_UnmarshalSourceUnit(t *testing.T) {
+	s := Source{}
+
+	unitJSON := `{
+		"Bucket": "my-test-bucket",
+		"Role": "my-test-role"
+	}`
+
+	unit, err := s.UnmarshalSourceUnit([]byte(unitJSON))
+	require.NoError(t, err, "UnmarshalSourceUnit should not return an error")
+
+	s3Unit, ok := unit.(S3SourceUnit)
+	require.True(t, ok, "Unmarshaled unit should be of type S3SourceUnit")
+
+	assert.Equal(t, "my-test-bucket", s3Unit.Bucket, "Bucket field should match")
+	assert.Equal(t, "my-test-role", s3Unit.Role, "Role field should match")
+}

pkg/sources/s3/unit.go

@@ -0,0 +1,23 @@
+package s3
+
+import (
+	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
+)
+
+const SourceUnitKindBucket sources.SourceUnitKind = "bucket"
+
+type S3SourceUnit struct {
+	Bucket string
+	Role   string
+}
+
+var _ sources.SourceUnit = S3SourceUnit{}
+
+func (s S3SourceUnit) SourceUnitID() (string, sources.SourceUnitKind) {
+	// The ID is the bucket name, and the kind is "bucket".
+	return s.Bucket, SourceUnitKindBucket
+}
+
+func (s S3SourceUnit) Display() string {
+	return s.Bucket
+}

pkg/sourcestest/sourcestest.go

@@ -2,6 +2,7 @@ package sourcestest
 
 import (
 	"fmt"
+	"sync"
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
@@ -59,3 +60,32 @@ func (ErrReporter) ChunkOk(context.Context, sources.Chunk) error {
 func (ErrReporter) ChunkErr(context.Context, error) error {
 	return fmt.Errorf("ErrReporter: ChunkErr error")
 }
+
+// SafeTestReporter is a helper struct that implements both UnitReporter and
+// ChunkReporter by recording the values passed in the methods with thread safety.
+type SafeTestReporter struct {
+	TestReporter
+
+	mu sync.Mutex
+}
+
+func (t *SafeTestReporter) UnitOk(_ context.Context, unit sources.SourceUnit) error {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	return t.TestReporter.UnitOk(nil, unit)
+}
+func (t *SafeTestReporter) UnitErr(_ context.Context, err error) error {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	return t.TestReporter.UnitErr(nil, err)
+}
+func (t *SafeTestReporter) ChunkOk(_ context.Context, chunk sources.Chunk) error {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	return t.TestReporter.ChunkOk(nil, chunk)
+}
+func (t *SafeTestReporter) ChunkErr(_ context.Context, err error) error {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	return t.TestReporter.ChunkErr(nil, err)
+}