Mirrors/trufflehog
2
0
Fork 0
mirror of https://github.com/trufflesecurity/trufflehog.git synced 2024-11-10 07:04:24 +00:00
Code Issues Projects Releases Packages Wiki Activity

Check for SourceUnit support dynamically in the SourceManager (#2205)

Browse source 
* Check for SourceUnit support dynamically in the SourceManager

* Only call the function if we can use source units
This commit is contained in:
Miccah 2023-12-14 11:48:15 -08:00 committed by GitHub
parent 06b137fd93
commit f6bbc59bf6
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 11 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
pkg/sources/source_manager.go
View file

@ -24,7 +24,8 @@ type SourceManager struct {
// Max number of units to scan concurrently per source.
concurrentUnits int
// Run the sources using source unit enumeration / chunking if available.
useSourceUnits bool
// Checked at runtime to allow feature flagging.
useSourceUnitsFunc func() bool
// Downstream chunks channel to be scanned.
outputChunks chan *Chunk
// Set when Wait() returns.
@ -67,7 +68,13 @@ func WithBufferedOutput(size int) func(*SourceManager) {
// WithSourceUnits enables using source unit enumeration and chunking if the
// source supports it.
func WithSourceUnits() func(*SourceManager) {
return func(mgr *SourceManager) { mgr.useSourceUnits = true }
return func(mgr *SourceManager) {
mgr.useSourceUnitsFunc = func() bool { return true }
}
}
func WithSourceUnitsFunc(f func() bool) func(*SourceManager) {
return func(mgr *SourceManager) { mgr.useSourceUnitsFunc = f }
}
// WithConcurrentUnits limits the number of units to be scanned concurrently.
@ -229,7 +236,8 @@ func (s *SourceManager) run(ctx context.Context, source Source, report *JobProgr
"source_type", source.Type().String(),
)
// Check for the preferred method of tracking source units.
if enumChunker, ok := source.(SourceUnitEnumChunker); ok && s.useSourceUnits && len(targets) == 0 {
canUseSourceUnits := len(targets) == 0 && s.useSourceUnitsFunc != nil
if enumChunker, ok := source.(SourceUnitEnumChunker); ok && canUseSourceUnits && s.useSourceUnitsFunc() {
return s.runWithUnits(ctx, enumChunker, report)
}
return s.runWithoutUnits(ctx, source, report, targets...)