Skip to content

Commit

Permalink
[chore] Add some doc comments to source manager (#2434)
Browse files Browse the repository at this point in the history
  • Loading branch information
mcastorina authored Feb 13, 2024
1 parent e8006f1 commit 216a29d
Showing 1 changed file with 20 additions and 0 deletions.
20 changes: 20 additions & 0 deletions pkg/sources/source_manager.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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 }
}
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 216a29d

Please sign in to comment.