From 1276d262f2434dbd0a3cacdb91a49f443549921c Mon Sep 17 00:00:00 2001 From: 0x1 <13666360+0x1@users.noreply.github.com> Date: Mon, 25 Nov 2024 14:13:03 -0500 Subject: [PATCH] [scan-9] Update enumeration logic (#3626) * renaming to enumeration * update enumeration * comments * remove commented out func --------- Co-authored-by: Miccah Castorina --- pkg/engine/circleci.go | 2 +- pkg/engine/docker.go | 2 +- pkg/engine/elasticsearch.go | 2 +- pkg/engine/filesystem.go | 2 +- pkg/engine/gcs.go | 2 +- pkg/engine/git.go | 2 +- pkg/engine/github.go | 2 +- pkg/engine/github_experimental.go | 2 +- pkg/engine/gitlab.go | 2 +- pkg/engine/huggingface.go | 2 +- pkg/engine/jenkins.go | 2 +- pkg/engine/postman.go | 2 +- pkg/engine/s3.go | 2 +- pkg/engine/syslog.go | 2 +- pkg/engine/travisci.go | 2 +- pkg/sources/source_manager.go | 122 ++++++++++++++++++++++++++++- pkg/sources/source_manager_test.go | 24 +++--- pkg/sources/sources.go | 24 ++++++ 18 files changed, 171 insertions(+), 29 deletions(-) diff --git a/pkg/engine/circleci.go b/pkg/engine/circleci.go index 97627dd77..c2ca134e5 100644 --- 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 index 07f5376eb..1e55817a9 100644 --- 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 index bc6c49488..10bdf7193 100644 --- 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 index fe61d4948..90cc1521f 100644 --- 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 index 2284c5c69..7ad54f273 100644 --- 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 index de487591f..90adb0b0b 100644 --- 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 index b3908afc5..251843dfc 100644 --- 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 index e88fa3f1f..0354a3637 100644 --- 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 index 9baa07b2f..ad5f91cda 100644 --- 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 index f11fc851b..9a20fe27f 100644 --- 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 index 363e657a3..b488951a2 100644 --- 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 index ad71e07ac..9f09ac6df 100644 --- 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 index be74e0ec3..1f6fe5b6d 100644 --- 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 index 6c8518175..a857fd90c 100644 --- 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 index 987ff1342..790ced12d 100644 --- 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 index 29ac4ce56..3a3ce736d 100644 --- 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 index d87d77c27..4706ce141 100644 --- 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,7 +141,7 @@ 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) } @@ -149,7 +149,7 @@ 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 index 0a140e180..b1190209c 100644 --- 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.