diff --git a/pkg/sources/source_manager.go b/pkg/sources/source_manager.go index c21a925cac19..d3596de798d9 100644 --- a/pkg/sources/source_manager.go +++ b/pkg/sources/source_manager.go @@ -16,6 +16,8 @@ import ( "github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb" ) +// SourceManager provides an interface for starting and managing running +// sources. type SourceManager struct { api apiClient hooks []JobProgressHook @@ -49,6 +51,8 @@ func WithAPI(api apiClient) func(*SourceManager) { return func(mgr *SourceManager) { mgr.api = api } } +// WithReportHook adds a hook to the SourceManager's reporting feature for +// customizing data aggregation. func WithReportHook(hook JobProgressHook) func(*SourceManager) { return func(mgr *SourceManager) { mgr.hooks = append(mgr.hooks, hook) @@ -82,6 +86,10 @@ func WithSourceUnits() func(*SourceManager) { } } +// WithSourceUnitsFunc dynamically configures whether to use source unit +// enumeration and chunking if the source supports it. If the function returns +// true and the source supports it, then units will be used. Otherwise, the +// legacy scanning method will be used. func WithSourceUnitsFunc(f func() bool) func(*SourceManager) { return func(mgr *SourceManager) { mgr.useSourceUnitsFunc = f } } @@ -394,33 +402,45 @@ func (api *headlessAPI) GetIDs(context.Context, string, sourcespb.SourceType) (S } // mgrUnitReporter implements the UnitReporter interface. +var _ UnitReporter = (*mgrUnitReporter)(nil) + type mgrUnitReporter struct { unitCh chan SourceUnit report *JobProgress } +// UnitOk implements the UnitReporter interface by recording the unit in the +// report and sending it on the SourceUnit channel. func (s *mgrUnitReporter) UnitOk(ctx context.Context, unit SourceUnit) error { s.report.ReportUnit(unit) return common.CancellableWrite(ctx, s.unitCh, unit) } +// UnitErr implements the UnitReporter interface by recording the error in the +// report. func (s *mgrUnitReporter) UnitErr(ctx context.Context, err error) error { s.report.ReportError(err) return nil } // mgrChunkReporter implements the ChunkReporter interface. +var _ ChunkReporter = (*mgrChunkReporter)(nil) + type mgrChunkReporter struct { unit SourceUnit chunkCh chan *Chunk report *JobProgress } +// ChunkOk implements the ChunkReporter interface by recording the chunk and +// its associated unit in the report and sending it on the Chunk channel. func (s *mgrChunkReporter) ChunkOk(ctx context.Context, chunk Chunk) error { s.report.ReportChunk(s.unit, &chunk) return common.CancellableWrite(ctx, s.chunkCh, &chunk) } +// ChunkErr implements the ChunkReporter interface by recording the error and +// its associated unit in the report. func (s *mgrChunkReporter) ChunkErr(ctx context.Context, err error) error { s.report.ReportError(ChunkError{s.unit, err}) return nil