trufflesecurity/trufflehog
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Keep filtered objects in the S3 progress total, since the scan counts them as done

Browse Source
This commit is contained in:
Shahzad Haider
2026-09-19 00:14:36 +05:00
parent 76d47d6867
commit 0c529863a1
5 changed files with 59 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/sources/s3/progress.go
+15 -8
View File
@@ -60,7 +60,7 @@ func (s *Source) countBucket(ctx context.Context, client *s3.Client, bucket stri
// Publish after the decrement, because percent reports nothing while a count is in flight. Without it
// the bar stays at 0 until the next object finishes, which for large objects is minutes.
defer func() {
s.progress.listingsInFlight.Add(-1)
s.objectProgress.listingsInFlight.Add(-1)
s.publishProgress()
}()
@@ -72,7 +72,7 @@ func (s *Source) countBucket(ctx context.Context, client *s3.Client, bucket stri
if ctx.Err() != nil {
return
}
s.progress.countFailed.Store(true)
s.objectProgress.countFailed.Store(true)
ctx.Logger().V(2).Info("could not count objects for progress", "bucket", bucket, "err", err)
return
}
@@ -90,15 +90,22 @@ func (s *Source) countBucket(ctx context.Context, client *s3.Client, bucket stri
doneBytes += size
}
}
s.progress.addTotal(objects, bytes)
s.progress.addDone(doneObjects, doneBytes)
s.objectProgress.addTotal(objects, bytes)
s.objectProgress.addDone(doneObjects, doneBytes)
}
}
// countsTowardProgress reports whether an object belongs in the progress ratio. Objects this scan will
// never download are left out of both the total and the done count, so that the percent tracks the bytes
// actually fetched rather than jumping whenever a skipped object goes by.
//
// The order of the tests mirrors pageChunker, which checks the object filter before anything else and
// counts what the filter skips as done. A filtered object therefore belongs in the total as well,
// whatever its storage class or size, or the done count would climb past a total it was never in.
func (s *Source) countsTowardProgress(obj s3types.Object) bool {
if !s.objectFilter.shouldInclude(*obj.Key) {
return true
}
if obj.StorageClass == s3types.ObjectStorageClassGlacier || obj.StorageClass == s3types.ObjectStorageClassGlacierIr {
return false
}
@@ -107,15 +114,15 @@ func (s *Source) countsTowardProgress(obj s3types.Object) bool {
// objectDone records that an object has been scanned or skipped.
func (s *Source) objectDone(size int64) {
s.progress.addDone(1, uint64(max(size, 0)))
s.objectProgress.addDone(1, uint64(max(size, 0)))
s.publishProgress()
}
func (s *Source) publishProgress() {
s.SetProgressPercent(
s.progress.percent(),
clampInt32(s.progress.objectsDone.Load()),
clampInt32(s.progress.objectsTotal.Load()),
s.objectProgress.percent(),
clampInt32(s.objectProgress.objectsDone.Load()),
clampInt32(s.objectProgress.objectsTotal.Load()),
progressMessage,
)
}
pkg/sources/s3/progress_test.go
+30 -13
View File
@@ -5,31 +5,48 @@ import (
s3types "github.com/aws/aws-sdk-go-v2/service/s3/types"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestSource_CountsTowardProgress(t *testing.T) {
s := Source{maxObjectSize: 1024}
size := int64(512)
oversize := int64(2048)
filter, err := newObjectFilter(nil, []string{"archive/"}, nil, nil)
require.NoError(t, err)
s := Source{maxObjectSize: 1024, objectFilter: filter}
key, filteredKey := "keep/obj.txt", "archive/obj.txt"
size, oversize := int64(512), int64(2048)
tests := []struct {
name string
obj s3types.Object
want bool
}{
{name: "standard object", obj: s3types.Object{Size: &size}, want: true},
{name: "at the size limit", obj: s3types.Object{Size: &s.maxObjectSize}, want: true},
{name: "over the size limit", obj: s3types.Object{Size: &oversize}, want: false},
{name: "standard object", obj: s3types.Object{Key: &key, Size: &size}, want: true},
{name: "at the size limit", obj: s3types.Object{Key: &key, Size: &s.maxObjectSize}, want: true},
{name: "over the size limit", obj: s3types.Object{Key: &key, Size: &oversize}, want: false},
{
name: "glacier",
obj: s3types.Object{Size: &size, StorageClass: s3types.ObjectStorageClassGlacier},
obj: s3types.Object{Key: &key, Size: &size, StorageClass: s3types.ObjectStorageClassGlacier},
want: false,
},
{
name: "glacier instant retrieval",
obj: s3types.Object{Size: &size, StorageClass: s3types.ObjectStorageClassGlacierIr},
obj: s3types.Object{Key: &key, Size: &size, StorageClass: s3types.ObjectStorageClassGlacierIr},
want: false,
},
// pageChunker counts a filtered object as done before it looks at storage class or size, so the
// total has to include it too.
{name: "filtered", obj: s3types.Object{Key: &filteredKey, Size: &size}, want: true},
{
name: "filtered and over the size limit",
obj: s3types.Object{Key: &filteredKey, Size: &oversize},
want: true,
},
{
name: "filtered and glacier",
obj: s3types.Object{Key: &filteredKey, Size: &size, StorageClass: s3types.ObjectStorageClassGlacier},
want: true,
},
}
for _, tt := range tests {
@@ -104,17 +121,17 @@ func TestScanProgress_Percent(t *testing.T) {
}
func TestSource_PublishProgressAfterCountPass(t *testing.T) {
s := Source{progress: &scanProgress{}}
s.progress.listingsInFlight.Add(1)
s.progress.addTotal(4, 4000)
s.progress.addDone(2, 2000)
s := Source{objectProgress: &scanProgress{}}
s.objectProgress.listingsInFlight.Add(1)
s.objectProgress.addTotal(4, 4000)
s.objectProgress.addDone(2, 2000)
// A publish while the count is still running cannot know the total yet.
s.publishProgress()
assert.Zero(t, s.GetProgress().PercentComplete)
// The count pass publishes once it has finished, so the bar does not wait for the next object.
s.progress.listingsInFlight.Add(-1)
s.objectProgress.listingsInFlight.Add(-1)
s.publishProgress()
assert.EqualValues(t, 50, s.GetProgress().PercentComplete)
}
pkg/sources/s3/s3.go
+4 -4
View File
@@ -60,8 +60,8 @@ type Source struct {
endpoint *url.URL
// objectFilter is never nil after Init.
objectFilter *objectFilter
// progress is never nil after Init.
progress *scanProgress
// objectProgress is never nil after Init.
objectProgress *scanProgress
}
// Ensure the Source satisfies the interfaces at compile time
@@ -93,7 +93,7 @@ func (s *Source) Init(
s.verify = verify
s.concurrency = concurrency
s.errorCount = &sync.Map{}
s.progress = &scanProgress{}
s.objectProgress = &scanProgress{}
s.jobPool = &errgroup.Group{}
s.jobPool.SetLimit(concurrency)
@@ -463,7 +463,7 @@ func (s *Source) scanBucket(
if checkpointer.isUnitScan {
countCtx, cancel := context.WithCancel(ctx)
defer cancel()
s.progress.listingsInFlight.Add(1)
s.objectProgress.listingsInFlight.Add(1)
go s.countBucket(countCtx, regionalClient, bucket, startAfter)
}
pkg/sources/s3/s3_integration_test.go
+4 -4
View File
@@ -665,10 +665,10 @@ func TestSource_ChunkUnit_ReportsPercent(t *testing.T) {
require.Len(t, reporter.Units, 1)
require.NoError(t, s.ChunkUnit(ctx, reporter.Units[0], &reporter))
require.Eventually(t, func() bool { return s.progress.listingsInFlight.Load() == 0 }, 5*time.Second, 10*time.Millisecond)
assert.False(t, s.progress.countFailed.Load())
assert.Positive(t, s.progress.objectsTotal.Load())
assert.Equal(t, s.progress.objectsTotal.Load(), s.progress.objectsDone.Load())
require.Eventually(t, func() bool { return s.objectProgress.listingsInFlight.Load() == 0 }, 5*time.Second, 10*time.Millisecond)
assert.False(t, s.objectProgress.countFailed.Load())
assert.Positive(t, s.objectProgress.objectsTotal.Load())
assert.Equal(t, s.objectProgress.objectsTotal.Load(), s.objectProgress.objectsDone.Load())
assert.EqualValues(t, 99, s.GetProgress().PercentComplete)
unitID, _ := reporter.Units[0].SourceUnitID()
pkg/sources/s3/s3_test.go
+6 -6
View File
@@ -552,7 +552,7 @@ func TestSource_PageChunker_UnitScanCountsSkippedObjectsAsDone(t *testing.T) {
size := int64(objectSize)
page.Contents[i] = s3types.Object{Key: &key, Size: &size}
}
s.progress.addTotal(objectCount, objectCount*objectSize)
s.objectProgress.addTotal(objectCount, objectCount*objectSize)
// Every object is filtered, so pageChunker never reaches GetObject and needs no client.
var scanned, filtered uint64
@@ -564,8 +564,8 @@ func TestSource_PageChunker_UnitScanCountsSkippedObjectsAsDone(t *testing.T) {
NewCheckpointer(ctx, &s.Progress, true),
)
assert.EqualValues(t, objectCount, s.progress.objectsDone.Load())
assert.EqualValues(t, objectCount*objectSize, s.progress.bytesDone.Load())
assert.EqualValues(t, objectCount, s.objectProgress.objectsDone.Load())
assert.EqualValues(t, objectCount*objectSize, s.objectProgress.bytesDone.Load())
assert.EqualValues(t, 99, s.GetProgress().PercentComplete)
assert.EqualValues(t, objectCount, s.GetProgress().SectionsCompleted)
// Filtered objects advance the checkpoint to the last key of the page, and the percent update
@@ -599,8 +599,8 @@ func TestSource_PageChunker_UnitScanLeavesUncountedObjectsOutOfProgress(t *testi
NewCheckpointer(ctx, &s.Progress, true),
)
assert.Zero(t, s.progress.objectsDone.Load())
assert.Zero(t, s.progress.bytesDone.Load())
assert.Zero(t, s.objectProgress.objectsDone.Load())
assert.Zero(t, s.objectProgress.bytesDone.Load())
}
func TestSource_PageChunker_LegacyScanLeavesProgressUntouched(t *testing.T) {
@@ -628,7 +628,7 @@ func TestSource_PageChunker_LegacyScanLeavesProgressUntouched(t *testing.T) {
NewCheckpointer(ctx, &s.Progress, false),
)
assert.Zero(t, s.progress.objectsDone.Load())
assert.Zero(t, s.objectProgress.objectsDone.Load())
assert.EqualValues(t, 1, s.GetProgress().SectionsRemaining)
assert.Equal(t, "Bucket: test-bucket", s.GetProgress().Message)
}