[INS-104] Support units in S3 source

pkg/sources/s3/s3.go

@@ -62,6 +62,7 @@ type Source struct {
 var _ sources.Source = (*Source)(nil)
 var _ sources.SourceUnitUnmarshaller = (*Source)(nil)
 var _ sources.Validator = (*Source)(nil)
+var _ sources.SourceUnitEnumChunker = (*Source)(nil)
 
 // Type returns the type of source
 func (s *Source) Type() sourcespb.SourceType { return SourceType }
@@ -316,16 +317,7 @@ func (s *Source) scanBuckets(
 
 	bucketsToScanCount := len(bucketsToScan)
 	for bucketIdx := pos.index; bucketIdx < bucketsToScanCount; bucketIdx++ {
-		s.metricsCollector.RecordBucketForRole(role)
 		bucket := bucketsToScan[bucketIdx]
-		ctx := context.WithValue(ctx, "bucket", bucket)
-
-		if common.IsDone(ctx) {
-			ctx.Logger().Error(ctx.Err(), "context done, while scanning bucket")
-			return
-		}
-
-		ctx.Logger().V(3).Info("Scanning bucket")
 
 		s.SetProgressComplete(
 			bucketIdx,
@@ -334,53 +326,16 @@ func (s *Source) scanBuckets(
 			s.Progress.EncodedResumeInfo,
 		)
 
-		regionalClient, err := s.getRegionalClientForBucket(ctx, client, role, bucket)
-		if err != nil {
-			ctx.Logger().Error(err, "could not get regional client for bucket")
-			continue
-		}
-
-		errorCount := sync.Map{}
-
-		input := &s3.ListObjectsV2Input{Bucket: &bucket}
+		var startAfter *string
 		if bucket == pos.bucket && pos.startAfter != "" {
-			input.StartAfter = &pos.startAfter
+			startAfter = &pos.startAfter
 			ctx.Logger().V(3).Info(
 				"Resuming bucket scan",
 				"start_after", pos.startAfter,
 			)
 		}
 
-		pageNumber := 1
-		paginator := s3.NewListObjectsV2Paginator(regionalClient, input)
-		for paginator.HasMorePages() {
-			output, err := paginator.NextPage(ctx)
-			if err != nil {
-				if role == "" {
-					ctx.Logger().Error(err, "could not list objects in bucket")
-				} else {
-					// Our documentation blesses specifying a role to assume without specifying buckets to scan, which will
-					// often cause this to happen a lot (because in that case the scanner tries to scan every bucket in the
-					// account, but the role probably doesn't have access to all of them). This makes it expected behavior
-					// and therefore not an error.
-					ctx.Logger().V(3).Info("could not list objects in bucket", "err", err)
-				}
-				break
-			}
-			pageMetadata := pageMetadata{
-				bucket:     bucket,
-				pageNumber: pageNumber,
-				client:     regionalClient,
-				page:       output,
-			}
-			processingState := processingState{
-				errorCount:  &errorCount,
-				objectCount: &objectCount,
-			}
-			s.pageChunker(ctx, pageMetadata, processingState, chunksChan)
-
-			pageNumber++
-		}
+		s.scanBucket(ctx, client, role, bucket, sources.ChanReporter{Ch: chunksChan}, startAfter, &objectCount)
 	}
 
 	s.SetProgressComplete(
@@ -391,6 +346,75 @@ func (s *Source) scanBuckets(
 	)
 }
 
+func (s *Source) scanBucket(
+	ctx context.Context,
+	client *s3.Client,
+	role string,
+	bucket string,
+	reporter sources.ChunkReporter,
+	startAfter *string,
+	objectCount *uint64,
+) {
+	s.metricsCollector.RecordBucketForRole(role)
+
+	ctx = context.WithValue(ctx, "bucket", bucket)
+
+	if common.IsDone(ctx) {
+		ctx.Logger().Error(ctx.Err(), "context done, while scanning bucket")
+		return
+	}
+
+	ctx.Logger().V(3).Info("Scanning bucket")
+
+	regionalClient, err := s.getRegionalClientForBucket(ctx, client, role, bucket)
+	if err != nil {
+		ctx.Logger().Error(err, "could not get regional client for bucket")
+		return
+	}
+
+	errorCount := sync.Map{}
+
+	input := &s3.ListObjectsV2Input{Bucket: &bucket}
+	if startAfter != nil {
+		input.StartAfter = startAfter
+	}
+
+	pageNumber := 1
+	paginator := s3.NewListObjectsV2Paginator(regionalClient, input)
+	if objectCount == nil {
+		var newObjectCount uint64
+		objectCount = &newObjectCount
+	}
+	for paginator.HasMorePages() {
+		output, err := paginator.NextPage(ctx)
+		if err != nil {
+			if role == "" {
+				ctx.Logger().Error(err, "could not list objects in bucket")
+			} else {
+				// Our documentation blesses specifying a role to assume without specifying buckets to scan, which will
+				// often cause this to happen a lot (because in that case the scanner tries to scan every bucket in the
+				// account, but the role probably doesn't have access to all of them). This makes it expected behavior
+				// and therefore not an error.
+				ctx.Logger().V(3).Info("could not list objects in bucket", "err", err)
+			}
+			break
+		}
+		pageMetadata := pageMetadata{
+			bucket:     bucket,
+			pageNumber: pageNumber,
+			client:     regionalClient,
+			page:       output,
+		}
+		processingState := processingState{
+			errorCount:  &errorCount,
+			objectCount: objectCount,
+		}
+		s.pageChunker(ctx, pageMetadata, processingState, reporter)
+
+		pageNumber++
+	}
+}
+
 // 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.Client, roleArn string, buckets []string) error {
@@ -429,14 +453,12 @@ func (s *Source) pageChunker(
 	ctx context.Context,
 	metadata pageMetadata,
 	state processingState,
-	chunksChan chan *sources.Chunk,
+	reporter sources.ChunkReporter,
 ) {
 	s.checkpointer.Reset() // Reset the checkpointer for each PAGE
 	ctx = context.WithValues(ctx, "bucket", metadata.bucket, "page_number", metadata.pageNumber)
-
 	for objIdx, obj := range metadata.page.Contents {
 		ctx = context.WithValues(ctx, "key", *obj.Key, "size", *obj.Size)
-
 		if common.IsDone(ctx) {
 			return
 		}
@@ -572,12 +594,11 @@ func (s *Source) pageChunker(
 				Verify: s.verify,
 			}
 
-			if err := handlers.HandleFile(ctx, res.Body, chunkSkel, sources.ChanReporter{Ch: chunksChan}); err != nil {
+			if err := handlers.HandleFile(ctx, res.Body, chunkSkel, reporter); err != nil {
 				ctx.Logger().Error(err, "error handling file")
 				s.metricsCollector.RecordObjectError(metadata.bucket)
 				return nil
 			}
-
 			atomic.AddUint64(state.objectCount, 1)
 			ctx.Logger().V(5).Info("S3 object scanned.", "object_count", state.objectCount)
 			nErr, ok = state.errorCount.Load(prefix)
@@ -587,17 +608,14 @@ func (s *Source) pageChunker(
 			if nErr.(int) > 0 {
 				state.errorCount.Store(prefix, 0)
 			}
-
 			// Update progress after successful processing.
 			if err := s.checkpointer.UpdateObjectCompletion(ctx, objIdx, metadata.bucket, metadata.page.Contents); err != nil {
 				ctx.Logger().Error(err, "could not update progress for scanned object")
 			}
 			s.metricsCollector.RecordObjectScanned(metadata.bucket, float64(*obj.Size))
-
 			return nil
 		})
 	}
-
 	_ = s.jobPool.Wait()
 }
 
@@ -681,3 +699,62 @@ func (s *Source) visitRoles(
 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 {
+	visitor := func(c context.Context, defaultRegionClient *s3.Client, roleArn string, buckets []string) error {
+		for _, bucket := range buckets {
+			if common.IsDone(ctx) {
+				return ctx.Err()
+			}
+
+			ctx.Logger().V(5).Info("Enumerating bucket", "bucket", bucket)
+
+			unit := S3SourceUnit{
+				Bucket: bucket,
+				Role:   roleArn,
+			}
+
+			if err := reporter.UnitOk(ctx, unit); err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+
+	return s.visitRoles(ctx, visitor)
+}
+
+func (s *Source) ChunkUnit(ctx context.Context, unit sources.SourceUnit, reporter sources.ChunkReporter) error {
+
+	s3unit, ok := unit.(S3SourceUnit)
+	if !ok {
+		return fmt.Errorf("expected *S3SourceUnit, got %T", unit)
+	}
+	bucket := s3unit.Bucket
+
+	defaultClient, err := s.newClient(ctx, defaultAWSRegion, s3unit.Role)
+	if err != nil {
+		return fmt.Errorf("could not create s3 client: %w", err)
+	}
+
+	s.scanBucket(ctx, defaultClient, s3unit.Role, bucket, reporter, nil, nil)
+	return nil
+}

pkg/sources/s3/s3_integration_test.go

@@ -18,6 +18,7 @@ import (
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/credentialspb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/sourcestest"
 )
 
 func TestSource_ChunksCount(t *testing.T) {
@@ -391,3 +392,91 @@ func TestSourceChunksResumptionMultipleBucketsIgnoredBucket(t *testing.T) {
 
 	assert.Equal(t, 103, count, "Should have processed all remaining data on resume")
 }
+
+func TestSource_Enumerate(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
+	defer cancel()
+
+	secret, err := common.GetTestSecret(ctx)
+	if err != nil {
+		t.Fatal(fmt.Errorf("failed to access secret: %v", err))
+	}
+
+	s3key := secret.MustGetField("AWS_S3_KEY")
+	s3secret := secret.MustGetField("AWS_S3_SECRET")
+
+	connection := &sourcespb.S3{
+		Credential: &sourcespb.S3_AccessKey{
+			AccessKey: &credentialspb.KeySecret{
+				Key:    s3key,
+				Secret: s3secret,
+			},
+		},
+		Buckets: []string{"truffletestbucket"},
+	}
+
+	conn, err := anypb.New(connection)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	s := Source{}
+	err = s.Init(ctx, "test enumerate", 0, 0, false, conn, 1)
+	assert.NoError(t, err)
+
+	reporter := sourcestest.TestReporter{}
+	err = s.Enumerate(ctx, &reporter)
+	assert.NoError(t, err)
+
+	assert.Equal(t, len(reporter.Units), 1)
+	assert.Equal(t, 0, len(reporter.UnitErrs), "Expected no errors during enumeration")
+
+	for _, unit := range reporter.Units {
+		id, _ := unit.SourceUnitID()
+		assert.NotEmpty(t, id, "Unit ID should not be empty")
+	}
+}
+
+func TestSource_ChunkUnit(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
+	defer cancel()
+
+	secret, err := common.GetTestSecret(ctx)
+	if err != nil {
+		t.Fatal(fmt.Errorf("failed to access secret: %v", err))
+	}
+
+	s3key := secret.MustGetField("AWS_S3_KEY")
+	s3secret := secret.MustGetField("AWS_S3_SECRET")
+
+	connection := &sourcespb.S3{
+		Credential: &sourcespb.S3_AccessKey{
+			AccessKey: &credentialspb.KeySecret{
+				Key:    s3key,
+				Secret: s3secret,
+			},
+		},
+		Buckets: []string{"truffletestbucket"},
+	}
+
+	conn, err := anypb.New(connection)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	s := Source{}
+	err = s.Init(ctx, "test enumerate", 0, 0, false, conn, 1)
+	assert.NoError(t, err)
+
+	reporter := sourcestest.TestReporter{}
+	err = s.Enumerate(ctx, &reporter)
+	assert.NoError(t, err)
+
+	for _, unit := range reporter.Units {
+		err = s.ChunkUnit(ctx, unit, &reporter)
+		assert.NoError(t, err, "Expected no error during ChunkUnit")
+	}
+
+	assert.Equal(t, 103, len(reporter.Chunks))
+	assert.Equal(t, 0, len(reporter.ChunkErrs))
+}