[scan-9] Update enumeration logic (#3626)

* renaming to enumeration * update enumeration * comments * remove commented out func --------- Co-authored-by: Miccah Castorina <m.castorina93@gmail.com>
trufflesecurity · Nov 25, 2024 · 1276d26 · 1276d26
1 parent f119adc
commit 1276d26
Show file tree

Hide file tree

Showing 18 changed files with 171 additions and 29 deletions.
diff --git a/pkg/engine/circleci.go b/pkg/engine/circleci.go
@@ -34,5 +34,5 @@ func (e *Engine) ScanCircleCI(ctx context.Context, token string) (sources.JobPro
 	if err := circleSource.Init(ctx, "trufflehog - Circle CI", jobID, sourceID, true, &conn, runtime.NumCPU()); err != nil {
 		return sources.JobProgressRef{}, err
 	}
-	return e.sourceManager.Run(ctx, sourceName, circleSource)
+	return e.sourceManager.EnumerateAndScan(ctx, sourceName, circleSource)
 }
diff --git a/pkg/engine/docker.go b/pkg/engine/docker.go
@@ -39,5 +39,5 @@ func (e *Engine) ScanDocker(ctx context.Context, c sources.DockerConfig) (source
 	if err := dockerSource.Init(ctx, sourceName, jobID, sourceID, true, &conn, runtime.NumCPU()); err != nil {
 		return sources.JobProgressRef{}, err
 	}
-	return e.sourceManager.Run(ctx, sourceName, dockerSource)
+	return e.sourceManager.EnumerateAndScan(ctx, sourceName, dockerSource)
 }
diff --git a/pkg/engine/elasticsearch.go b/pkg/engine/elasticsearch.go
@@ -41,5 +41,5 @@ func (e *Engine) ScanElasticsearch(ctx context.Context, c sources.ElasticsearchC
 	if err := elasticsearchSource.Init(ctx, sourceName, jobID, sourceID, true, &conn, runtime.NumCPU()); err != nil {
 		return sources.JobProgressRef{}, err
 	}
-	return e.sourceManager.Run(ctx, sourceName, elasticsearchSource)
+	return e.sourceManager.EnumerateAndScan(ctx, sourceName, elasticsearchSource)
 }
diff --git a/pkg/engine/filesystem.go b/pkg/engine/filesystem.go
@@ -33,5 +33,5 @@ func (e *Engine) ScanFileSystem(ctx context.Context, c sources.FilesystemConfig)
 	if err := fileSystemSource.Init(ctx, sourceName, jobID, sourceID, true, &conn, runtime.NumCPU()); err != nil {
 		return sources.JobProgressRef{}, err
 	}
-	return e.sourceManager.Run(ctx, sourceName, fileSystemSource)
+	return e.sourceManager.EnumerateAndScan(ctx, sourceName, fileSystemSource)
 }
diff --git a/pkg/engine/gcs.go b/pkg/engine/gcs.go
@@ -51,7 +51,7 @@ func (e *Engine) ScanGCS(ctx context.Context, c sources.GCSConfig) (sources.JobP
 	if err := gcsSource.Init(ctx, sourceName, jobID, sourceID, true, &conn, int(c.Concurrency)); err != nil {
 		return sources.JobProgressRef{}, err
 	}
-	return e.sourceManager.Run(ctx, sourceName, gcsSource)
+	return e.sourceManager.EnumerateAndScan(ctx, sourceName, gcsSource)
 }
 
 func isAuthValid(ctx context.Context, c sources.GCSConfig, connection *sourcespb.GCS) bool {

diff --git a/pkg/engine/git.go b/pkg/engine/git.go
@@ -39,5 +39,5 @@ func (e *Engine) ScanGit(ctx context.Context, c sources.GitConfig) (sources.JobP
 		return sources.JobProgressRef{}, err
 	}
 
-	return e.sourceManager.Run(ctx, sourceName, gitSource)
+	return e.sourceManager.EnumerateAndScan(ctx, sourceName, gitSource)
 }
diff --git a/pkg/engine/github.go b/pkg/engine/github.go
@@ -59,5 +59,5 @@ func (e *Engine) ScanGitHub(ctx context.Context, c sources.GithubConfig) (source
 		return sources.JobProgressRef{}, err
 	}
 	githubSource.WithScanOptions(scanOptions)
-	return e.sourceManager.Run(ctx, sourceName, githubSource)
+	return e.sourceManager.EnumerateAndScan(ctx, sourceName, githubSource)
 }
diff --git a/pkg/engine/github_experimental.go b/pkg/engine/github_experimental.go
@@ -60,5 +60,5 @@ func (e *Engine) ScanGitHubExperimental(ctx context.Context, c sources.GitHubExp
 		return sources.JobProgressRef{}, err
 	}
 	githubExperimentalSource.WithScanOptions(scanOptions)
-	return e.sourceManager.Run(ctx, sourceName, githubExperimentalSource)
+	return e.sourceManager.EnumerateAndScan(ctx, sourceName, githubExperimentalSource)
 }
diff --git a/pkg/engine/gitlab.go b/pkg/engine/gitlab.go
@@ -66,5 +66,5 @@ func (e *Engine) ScanGitLab(ctx context.Context, c sources.GitlabConfig) (source
 		return sources.JobProgressRef{}, err
 	}
 	gitlabSource.WithScanOptions(scanOptions)
-	return e.sourceManager.Run(ctx, sourceName, gitlabSource)
+	return e.sourceManager.EnumerateAndScan(ctx, sourceName, gitlabSource)
 }
diff --git a/pkg/engine/huggingface.go b/pkg/engine/huggingface.go
@@ -76,5 +76,5 @@ func (e *Engine) ScanHuggingface(ctx context.Context, c HuggingfaceConfig) (sour
 	if err := huggingfaceSource.Init(ctx, sourceName, jobID, sourceID, true, &conn, c.Concurrency); err != nil {
 		return sources.JobProgressRef{}, err
 	}
-	return e.sourceManager.Run(ctx, sourceName, huggingfaceSource)
+	return e.sourceManager.EnumerateAndScan(ctx, sourceName, huggingfaceSource)
 }
diff --git a/pkg/engine/jenkins.go b/pkg/engine/jenkins.go
@@ -77,5 +77,5 @@ func (e *Engine) ScanJenkins(ctx context.Context, jenkinsConfig JenkinsConfig) (
 	if err := jenkinsSource.Init(ctx, "trufflehog - Jenkins", jobID, sourceID, true, &conn, runtime.NumCPU()); err != nil {
 		return sources.JobProgressRef{}, err
 	}
-	return e.sourceManager.Run(ctx, sourceName, jenkinsSource)
+	return e.sourceManager.EnumerateAndScan(ctx, sourceName, jenkinsSource)
 }
diff --git a/pkg/engine/postman.go b/pkg/engine/postman.go
@@ -61,5 +61,5 @@ func (e *Engine) ScanPostman(ctx context.Context, c sources.PostmanConfig) (sour
 	if err := postmanSource.Init(ctx, sourceName, jobID, sourceID, true, &conn, c.Concurrency); err != nil {
 		return sources.JobProgressRef{}, err
 	}
-	return e.sourceManager.Run(ctx, sourceName, postmanSource)
+	return e.sourceManager.EnumerateAndScan(ctx, sourceName, postmanSource)
 }
diff --git a/pkg/engine/s3.go b/pkg/engine/s3.go
@@ -68,5 +68,5 @@ func (e *Engine) ScanS3(ctx context.Context, c sources.S3Config) (sources.JobPro
 	if err := s3Source.Init(ctx, sourceName, jobID, sourceID, true, &conn, runtime.NumCPU()); err != nil {
 		return sources.JobProgressRef{}, err
 	}
-	return e.sourceManager.Run(ctx, sourceName, s3Source)
+	return e.sourceManager.EnumerateAndScan(ctx, sourceName, s3Source)
 }
diff --git a/pkg/engine/syslog.go b/pkg/engine/syslog.go
@@ -49,5 +49,5 @@ func (e *Engine) ScanSyslog(ctx context.Context, c sources.SyslogConfig) (source
 	}
 	syslogSource.InjectConnection(connection)
 
-	return e.sourceManager.Run(ctx, sourceName, syslogSource)
+	return e.sourceManager.EnumerateAndScan(ctx, sourceName, syslogSource)
 }
diff --git a/pkg/engine/travisci.go b/pkg/engine/travisci.go
@@ -34,5 +34,5 @@ func (e *Engine) ScanTravisCI(ctx context.Context, token string) (sources.JobPro
 	if err := travisSource.Init(ctx, sourceName, jobID, sourceID, true, &conn, runtime.NumCPU()); err != nil {
 		return sources.JobProgressRef{}, err
 	}
-	return e.sourceManager.Run(ctx, sourceName, travisSource)
+	return e.sourceManager.EnumerateAndScan(ctx, sourceName, travisSource)
 }
diff --git a/pkg/sources/source_manager.go b/pkg/sources/source_manager.go
@@ -123,10 +123,10 @@ func (s *SourceManager) GetIDs(ctx context.Context, sourceName string, kind sour
 	return s.api.GetIDs(ctx, sourceName, kind)
 }
 
-// Run blocks until a resource is available to run the source, then
+// EnumerateAndScan blocks until a resource is available to run the source, then
 // asynchronously runs it. Error information is stored and accessible via the
 // JobProgressRef as it becomes available.
-func (s *SourceManager) Run(ctx context.Context, sourceName string, source Source, targets ...ChunkingTarget) (JobProgressRef, error) {
+func (s *SourceManager) EnumerateAndScan(ctx context.Context, sourceName string, source Source, targets ...ChunkingTarget) (JobProgressRef, error) {
 	sourceID, jobID := source.SourceID(), source.JobID()
 	// Do preflight checks before waiting on the pool.
 	if err := s.preflightChecks(ctx); err != nil {
@@ -169,6 +169,54 @@ func (s *SourceManager) Run(ctx context.Context, sourceName string, source Sourc
 	return progress.Ref(), nil
 }
 
+func (s *SourceManager) Enumerate(ctx context.Context, sourceName string, source Source, reporter UnitReporter) (JobProgressRef, error) {
+	sourceID, jobID := source.SourceID(), source.JobID()
+	// Do preflight checks before waiting on the pool.
+	if err := s.preflightChecks(ctx); err != nil {
+		return JobProgressRef{
+			SourceName: sourceName,
+			SourceID:   sourceID,
+			JobID:      jobID,
+		}, err
+	}
+
+	// Create a JobProgress object for tracking progress.
+	sem := s.sem
+	ctx, cancel := context.WithCancelCause(ctx)
+	progress := NewJobProgress(jobID, sourceID, sourceName, WithHooks(s.hooks...), WithCancel(cancel))
+	if err := sem.Acquire(ctx, 1); err != nil {
+		// Context cancelled.
+		progress.ReportError(Fatal{err})
+		return progress.Ref(), Fatal{err}
+	}
+
+	// Wrap the passed in reporter so we update the progress information.
+	reporter = baseUnitReporter{
+		child:    reporter,
+		progress: progress,
+	}
+
+	s.wg.Add(1)
+	go func() {
+		// Call Finish after the semaphore has been released.
+		defer progress.Finish()
+		defer sem.Release(1)
+		defer s.wg.Done()
+		ctx := context.WithValues(ctx,
+			"source_manager_worker_id", common.RandomID(5),
+		)
+		defer common.Recover(ctx)
+		defer cancel(nil)
+		if err := s.enumerate(ctx, source, progress, reporter); err != nil {
+			select {
+			case s.firstErr <- err:
+			default:
+			}
+		}
+	}()
+	return progress.Ref(), nil
+}
+
 // Chunks returns the read only channel of all the chunks produced by all of
 // the sources managed by this manager.
 func (s *SourceManager) Chunks() <-chan *Chunk {
@@ -286,6 +334,75 @@ func (s *SourceManager) run(ctx context.Context, source Source, report *JobProgr
 	return s.runWithoutUnits(ctx, source, report, targets...)
 }
 
+// enumerate is a helper method to enumerate a Source.
+func (s *SourceManager) enumerate(ctx context.Context, source Source, report *JobProgress, reporter UnitReporter) error {
+	report.Start(time.Now())
+	defer func() { report.End(time.Now()) }()
+
+	defer func() {
+		if err := context.Cause(ctx); err != nil {
+			report.ReportError(Fatal{err})
+		}
+	}()
+
+	report.TrackProgress(source.GetProgress())
+	if ctx.Value("job_id") == "" {
+		ctx = context.WithValue(ctx, "job_id", report.JobID)
+	}
+	if ctx.Value("source_id") == "" {
+		ctx = context.WithValue(ctx, "source_id", report.SourceID)
+	}
+	if ctx.Value("source_name") == "" {
+		ctx = context.WithValue(ctx, "source_name", report.SourceName)
+	}
+	if ctx.Value("source_type") == "" {
+		ctx = context.WithValue(ctx, "source_type", source.Type().String())
+	}
+
+	// Check for the preferred method of tracking source units.
+	canUseSourceUnits := s.useSourceUnitsFunc != nil
+	if enumChunker, ok := source.(SourceUnitEnumerator); ok && canUseSourceUnits && s.useSourceUnitsFunc() {
+		ctx.Logger().Info("running source",
+			"with_units", true)
+		return s.enumerateWithUnits(ctx, enumChunker, report, reporter)
+	}
+	return fmt.Errorf("Enumeration not supported or configured for source: %s", source.Type().String())
+}
+
+// enumerateWithUnits is a helper method to enumerate a Source that is also a
+// SourceUnitEnumerator. This allows better introspection of what is getting
+// enumerated and any errors encountered.
+func (s *SourceManager) enumerateWithUnits(ctx context.Context, source SourceUnitEnumerator, report *JobProgress, reporter UnitReporter) error {
+	// Create a function that will save the first error encountered (if
+	// any) and discard the rest.
+	fatalErr := make(chan error, 1)
+	catchFirstFatal := func(err error) {
+		select {
+		case fatalErr <- err:
+		default:
+		}
+	}
+
+	// Produce units.
+	func() {
+		// TODO: Catch panics and add to report.
+		report.StartEnumerating(time.Now())
+		defer func() { report.EndEnumerating(time.Now()) }()
+		ctx.Logger().V(2).Info("enumerating source with units")
+		if err := source.Enumerate(ctx, reporter); err != nil {
+			report.ReportError(Fatal{err})
+			catchFirstFatal(Fatal{err})
+		}
+	}()
+
+	select {
+	case err := <-fatalErr:
+		return err
+	default:
+		return nil
+	}
+}
+
 // runWithoutUnits is a helper method to run a Source. It has coarse-grained
 // job reporting.
 func (s *SourceManager) runWithoutUnits(ctx context.Context, source Source, report *JobProgress, targets ...ChunkingTarget) error {
@@ -302,6 +419,7 @@ func (s *SourceManager) runWithoutUnits(ctx context.Context, source Source, repo
 			s.outputChunks <- chunk
 		}
 	}()
+
 	// Don't return from this function until the goroutine has finished
 	// outputting chunks to the downstream channel. Closing the channel
 	// will stop the goroutine, so that needs to happen first in the defer

diff --git a/pkg/sources/source_manager_test.go b/pkg/sources/source_manager_test.go
@@ -114,7 +114,7 @@ func TestSourceManagerRun(t *testing.T) {
 	source, err := buildDummy(&counterChunker{count: 1})
 	assert.NoError(t, err)
 	for i := 0; i < 3; i++ {
-		ref, err := mgr.Run(context.Background(), "dummy", source)
+		ref, err := mgr.EnumerateAndScan(context.Background(), "dummy", source)
 		<-ref.Done()
 		assert.NoError(t, err)
 		assert.NoError(t, ref.Snapshot().FatalError())
@@ -132,7 +132,7 @@ func TestSourceManagerWait(t *testing.T) {
 	source, err := buildDummy(&counterChunker{count: 1})
 	assert.NoError(t, err)
 	// Asynchronously run the source.
-	_, err = mgr.Run(context.Background(), "dummy", source)
+	_, err = mgr.EnumerateAndScan(context.Background(), "dummy", source)
 	assert.NoError(t, err)
 	// Read the 1 chunk we're expecting so Waiting completes.
 	<-mgr.Chunks()
@@ -141,15 +141,15 @@ func TestSourceManagerWait(t *testing.T) {
 	// Run should return an error now.
 	_, err = buildDummy(&counterChunker{count: 1})
 	assert.NoError(t, err)
-	_, err = mgr.Run(context.Background(), "dummy", source)
+	_, err = mgr.EnumerateAndScan(context.Background(), "dummy", source)
 	assert.Error(t, err)
 }
 
 func TestSourceManagerError(t *testing.T) {
 	mgr := NewManager()
 	source, err := buildDummy(errorChunker{fmt.Errorf("oops")})
 	assert.NoError(t, err)
-	ref, err := mgr.Run(context.Background(), "dummy", source)
+	ref, err := mgr.EnumerateAndScan(context.Background(), "dummy", source)
 	assert.NoError(t, err)
 	<-ref.Done()
 	assert.Error(t, ref.Snapshot().FatalError())
@@ -165,7 +165,7 @@ func TestSourceManagerReport(t *testing.T) {
 		mgr := NewManager(opts...)
 		source, err := buildDummy(&counterChunker{count: 4})
 		assert.NoError(t, err)
-		ref, err := mgr.Run(context.Background(), "dummy", source)
+		ref, err := mgr.EnumerateAndScan(context.Background(), "dummy", source)
 		assert.NoError(t, err)
 		<-ref.Done()
 		assert.Equal(t, 0, len(ref.Snapshot().Errors))
@@ -230,7 +230,7 @@ func TestSourceManagerNonFatalError(t *testing.T) {
 	mgr := NewManager(WithBufferedOutput(8), WithSourceUnits())
 	source, err := buildDummy(&unitChunker{input})
 	assert.NoError(t, err)
-	ref, err := mgr.Run(context.Background(), "dummy", source)
+	ref, err := mgr.EnumerateAndScan(context.Background(), "dummy", source)
 	assert.NoError(t, err)
 	<-ref.Done()
 	report := ref.Snapshot()
@@ -247,7 +247,7 @@ func TestSourceManagerContextCancelled(t *testing.T) {
 	assert.NoError(t, err)
 
 	ctx, cancel := context.WithCancel(context.Background())
-	ref, err := mgr.Run(ctx, "dummy", source)
+	ref, err := mgr.EnumerateAndScan(ctx, "dummy", source)
 	assert.NoError(t, err)
 
 	cancel()
@@ -291,7 +291,7 @@ func TestSourceManagerCancelRun(t *testing.T) {
 	}})
 	assert.NoError(t, err)
 
-	ref, err := mgr.Run(context.Background(), "dummy", source)
+	ref, err := mgr.EnumerateAndScan(context.Background(), "dummy", source)
 	assert.NoError(t, err)
 
 	cancelErr := fmt.Errorf("abort! abort!")
@@ -313,7 +313,7 @@ func TestSourceManagerAvailableCapacity(t *testing.T) {
 	assert.NoError(t, err)
 
 	assert.Equal(t, 1337, mgr.AvailableCapacity())
-	ref, err := mgr.Run(context.Background(), "dummy", source)
+	ref, err := mgr.EnumerateAndScan(context.Background(), "dummy", source)
 	assert.NoError(t, err)
 
 	<-start // Wait for start signal.
@@ -338,7 +338,7 @@ func TestSourceManagerUnitHook(t *testing.T) {
 	)
 	source, err := buildDummy(&unitChunker{input})
 	assert.NoError(t, err)
-	ref, err := mgr.Run(context.Background(), "dummy", source)
+	ref, err := mgr.EnumerateAndScan(context.Background(), "dummy", source)
 	assert.NoError(t, err)
 	<-ref.Done()
 	assert.NoError(t, mgr.Wait())
@@ -399,7 +399,7 @@ func TestSourceManagerUnitHookBackPressure(t *testing.T) {
 	)
 	source, err := buildDummy(&unitChunker{input})
 	assert.NoError(t, err)
-	ref, err := mgr.Run(context.Background(), "dummy", source)
+	ref, err := mgr.EnumerateAndScan(context.Background(), "dummy", source)
 	assert.NoError(t, err)
 
 	var metrics []UnitMetrics
@@ -428,7 +428,7 @@ func TestSourceManagerUnitHookNoUnits(t *testing.T) {
 	source, err := buildDummy(&counterChunker{count: 5})
 	assert.NoError(t, err)
 
-	ref, err := mgr.Run(context.Background(), "dummy", source)
+	ref, err := mgr.EnumerateAndScan(context.Background(), "dummy", source)
 	assert.NoError(t, err)
 	<-ref.Done()
 	assert.NoError(t, mgr.Wait())

diff --git a/pkg/sources/sources.go b/pkg/sources/sources.go
@@ -103,6 +103,30 @@ type SourceUnitEnumerator interface {
 	Enumerate(ctx context.Context, reporter UnitReporter) error
 }
 
+// BaseUnitReporter is a helper struct that implements the UnitReporter interface
+// and includes a JobProgress reference.
+type baseUnitReporter struct {
+	child       UnitReporter
+	progress *JobProgress
+}
+
+func (b baseUnitReporter) UnitOk(ctx context.Context, unit SourceUnit) error {
+	b.progress.ReportUnit(unit)
+	if b.child != nil {
+		return b.child.UnitOk(ctx, unit)
+	}
+	return nil
+}
+
+func (b baseUnitReporter) UnitErr(ctx context.Context, err error) error {
+	b.progress.ReportError(err)
+	if b.child != nil {
+		return b.child.UnitErr(ctx, err)
+	}
+	return nil
+}
+
+
 // UnitReporter defines the interface a source will use to report whether a
 // unit was found during enumeration. Either method may be called any number of
 // times. Implementors of this interface should allow for concurrent calls.