[INS-104] Support units in S3 source (#4560)

* implemented Source unit for S3.
Implemented. integration test

* use bucket as source unit

* remove code duplication, reuse from Chunks

* remove unnecessary change

* remove unused functions

* revisit tests

* revert unnecessary change

* change SourceUnitKind to s3_bucket

* handle nil objectCount inside scanBucket

* handle nil objectCount outside loop

* add bucket to resume log

* add bucket and role to error log, remove enumerating log

* implement sub unit resumption

* add comment to checkpointer for unit scans

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

* add role to SourceUnitID

* Revert "add role to SourceUnitID"

This reverts commit 549e6be.

* add role to source unit ID, keep track of resumption using source unit ID instead of just bucket name

* rename bucket -> unitID in UnmarshalSourceUnit

---------

Co-authored-by: Amaan Ullah <aman.ullah.jalal@trufflesec.com>

Author: mustansir14

pkg/sources/s3/checkpointer.go

@@ -11,6 +11,8 @@ import (
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
 )
 
+// TODO [INS-207] Add role to legacy scan resumption info
+
 // Checkpointer maintains resumption state for S3 bucket scanning,
 // enabling resumable scans by tracking which objects have been successfully processed.
 // It provides checkpoints that can be used to resume interrupted scans without missing objects.
@@ -33,6 +35,10 @@ import (
 // resuming from the correct bucket. The scan will continue from the last checkpointed object
 // in that bucket.
 //
+// Unit scans are also supported. The encoded resume info in this case tracks the last processed object
+// for each unit separately by using the SetEncodedResumeInfoFor method on Progress. To use the
+// checkpointer for unit scans, call SetIsUnitScan(true) before starting the scan.
+//
 // For example, if scanning is interrupted after processing 1500 objects across 2 pages:
 // Page 1 (objects 0-999): Fully processed, checkpoint saved at object 999
 // Page 2 (objects 1000-1999): Partially processed through 1600, but only consecutive through 1499
@@ -56,6 +62,8 @@ type Checkpointer struct {
 	// progress holds the scan's overall progress state and enables persistence.
 	// The EncodedResumeInfo field stores the JSON-encoded ResumeInfo checkpoint.
 	progress *sources.Progress // Reference to source's Progress
+
+	isUnitScan bool // Indicates if scanning is done in unit scan mode
 }
 
 const defaultMaxObjectsPerPage = 1000
@@ -153,9 +161,10 @@ func (p *Checkpointer) UpdateObjectCompletion(
 	ctx context.Context,
 	completedIdx int,
 	bucket string,
+	role string,
 	pageContents []s3types.Object,
 ) error {
-	ctx = context.WithValues(ctx, "bucket", bucket, "completedIdx", completedIdx)
+	ctx = context.WithValues(ctx, "bucket", bucket, "role", role, "completedIdx", completedIdx)
 	ctx.Logger().V(5).Info("Updating progress")
 
 	if completedIdx >= len(p.completedObjects) {
@@ -184,7 +193,7 @@ func (p *Checkpointer) UpdateObjectCompletion(
 	}
 	obj := pageContents[checkpointIdx]
 
-	return p.updateCheckpoint(bucket, *obj.Key)
+	return p.updateCheckpoint(bucket, role, *obj.Key)
 }
 
 // advanceLowestIncompleteIdx moves the lowest incomplete index forward to the next incomplete object.
@@ -198,7 +207,14 @@ func (p *Checkpointer) advanceLowestIncompleteIdx() {
 
 // updateCheckpoint persists the current resumption state.
 // Must be called with lock held.
-func (p *Checkpointer) updateCheckpoint(bucket string, lastKey string) error {
+func (p *Checkpointer) updateCheckpoint(bucket string, role string, lastKey string) error {
+	if p.isUnitScan {
+		unitID := constructS3SourceUnitID(bucket, role)
+		// track sub-unit resumption state
+		p.progress.SetEncodedResumeInfoFor(unitID, lastKey)
+		return nil
+	}
+
 	encoded, err := json.Marshal(&ResumeInfo{CurrentBucket: bucket, StartAfter: lastKey})
 	if err != nil {
 		return fmt.Errorf("failed to encode resume info: %w", err)
@@ -212,3 +228,11 @@ func (p *Checkpointer) updateCheckpoint(bucket string, lastKey string) error {
 	)
 	return nil
 }
+
+// SetIsUnitScan sets whether the checkpointer is operating in unit scan mode.
+func (p *Checkpointer) SetIsUnitScan(isUnitScan bool) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	p.isUnitScan = isUnitScan
+}

pkg/sources/s3/checkpointer_test.go

@@ -31,7 +31,7 @@ func TestCheckpointerResumption(t *testing.T) {
 
 	// Process first 6 objects.
 	for i := range 6 {
-		err := tracker.UpdateObjectCompletion(ctx, i, "test-bucket", firstPage.Contents)
+		err := tracker.UpdateObjectCompletion(ctx, i, "test-bucket", "", firstPage.Contents)
 		assert.NoError(t, err)
 	}
 
@@ -50,7 +50,7 @@ func TestCheckpointerResumption(t *testing.T) {
 
 	// Process remaining objects.
 	for i := range len(resumePage.Contents) {
-		err := resumeTracker.UpdateObjectCompletion(ctx, i, "test-bucket", resumePage.Contents)
+		err := resumeTracker.UpdateObjectCompletion(ctx, i, "test-bucket", "", resumePage.Contents)
 		assert.NoError(t, err)
 	}
 
@@ -244,7 +244,7 @@ func TestCheckpointerUpdate(t *testing.T) {
 				}
 			}
 
-			err := tracker.UpdateObjectCompletion(ctx, tt.completedIdx, "test-bucket", page.Contents)
+			err := tracker.UpdateObjectCompletion(ctx, tt.completedIdx, "test-bucket", "", page.Contents)
 			assert.NoError(t, err, "Unexpected error updating progress")
 
 			var info ResumeInfo
@@ -258,6 +258,36 @@ func TestCheckpointerUpdate(t *testing.T) {
 	}
 }
 
+func TestCheckpointerUpdateUnitScan(t *testing.T) {
+	ctx := context.Background()
+	progress := new(sources.Progress)
+	tracker := NewCheckpointer(ctx, progress)
+	tracker.SetIsUnitScan(true)
+
+	page := &s3.ListObjectsV2Output{
+		Contents: make([]s3types.Object, 3),
+	}
+	for i := range 3 {
+		key := fmt.Sprintf("key-%d", i)
+		page.Contents[i] = s3types.Object{Key: &key}
+	}
+
+	// Complete first object.
+	err := tracker.UpdateObjectCompletion(ctx, 0, "test-bucket", "test-role", page.Contents)
+	assert.NoError(t, err, "Unexpected error updating progress")
+
+	var info map[string]string
+	err = json.Unmarshal([]byte(progress.EncodedResumeInfo), &info)
+	var gotUnitID, gotStartAfter string
+	for k, v := range info {
+		gotUnitID = k
+		gotStartAfter = v
+	}
+	assert.NoError(t, err, "Failed to decode resume info")
+	assert.Equal(t, "test-role|test-bucket", gotUnitID, "Incorrect unit ID")
+	assert.Equal(t, "key-0", gotStartAfter, "Incorrect resume point")
+}
+
 func TestComplete(t *testing.T) {
 	tests := []struct {
 		name         string