[feat] - S3 Progress Tracker (#3568)
* add config option for s3 resumption * updates * initial progress tracking logic * more testing * revert s3 source file * UpdateScanProgress tests * adjust * updates * invert * updates * updates * fix * update * adjust test * fix * remove progress tracking * address comments * add comments
This commit is contained in:
ahrav
@@ -0,0 +1,216 @@
|
||||
package s3
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sync"
|
||||
|
||||
"github.com/aws/aws-sdk-go/service/s3"
|
||||
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/context"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
|
||||
)
|
||||
|
||||
// ProgressTracker maintains scan progress 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.
|
||||
//
|
||||
// S3 buckets are organized as flat namespaces of objects identified by unique keys.
|
||||
// When listing objects, S3 returns paginated results with a maximum of 1000 objects per page.
|
||||
// The ListObjectsV2 API accepts a 'StartAfter' parameter that allows resuming the listing
|
||||
// from a specific object key.
|
||||
//
|
||||
// The tracker maintains state for the current page of objects (up to 1000) using a boolean array
|
||||
// to track completion status and an ordered list to record the sequence of completions.
|
||||
// This enables finding the highest consecutive completed index as a "low water mark".
|
||||
//
|
||||
// The key of the object at this index is encoded with the current bucket into a ResumeInfo checkpoint
|
||||
// and persisted in the Progress.EncodedResumeInfo field as JSON. If a scan is interrupted, it can
|
||||
// resume from the last checkpoint by using that key as StartAfter.
|
||||
//
|
||||
// The low water mark approach ensures scan reliability by only checkpointing consecutively completed
|
||||
// objects. For example, if objects 0-5 and 7-8 are complete but 6 is incomplete, only objects 0-5
|
||||
// will be checkpointed. While this may result in re-scanning objects 7-8 when resuming, it guarantees
|
||||
// no objects are missed in case of interruption.
|
||||
//
|
||||
// When scanning multiple buckets, the current bucket is tracked in the checkpoint to enable
|
||||
// resuming from the correct bucket. The scan will continue from the last checkpointed object
|
||||
// in that bucket.
|
||||
//
|
||||
// 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
|
||||
// On resume: StartAfter=object1499 in saved bucket, scanning continues from object 1500
|
||||
type ProgressTracker struct {
|
||||
enabled bool
|
||||
|
||||
// completedObjects tracks which indices in the current page have been processed.
|
||||
sync.Mutex
|
||||
completedObjects []bool
|
||||
completionOrder []int // Track the order in which objects complete
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
const defaultMaxObjectsPerPage = 1000
|
||||
|
||||
// NewProgressTracker creates a new progress tracker for S3 scanning operations.
|
||||
// The enabled parameter determines if progress tracking is active, and progress
|
||||
// provides the underlying mechanism for persisting scan state.
|
||||
func NewProgressTracker(ctx context.Context, enabled bool, progress *sources.Progress) *ProgressTracker {
|
||||
ctx.Logger().Info("Creating progress tracker")
|
||||
|
||||
return &ProgressTracker{
|
||||
// We are resuming if we have completed objects from a previous scan.
|
||||
completedObjects: make([]bool, defaultMaxObjectsPerPage),
|
||||
completionOrder: make([]int, 0, defaultMaxObjectsPerPage),
|
||||
enabled: enabled,
|
||||
progress: progress,
|
||||
}
|
||||
}
|
||||
|
||||
// Reset prepares the tracker for a new page of objects by clearing the completion state.
|
||||
func (p *ProgressTracker) Reset() {
|
||||
if !p.enabled {
|
||||
return
|
||||
}
|
||||
|
||||
p.Lock()
|
||||
defer p.Unlock()
|
||||
// Store the current completed count before moving to next page.
|
||||
p.completedObjects = make([]bool, defaultMaxObjectsPerPage)
|
||||
p.completionOrder = make([]int, 0, defaultMaxObjectsPerPage)
|
||||
}
|
||||
|
||||
// ResumeInfo represents the state needed to resume an interrupted operation.
|
||||
// It contains the necessary information to continue processing from the last
|
||||
// successfully processed item.
|
||||
type ResumeInfo struct {
|
||||
CurrentBucket string `json:"current_bucket"` // Current bucket being scanned
|
||||
StartAfter string `json:"start_after"` // Last processed object key
|
||||
}
|
||||
|
||||
// GetResumePoint retrieves the last saved checkpoint state if one exists.
|
||||
// It returns nil if progress tracking is disabled or no resume state exists.
|
||||
// This method decodes the stored resume information and validates it contains
|
||||
// the minimum required data to enable resumption.
|
||||
func (p *ProgressTracker) GetResumePoint(ctx context.Context) (ResumeInfo, error) {
|
||||
resume := ResumeInfo{}
|
||||
|
||||
if !p.enabled || p.progress.EncodedResumeInfo == "" {
|
||||
return resume, nil
|
||||
}
|
||||
|
||||
var resumeInfo ResumeInfo
|
||||
if err := json.Unmarshal([]byte(p.progress.EncodedResumeInfo), &resumeInfo); err != nil {
|
||||
return resume, fmt.Errorf("failed to decode resume info: %w", err)
|
||||
}
|
||||
|
||||
if resumeInfo.CurrentBucket == "" {
|
||||
ctx.Logger().V(2).Info("resume info is missing current bucket, resuming from the beginning")
|
||||
return resume, nil
|
||||
}
|
||||
|
||||
return ResumeInfo{CurrentBucket: resumeInfo.CurrentBucket, StartAfter: resumeInfo.StartAfter}, nil
|
||||
}
|
||||
|
||||
// Complete marks the entire scanning operation as finished and clears the resume state.
|
||||
// This should only be called once all scanning operations are complete.
|
||||
func (p *ProgressTracker) Complete(_ context.Context, message string) error {
|
||||
// Preserve existing progress counters while clearing resume state.
|
||||
p.progress.SetProgressComplete(
|
||||
int(p.progress.SectionsCompleted),
|
||||
int(p.progress.SectionsRemaining),
|
||||
message,
|
||||
"", // Clear resume info as scanning is complete
|
||||
)
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateObjectProgress records successfully processed objects within the current page
|
||||
// and maintains fine-grained resumption checkpoints. It uses a conservative tracking
|
||||
// strategy that ensures no objects are missed by only checkpointing consecutively
|
||||
// completed objects.
|
||||
//
|
||||
// This method manages the detailed object-level progress tracking and creates
|
||||
// checkpoints that enable resumption of interrupted scans.
|
||||
//
|
||||
// This approach ensures scan reliability by only checkpointing consecutively completed
|
||||
// objects. While this may result in re-scanning some objects when resuming, it guarantees
|
||||
// no objects are missed in case of interruption.
|
||||
//
|
||||
// For example, consider scanning a page of 10 objects where objects 0-5 and 7-8 complete
|
||||
// successfully but object 6 fails:
|
||||
// - Objects completed: [0,1,2,3,4,5,7,8]
|
||||
// - The checkpoint will only include objects 0-5 since they are consecutive
|
||||
// - If scanning is interrupted and resumed:
|
||||
// - Scan resumes after object 5 (the last checkpoint)
|
||||
// - Objects 7-8 will be re-scanned even though they completed before
|
||||
// - This ensures object 6 is not missed
|
||||
func (p *ProgressTracker) UpdateObjectProgress(
|
||||
ctx context.Context,
|
||||
completedIdx int,
|
||||
bucket string,
|
||||
pageContents []*s3.Object,
|
||||
) error {
|
||||
if !p.enabled {
|
||||
return nil
|
||||
}
|
||||
|
||||
ctx = context.WithValues(ctx, "bucket", bucket, "completedIdx", completedIdx)
|
||||
ctx.Logger().V(5).Info("Updating progress")
|
||||
|
||||
if completedIdx >= len(p.completedObjects) {
|
||||
return fmt.Errorf("completed index %d exceeds maximum page size", completedIdx)
|
||||
}
|
||||
|
||||
p.Lock()
|
||||
defer p.Unlock()
|
||||
|
||||
// Only track completion if this is the first time this index is marked complete.
|
||||
if !p.completedObjects[completedIdx] {
|
||||
p.completedObjects[completedIdx] = true
|
||||
p.completionOrder = append(p.completionOrder, completedIdx)
|
||||
}
|
||||
|
||||
// Find the highest safe checkpoint we can create.
|
||||
lastSafeIdx := -1
|
||||
var safeIndices [defaultMaxObjectsPerPage]bool
|
||||
|
||||
// Mark all completed indices.
|
||||
for _, idx := range p.completionOrder {
|
||||
safeIndices[idx] = true
|
||||
}
|
||||
|
||||
// Find the highest consecutive completed index.
|
||||
for i := range len(p.completedObjects) {
|
||||
if !safeIndices[i] {
|
||||
break
|
||||
}
|
||||
lastSafeIdx = i
|
||||
}
|
||||
|
||||
// Update progress if we have at least one completed object.
|
||||
if lastSafeIdx < 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
obj := pageContents[lastSafeIdx]
|
||||
info := &ResumeInfo{CurrentBucket: bucket, StartAfter: *obj.Key}
|
||||
encoded, err := json.Marshal(info)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Purposefully avoid updating any progress counts.
|
||||
// Only update resume info.
|
||||
p.progress.SetProgressComplete(
|
||||
int(p.progress.SectionsCompleted),
|
||||
int(p.progress.SectionsRemaining),
|
||||
p.progress.Message,
|
||||
string(encoded),
|
||||
)
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,325 @@
|
||||
package s3
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"github.com/aws/aws-sdk-go/service/s3"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/context"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
|
||||
)
|
||||
|
||||
func TestProgressTrackerResumption(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
// First scan - process 6 objects then interrupt.
|
||||
initialProgress := &sources.Progress{}
|
||||
tracker := NewProgressTracker(ctx, true, initialProgress)
|
||||
|
||||
firstPage := &s3.ListObjectsV2Output{
|
||||
Contents: make([]*s3.Object, 12), // Total of 12 objects
|
||||
}
|
||||
for i := range 12 {
|
||||
key := fmt.Sprintf("key-%d", i)
|
||||
firstPage.Contents[i] = &s3.Object{Key: &key}
|
||||
}
|
||||
|
||||
// Process first 6 objects.
|
||||
for i := range 6 {
|
||||
err := tracker.UpdateObjectProgress(ctx, i, "test-bucket", firstPage.Contents)
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
// Verify resume info is set correctly.
|
||||
resumeInfo, err := tracker.GetResumePoint(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "test-bucket", resumeInfo.CurrentBucket)
|
||||
assert.Equal(t, "key-5", resumeInfo.StartAfter)
|
||||
|
||||
// Resume scan with existing progress.
|
||||
resumeTracker := NewProgressTracker(ctx, true, initialProgress)
|
||||
|
||||
resumePage := &s3.ListObjectsV2Output{
|
||||
Contents: firstPage.Contents[6:], // Remaining 6 objects
|
||||
}
|
||||
|
||||
// Process remaining objects.
|
||||
for i := range len(resumePage.Contents) {
|
||||
err := resumeTracker.UpdateObjectProgress(ctx, i, "test-bucket", resumePage.Contents)
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
// Verify final resume info.
|
||||
finalResumeInfo, err := resumeTracker.GetResumePoint(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "test-bucket", finalResumeInfo.CurrentBucket)
|
||||
assert.Equal(t, "key-11", finalResumeInfo.StartAfter)
|
||||
}
|
||||
|
||||
func TestProgressTrackerReset(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
enabled bool
|
||||
}{
|
||||
{name: "reset with enabled tracker", enabled: true},
|
||||
{name: "reset with disabled tracker", enabled: false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := context.Background()
|
||||
progress := new(sources.Progress)
|
||||
tracker := NewProgressTracker(ctx, tt.enabled, progress)
|
||||
|
||||
tracker.completedObjects[1] = true
|
||||
tracker.completedObjects[2] = true
|
||||
|
||||
tracker.Reset()
|
||||
|
||||
assert.Equal(t, defaultMaxObjectsPerPage, len(tracker.completedObjects),
|
||||
"Reset changed the length of completed objects")
|
||||
|
||||
if tt.enabled {
|
||||
// All values should be false after reset.
|
||||
for i, isCompleted := range tracker.completedObjects {
|
||||
assert.False(t, isCompleted,
|
||||
"Reset did not clear completed object at index %d", i)
|
||||
}
|
||||
|
||||
// Completion order should be empty.
|
||||
assert.Equal(t, 0, len(tracker.completionOrder),
|
||||
"Reset did not clear completion order")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetResumePoint(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
enabled bool
|
||||
progress *sources.Progress
|
||||
expectedResumeInfo ResumeInfo
|
||||
expectError bool
|
||||
}{
|
||||
{
|
||||
name: "valid resume info",
|
||||
enabled: true,
|
||||
progress: &sources.Progress{
|
||||
EncodedResumeInfo: `{"current_bucket":"test-bucket","start_after":"test-key"}`,
|
||||
},
|
||||
expectedResumeInfo: ResumeInfo{CurrentBucket: "test-bucket", StartAfter: "test-key"},
|
||||
},
|
||||
{
|
||||
name: "progress disabled",
|
||||
enabled: false,
|
||||
progress: &sources.Progress{
|
||||
EncodedResumeInfo: `{"current_bucket":"test-bucket","start_after":"test-key"}`,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "empty encoded resume info",
|
||||
enabled: true,
|
||||
progress: &sources.Progress{EncodedResumeInfo: ""},
|
||||
},
|
||||
{
|
||||
name: "empty current bucket",
|
||||
enabled: true,
|
||||
progress: &sources.Progress{
|
||||
EncodedResumeInfo: `{"current_bucket":"","start_after":"test-key"}`,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "unmarshal error",
|
||||
enabled: true,
|
||||
progress: &sources.Progress{
|
||||
EncodedResumeInfo: `{"current_bucket":123,"start_after":"test-key"}`, // Invalid JSON
|
||||
},
|
||||
expectedResumeInfo: ResumeInfo{},
|
||||
expectError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tracker := &ProgressTracker{enabled: tt.enabled, progress: tt.progress}
|
||||
|
||||
resumePoint, err := tracker.GetResumePoint(context.Background())
|
||||
if tt.expectError {
|
||||
assert.Error(t, err, "Expected an error decoding resume info")
|
||||
} else {
|
||||
assert.NoError(t, err, "Expected no error decoding resume info")
|
||||
}
|
||||
|
||||
assert.Equal(t, tt.expectedResumeInfo, resumePoint, "Unexpected resume point")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestProgressTrackerUpdateProgress(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
description string
|
||||
completedIdx int
|
||||
pageSize int
|
||||
preCompleted []int
|
||||
expectedKey string
|
||||
}{
|
||||
{
|
||||
name: "first object completed",
|
||||
description: "Basic case - completing first object",
|
||||
completedIdx: 0,
|
||||
pageSize: 3,
|
||||
expectedKey: "key-0",
|
||||
},
|
||||
{
|
||||
name: "completing missing middle",
|
||||
description: "Completing object when previous is done",
|
||||
completedIdx: 1,
|
||||
pageSize: 3,
|
||||
preCompleted: []int{0},
|
||||
expectedKey: "key-1",
|
||||
},
|
||||
{
|
||||
name: "all objects completed in order",
|
||||
description: "Completing final object in sequence",
|
||||
completedIdx: 2,
|
||||
pageSize: 3,
|
||||
preCompleted: []int{0, 1},
|
||||
expectedKey: "key-2",
|
||||
},
|
||||
{
|
||||
name: "last index in max page",
|
||||
description: "Edge case - maximum page size boundary",
|
||||
completedIdx: 999,
|
||||
pageSize: 1000,
|
||||
preCompleted: func() []int {
|
||||
indices := make([]int, 999)
|
||||
for i := range 999 {
|
||||
indices[i] = i
|
||||
}
|
||||
return indices
|
||||
}(),
|
||||
expectedKey: "key-999",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := context.Background()
|
||||
progress := new(sources.Progress)
|
||||
tracker := &ProgressTracker{
|
||||
enabled: true,
|
||||
progress: progress,
|
||||
completedObjects: make([]bool, tt.pageSize),
|
||||
completionOrder: make([]int, 0, tt.pageSize),
|
||||
}
|
||||
|
||||
page := &s3.ListObjectsV2Output{Contents: make([]*s3.Object, tt.pageSize)}
|
||||
for i := range tt.pageSize {
|
||||
key := fmt.Sprintf("key-%d", i)
|
||||
page.Contents[i] = &s3.Object{Key: &key}
|
||||
}
|
||||
|
||||
// Apply pre-completed indices in order.
|
||||
if tt.preCompleted != nil {
|
||||
for _, idx := range tt.preCompleted {
|
||||
tracker.completedObjects[idx] = true
|
||||
tracker.completionOrder = append(tracker.completionOrder, idx)
|
||||
}
|
||||
}
|
||||
|
||||
err := tracker.UpdateObjectProgress(ctx, tt.completedIdx, "test-bucket", page.Contents)
|
||||
assert.NoError(t, err, "Unexpected error updating progress")
|
||||
|
||||
var info ResumeInfo
|
||||
err = json.Unmarshal([]byte(progress.EncodedResumeInfo), &info)
|
||||
assert.NoError(t, err, "Failed to decode resume info")
|
||||
assert.Equal(t, tt.expectedKey, info.StartAfter, "Incorrect resume point")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestComplete(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
enabled bool
|
||||
initialState struct {
|
||||
resumeInfo string
|
||||
message string
|
||||
}
|
||||
completeMessage string
|
||||
wantState struct {
|
||||
resumeInfo string
|
||||
message string
|
||||
}
|
||||
}{
|
||||
{
|
||||
name: "marks completion with existing resume info",
|
||||
enabled: true,
|
||||
initialState: struct {
|
||||
resumeInfo string
|
||||
message string
|
||||
}{
|
||||
resumeInfo: `{"current_bucket":"test-bucket","start_after":"some-key"}`,
|
||||
message: "In progress",
|
||||
},
|
||||
completeMessage: "Scan complete",
|
||||
wantState: struct {
|
||||
resumeInfo string
|
||||
message string
|
||||
}{
|
||||
resumeInfo: "", // Should clear resume info
|
||||
message: "Scan complete",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "disabled tracker",
|
||||
enabled: false,
|
||||
initialState: struct {
|
||||
resumeInfo string
|
||||
message string
|
||||
}{
|
||||
resumeInfo: "",
|
||||
message: "Should not change",
|
||||
},
|
||||
completeMessage: "Completed",
|
||||
wantState: struct {
|
||||
resumeInfo string
|
||||
message string
|
||||
}{
|
||||
resumeInfo: "",
|
||||
message: "Completed",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := context.Background()
|
||||
progress := &sources.Progress{
|
||||
EncodedResumeInfo: tt.initialState.resumeInfo,
|
||||
Message: tt.initialState.message,
|
||||
}
|
||||
tracker := NewProgressTracker(ctx, tt.enabled, progress)
|
||||
|
||||
err := tracker.Complete(ctx, tt.completeMessage)
|
||||
assert.NoError(t, err)
|
||||
|
||||
assert.Equal(t, tt.wantState.resumeInfo, progress.EncodedResumeInfo)
|
||||
assert.Equal(t, tt.wantState.message, progress.Message)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -14,7 +14,6 @@ import (
|
||||
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/common"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/credentialspb"
|
||||
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/context"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
|
||||
|
||||
Reference in New Issue
Block a user