Use SourceManager in Engine

And subsequently all of the sources in the engine package.

pkg/detectors/okta/okta.go

@@ -19,8 +19,8 @@ type Scanner struct{}
 var _ detectors.Detector = (*Scanner)(nil)
 
 var (
-	domainPat = regexp.MustCompile(`\b[a-z0-9-]{1,40}\.okta(?:preview|-emea){0,1}\.com\b`)
-	tokenPat  = regexp.MustCompile(`\b00[a-zA-Z0-9_-]{40}\b`)
+	domainPat = regexp.MustCompile(`[a-z0-9-]{1,40}\.okta(?:preview|-emea){0,1}\.com`)
+	tokenPat  = regexp.MustCompile(`00[a-zA-Z0-9_-]{40}`)
 	// TODO: Oauth client secrets
 )
 

pkg/detectors/postgres/postgres.go

@@ -1,105 +0,0 @@
-package postgres
-
-import (
-	"context"
-	"fmt"
-	"net/http"
-	"regexp"
-	"strings"
-
-	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
-	"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
-	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
-)
-
-type Scanner struct {
-	client *http.Client
-}
-
-// Ensure the Scanner satisfies the interface at compile time.
-var _ detectors.Detector = (*Scanner)(nil)
-
-var (
-	defaultClient = common.SaneHttpClient()
-	// Make sure that your group is surrounded in boundary characters such as below to reduce false positives.
-	keyPat = regexp.MustCompile(detectors.PrefixRegex([]string{"postgres"}) + `\b([0-9a-zA-Z]{23}_[0-9a-zA-Z]{8})\b`)
-)
-
-// find various connection strings
-
-// postgresql://user:secret@localhost:5432/mydb
-// "host=localhost dbname=publishing user=www password=foo"
-
-// detectors.PrefixRegex([]string{"nytimes"})
-
-//var (
-//	connectionStringPats = map[string]*regexp.Regexp{
-//		"KeyValue":               regexp.MustCompile(`xoxb\-[0-9]{10,13}\-[0-9]{10,13}[a-zA-Z0-9\-]*`),
-//		"":              regexp.MustCompile(`xoxp\-[0-9]{10,13}\-[0-9]{10,13}[a-zA-Z0-9\-]*`),
-//		"Slack Workspace Access Token":  regexp.MustCompile(`xoxa\-[0-9]{10,13}\-[0-9]{10,13}[a-zA-Z0-9\-]*`),
-//		"Slack Workspace Refresh Token": regexp.MustCompile(`xoxr\-[0-9]{10,13}\-[0-9]{10,13}[a-zA-Z0-9\-]*`),
-//	}
-//	verifyURL = "https://slack.com/api/auth.test"
-//)
-
-// Keywords are used for efficiently pre-filtering chunks.
-// Use identifiers in the secret preferably, or the provider name.
-func (s Scanner) Keywords() []string {
-	return []string{"postgres://", ""}
-}
-
-// FromData will find and optionally verify Postgres secrets in a given set of bytes.
-func (s Scanner) FromData(ctx context.Context, verify bool, data []byte) (results []detectors.Result, err error) {
-	dataStr := string(data)
-
-	matches := keyPat.FindAllStringSubmatch(dataStr, -1)
-
-	for _, match := range matches {
-		if len(match) != 2 {
-			continue
-		}
-		resMatch := strings.TrimSpace(match[1])
-
-		s1 := detectors.Result{
-			DetectorType: detectorspb.DetectorType_Postgres,
-			Raw:          []byte(resMatch),
-		}
-
-		if verify {
-			client := s.client
-			if client == nil {
-				client = defaultClient
-			}
-			req, err := http.NewRequestWithContext(ctx, "GET", "https://eth-mainnet.g.postgres.com/v2/"+resMatch+"/getNFTs/?owner=vitalik.eth", nil)
-			if err != nil {
-				continue
-			}
-			res, err := client.Do(req)
-			if err == nil {
-				defer res.Body.Close()
-				if res.StatusCode >= 200 && res.StatusCode < 300 {
-					s1.Verified = true
-				} else {
-					// This function will check false positives for common test words, but also it will make sure the key appears 'random' enough to be a real key.
-					if detectors.IsKnownFalsePositive(resMatch, detectors.DefaultFalsePositives, true) {
-						continue
-					}
-
-					if res.StatusCode != 401 {
-						s1.VerificationError = fmt.Errorf("request to %v returned unexpected status %d", res.Request.URL, res.StatusCode)
-					}
-				}
-			} else {
-				s1.VerificationError = err
-			}
-		}
-
-		results = append(results, s1)
-	}
-
-	return results, nil
-}
-
-func (s Scanner) Type() detectorspb.DetectorType {
-	return detectorspb.DetectorType_Postgres
-}

pkg/detectors/postgres/postgres_test.go

@@ -1,160 +0,0 @@
-//go:build detectors
-// +build detectors
-
-package postgres
-
-import (
-	"context"
-	"fmt"
-	"testing"
-	"time"
-
-	"github.com/kylelemons/godebug/pretty"
-	"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
-
-	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
-	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
-)
-
-func TestPostgres_FromChunk(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
-	defer cancel()
-	testSecrets, err := common.GetSecret(ctx, "trufflehog-testing", "detectors4")
-	if err != nil {
-		t.Fatalf("could not get test secrets from GCP: %s", err)
-	}
-	secret := testSecrets.MustGetField("POSTGRES")
-	inactiveSecret := testSecrets.MustGetField("POSTGRES_INACTIVE")
-
-	type args struct {
-		ctx    context.Context
-		data   []byte
-		verify bool
-	}
-	tests := []struct {
-		name                string
-		s                   Scanner
-		args                args
-		want                []detectors.Result
-		wantErr             bool
-		wantVerificationErr bool
-	}{
-		{
-			name: "found, verified",
-			s:    Scanner{},
-			args: args{
-				ctx:    context.Background(),
-				data:   []byte(fmt.Sprintf("You can find a postgres secret %s within", secret)),
-				verify: true,
-			},
-			want: []detectors.Result{
-				{
-					DetectorType: detectorspb.DetectorType_Postgres,
-					Verified:     true,
-				},
-			},
-			wantErr:             false,
-			wantVerificationErr: false,
-		},
-		{
-			name: "found, unverified",
-			s:    Scanner{},
-			args: args{
-				ctx:    context.Background(),
-				data:   []byte(fmt.Sprintf("You can find a postgres secret %s within but not valid", inactiveSecret)), // the secret would satisfy the regex but not pass validation
-				verify: true,
-			},
-			want: []detectors.Result{
-				{
-					DetectorType: detectorspb.DetectorType_Postgres,
-					Verified:     false,
-				},
-			},
-			wantErr:             false,
-			wantVerificationErr: false,
-		},
-		{
-			name: "not found",
-			s:    Scanner{},
-			args: args{
-				ctx:    context.Background(),
-				data:   []byte("You cannot find the secret within"),
-				verify: true,
-			},
-			want:                nil,
-			wantErr:             false,
-			wantVerificationErr: false,
-		},
-		{
-			name: "found, would be verified if not for timeout",
-			s:    Scanner{client: common.SaneHttpClientTimeOut(1 * time.Microsecond)},
-			args: args{
-				ctx:    context.Background(),
-				data:   []byte(fmt.Sprintf("You can find a postgres secret %s within", secret)),
-				verify: true,
-			},
-			want: []detectors.Result{
-				{
-					DetectorType: detectorspb.DetectorType_Postgres,
-					Verified:     false,
-				},
-			},
-			wantErr:             false,
-			wantVerificationErr: true,
-		},
-		{
-			name: "found, verified but unexpected api surface",
-			s:    Scanner{client: common.ConstantStatusHttpClient(404)},
-			args: args{
-				ctx:    context.Background(),
-				data:   []byte(fmt.Sprintf("You can find a postgres secret %s within", secret)),
-				verify: true,
-			},
-			want: []detectors.Result{
-				{
-					DetectorType: detectorspb.DetectorType_Postgres,
-					Verified:     false,
-				},
-			},
-			wantErr:             false,
-			wantVerificationErr: true,
-		},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			got, err := tt.s.FromData(tt.args.ctx, tt.args.verify, tt.args.data)
-			if (err != nil) != tt.wantErr {
-				t.Errorf("Postgres.FromData() error = %v, wantErr %v", err, tt.wantErr)
-				return
-			}
-			for i := range got {
-				if len(got[i].Raw) == 0 {
-					t.Fatalf("no raw secret present: \n %+v", got[i])
-				}
-				got[i].Raw = nil
-				if (got[i].VerificationError != nil) != tt.wantVerificationErr {
-					t.Fatalf("verification error = %v, wantVerificationError %v", got[i].VerificationError, tt.wantVerificationErr)
-				}
-				got[i].VerificationError = nil
-			}
-			if diff := pretty.Compare(got, tt.want); diff != "" {
-				t.Errorf("Postgres.FromData() %s diff: (-got +want)\n%s", tt.name, diff)
-			}
-		})
-	}
-}
-
-func BenchmarkFromData(benchmark *testing.B) {
-	ctx := context.Background()
-	s := Scanner{}
-	for name, data := range detectors.MustGetBenchmarkData() {
-		benchmark.Run(name, func(b *testing.B) {
-			for n := 0; n < b.N; n++ {
-				_, err := s.FromData(ctx, false, data)
-				if err != nil {
-					b.Fatal(err)
-				}
-			}
-		})
-	}
-}

pkg/engine/circleci.go

@@ -1,16 +1,14 @@
 package engine
 
 import (
-	"fmt"
 	"runtime"
 
-	"github.com/go-errors/errors"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/anypb"
 
-	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources/circleci"
 )
 
@@ -29,23 +27,16 @@ func (e *Engine) ScanCircleCI(ctx context.Context, token string) error {
 		return err
 	}
 
-	circleSource := circleci.Source{}
-	ctx = context.WithValues(ctx,
-		"source_type", circleSource.Type().String(),
-		"source_name", "Circle CI",
-	)
-	err = circleSource.Init(ctx, "trufflehog - Circle CI", 0, int64(sourcespb.SourceType_SOURCE_TYPE_CIRCLECI), true, &conn, runtime.NumCPU())
+	handle, err := e.sourceManager.Enroll(ctx, "trufflehog - Circle CI", (&circleci.Source{}).Type(),
+		func(ctx context.Context, jobID, sourceID int64) (sources.Source, error) {
+			circleSource := &circleci.Source{}
+			if err := circleSource.Init(ctx, "trufflehog - Circle CI", jobID, sourceID, true, &conn, runtime.NumCPU()); err != nil {
+				return nil, err
+			}
+			return circleSource, nil
+		})
 	if err != nil {
-		return errors.WrapPrefix(err, "failed to init Circle CI source", 0)
+		return err
 	}
-
-	e.sourcesWg.Go(func() error {
-		defer common.RecoverWithExit(ctx)
-		err := circleSource.Chunks(ctx, e.ChunksChan())
-		if err != nil {
-			return fmt.Errorf("error scanning CircleCI: %w", err)
-		}
-		return nil
-	})
-	return nil
+	return e.sourceManager.ScheduleRun(ctx, handle)
 }

pkg/engine/docker.go

@@ -1,38 +1,27 @@
 package engine
 
 import (
-	"fmt"
 	"runtime"
 
-	"github.com/go-errors/errors"
 	"google.golang.org/protobuf/types/known/anypb"
 
-	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
-	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources/docker"
 )
 
 // ScanDocker scans a given docker connection.
 func (e *Engine) ScanDocker(ctx context.Context, conn *anypb.Any) error {
-	dockerSource := docker.Source{}
-
-	ctx = context.WithValues(ctx,
-		"source_type", dockerSource.Type().String(),
-		"source_name", "docker",
-	)
-	err := dockerSource.Init(ctx, "trufflehog - docker", 0, int64(sourcespb.SourceType_SOURCE_TYPE_DOCKER), true, conn, runtime.NumCPU())
+	handler, err := e.sourceManager.Enroll(ctx, "trufflehog - docker", (&docker.Source{}).Type(),
+		func(ctx context.Context, jobID, sourceID int64) (sources.Source, error) {
+			dockerSource := &docker.Source{}
+			if err := dockerSource.Init(ctx, "trufflehog - docker", jobID, sourceID, true, conn, runtime.NumCPU()); err != nil {
+				return nil, err
+			}
+			return dockerSource, nil
+		})
 	if err != nil {
-		return errors.WrapPrefix(err, "could not init docker source", 0)
+		return err
 	}
-
-	e.sourcesWg.Go(func() error {
-		defer common.RecoverWithExit(ctx)
-		err := dockerSource.Chunks(ctx, e.ChunksChan())
-		if err != nil {
-			return fmt.Errorf("error scanning docker image: %w", err)
-		}
-		return nil
-	})
-	return nil
+	return e.sourceManager.ScheduleRun(ctx, handler)
 }

pkg/engine/engine.go

@@ -11,7 +11,6 @@ import (
 	"time"
 
 	ahocorasick "github.com/petar-dambovaliev/aho-corasick"
-	"golang.org/x/sync/errgroup"
 	"google.golang.org/protobuf/proto"
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
@@ -27,14 +26,13 @@ import (
 
 type Engine struct {
 	concurrency     int
-	chunks          chan *sources.Chunk
+	sourceManager   *sources.SourceManager
 	results         chan detectors.ResultWithMetadata
 	decoders        []decoders.Decoder
 	detectors       map[bool][]detectors.Detector
 	chunksScanned   uint64
 	bytesScanned    uint64
 	detectorAvgTime sync.Map
-	sourcesWg       *errgroup.Group
 	workersWg       sync.WaitGroup
 	// filterUnverified is used to reduce the number of unverified results.
 	// If there are multiple unverified results for the same chunk for the same detector,
@@ -110,10 +108,8 @@ func filterDetectors(filterFunc func(detectors.Detector) bool, input []detectors
 
 func Start(ctx context.Context, options ...EngineOption) *Engine {
 	e := &Engine{
-		chunks:          make(chan *sources.Chunk),
 		results:         make(chan detectors.ResultWithMetadata),
 		detectorAvgTime: sync.Map{},
-		sourcesWg:       &errgroup.Group{},
 	}
 
 	for _, option := range options {
@@ -128,9 +124,6 @@ func Start(ctx context.Context, options ...EngineOption) *Engine {
 	}
 	ctx.Logger().V(2).Info("engine started", "workers", e.concurrency)
 
-	// Limit number of concurrent goroutines dedicated to chunking a source.
-	e.sourcesWg.SetLimit(e.concurrency)
-
 	if len(e.decoders) == 0 {
 		e.decoders = decoders.DefaultDecoders()
 	}
@@ -141,6 +134,9 @@ func Start(ctx context.Context, options ...EngineOption) *Engine {
 		e.detectors[false] = []detectors.Detector{}
 	}
 
+	// Setup source manager now that we know the concurrency.
+	e.sourceManager = sources.NewManager(sources.WithConcurrency(e.concurrency))
+
 	// build ahocorasick prefilter for efficient string matching
 	// on keywords
 	keywords := []string{}
@@ -184,7 +180,7 @@ func Start(ctx context.Context, options ...EngineOption) *Engine {
 	for i := 0; i < e.concurrency; i++ {
 		e.workersWg.Add(1)
 		go func() {
-			defer common.Recover(ctx)
+			defer common.RecoverWithExit(ctx)
 			defer e.workersWg.Done()
 			e.detectorWorker(ctx)
 		}()
@@ -198,13 +194,12 @@ func Start(ctx context.Context, options ...EngineOption) *Engine {
 // more sources may be scanned by the engine.
 func (e *Engine) Finish(ctx context.Context, logFunc func(error, string, ...any)) {
 	defer common.RecoverWithExit(ctx)
-	// wait for the sources to finish putting chunks onto the chunks channel
-	sourceErr := e.sourcesWg.Wait()
+	// Wait for the sources to finish putting chunks onto the chunks channel.
+	sourceErr := e.sourceManager.Wait()
 	if sourceErr != nil {
 		logFunc(sourceErr, "error occurred while collecting chunks")
 	}
 
-	close(e.chunks)
 	// wait for the workers to finish processing all of the chunks and putting
 	// results onto the results channel
 	e.workersWg.Wait()
@@ -217,8 +212,8 @@ func (e *Engine) Finish(ctx context.Context, logFunc func(error, string, ...any)
 	close(e.results)
 }
 
-func (e *Engine) ChunksChan() chan *sources.Chunk {
-	return e.chunks
+func (e *Engine) ChunksChan() <-chan *sources.Chunk {
+	return e.sourceManager.Chunks()
 }
 
 func (e *Engine) ResultsChan() chan detectors.ResultWithMetadata {
@@ -271,7 +266,7 @@ func (e *Engine) DetectorAvgTime() map[string][]time.Duration {
 }
 
 func (e *Engine) detectorWorker(ctx context.Context) {
-	for originalChunk := range e.chunks {
+	for originalChunk := range e.ChunksChan() {
 		for chunk := range sources.Chunker(originalChunk) {
 			var chunkResults []detectors.ResultWithMetadata
 			matchedKeywords := make(map[string]struct{})

pkg/engine/filesystem.go

@@ -1,14 +1,11 @@
 package engine
 
 import (
-	"fmt"
 	"runtime"
 
-	"github.com/go-errors/errors"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/anypb"
 
-	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
@@ -27,39 +24,17 @@ func (e *Engine) ScanFileSystem(ctx context.Context, c sources.FilesystemConfig)
 		return err
 	}
 
-	fileSystemSource := filesystem.Source{}
-	fileSystemSource.WithFilter(c.Filter)
-
-	// TODO: This is how it would be used; everything after these comments would be removed.
-	// handle, err := e.sourceManager.Enroll(ctx, "trufflehog - filesystem", fileSystemSource.Type(),
-	// 	func(ctx context.Context, jobID, sourceID int64) (sources.Source, error) {
-	// 		fileSystemSource := filesystem.Source{}
-	// 		fileSystemSource.WithFilter(c.Filter)
-	// 		if err := fileSystemSource.Init(ctx, "trufflehog - filesystem", jobID, sourceID, true, &conn, runtime.NumCPU()); err != nil {
-	// 			return nil, err
-	// 		}
-	// 		return &fileSystemSource, nil
-	// 	})
-	// if err != nil {
-	// 	return err
-	// }
-	// return e.sourceManager.ScheduleRun(ctx, handle)
-
-	ctx = context.WithValues(ctx,
-		"source_type", fileSystemSource.Type().String(),
-		"source_name", "filesystem",
-	)
-	err = fileSystemSource.Init(ctx, "trufflehog - filesystem", 0, int64(sourcespb.SourceType_SOURCE_TYPE_FILESYSTEM), true, &conn, runtime.NumCPU())
+	handle, err := e.sourceManager.Enroll(ctx, "trufflehog - filesystem", (&filesystem.Source{}).Type(),
+		func(ctx context.Context, jobID, sourceID int64) (sources.Source, error) {
+			fileSystemSource := &filesystem.Source{}
+			fileSystemSource.WithFilter(c.Filter)
+			if err := fileSystemSource.Init(ctx, "trufflehog - filesystem", jobID, sourceID, true, &conn, runtime.NumCPU()); err != nil {
+				return nil, err
+			}
+			return fileSystemSource, nil
+		})
 	if err != nil {
-		return errors.WrapPrefix(err, "could not init filesystem source", 0)
+		return err
 	}
-	e.sourcesWg.Go(func() error {
-		defer common.RecoverWithExit(ctx)
-		err := fileSystemSource.Chunks(ctx, e.ChunksChan())
-		if err != nil {
-			return fmt.Errorf("error scanning filesystem: %w", err)
-		}
-		return nil
-	})
-	return nil
+	return e.sourceManager.ScheduleRun(ctx, handle)
 }

pkg/engine/gcs.go

@@ -6,7 +6,6 @@ import (
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/anypb"
 
-	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
@@ -45,23 +44,18 @@ func (e *Engine) ScanGCS(ctx context.Context, c sources.GCSConfig) error {
 		return fmt.Errorf("failed to marshal GCS connection: %w", err)
 	}
 
-	source := gcs.Source{}
-	ctx = context.WithValues(ctx,
-		"source_type", source.Type().String(),
-		"source_name", "gcs",
-	)
-	if err = source.Init(ctx, "trufflehog - GCS", 0, 0, true, &conn, int(c.Concurrency)); err != nil {
-		return fmt.Errorf("failed to initialize GCS source: %w", err)
+	handle, err := e.sourceManager.Enroll(ctx, "trufflehog - GCS", (&gcs.Source{}).Type(),
+		func(ctx context.Context, jobID, sourceID int64) (sources.Source, error) {
+			source := &gcs.Source{}
+			if err := source.Init(ctx, "trufflehog - GCS", jobID, sourceID, true, &conn, int(c.Concurrency)); err != nil {
+				return nil, err
+			}
+			return source, nil
+		})
+	if err != nil {
+		return err
 	}
-
-	e.sourcesWg.Go(func() error {
-		defer common.RecoverWithExit(ctx)
-		if err := source.Chunks(ctx, e.ChunksChan()); err != nil {
-			return fmt.Errorf("could not scan GCS: %w", err)
-		}
-		return nil
-	})
-	return nil
+	return e.sourceManager.ScheduleRun(ctx, handle)
 }
 
 func isAuthValid(ctx context.Context, c sources.GCSConfig, connection *sourcespb.GCS) bool {

pkg/engine/git.go

@@ -1,14 +1,13 @@
 package engine
 
 import (
-	"fmt"
 	"runtime"
 
 	gogit "github.com/go-git/go-git/v5"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/anypb"
 
-	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
-	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/source_metadatapb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources/git"
@@ -22,16 +21,6 @@ func (e *Engine) ScanGit(ctx context.Context, c sources.GitConfig) error {
 		git.ScanOptionLogOptions(logOptions),
 	}
 
-	options := &gogit.PlainOpenOptions{
-		DetectDotGit:          true,
-		EnableDotGitCommonDir: true,
-	}
-
-	repo, err := gogit.PlainOpenWithOptions(c.RepoPath, options)
-	if err != nil {
-		return fmt.Errorf("could not open repo: %s: %w", c.RepoPath, err)
-	}
-
 	if c.MaxDepth != 0 {
 		opts = append(opts, git.ScanOptionMaxDepth(int64(c.MaxDepth)))
 	}
@@ -46,33 +35,34 @@ func (e *Engine) ScanGit(ctx context.Context, c sources.GitConfig) error {
 	}
 	scanOptions := git.NewScanOptions(opts...)
 
-	gitSource := git.NewGit(sourcespb.SourceType_SOURCE_TYPE_GIT, 0, 0, "trufflehog - git", true, runtime.NumCPU(),
-		func(file, email, commit, timestamp, repository string, line int64) *source_metadatapb.MetaData {
-			return &source_metadatapb.MetaData{
-				Data: &source_metadatapb.MetaData_Git{
-					Git: &source_metadatapb.Git{
-						Commit:     commit,
-						File:       file,
-						Email:      email,
-						Repository: repository,
-						Timestamp:  timestamp,
-						Line:       line,
-					},
-				},
-			}
-		})
+	connection := &sourcespb.Git{
+		// Using Directories here allows us to not pass any
+		// authentication. Also by this point, the c.RepoPath should
+		// still have been prepared and downloaded to a temporary
+		// directory if it was a URL.
+		Directories: []string{c.RepoPath},
+	}
+	var conn anypb.Any
+	err := anypb.MarshalFrom(&conn, connection, proto.MarshalOptions{})
+	if err != nil {
+		ctx.Logger().Error(err, "failed to marshal git connection")
+		return err
+	}
 
-	ctx = context.WithValues(ctx,
-		"source_type", sourcespb.SourceType_SOURCE_TYPE_GIT.String(),
-		"source_name", "git",
-	)
-	e.sourcesWg.Go(func() error {
-		defer common.RecoverWithExit(ctx)
-		err := gitSource.ScanRepo(ctx, repo, c.RepoPath, scanOptions, e.ChunksChan())
-		if err != nil {
-			return fmt.Errorf("could not scan repo: %w", err)
-		}
-		return nil
-	})
-	return nil
+	handle, err := e.sourceManager.Enroll(ctx, "trufflehog - git", (&git.Source{}).Type(),
+		func(ctx context.Context, jobID, sourceID int64) (sources.Source, error) {
+			gitSource := &git.Source{}
+			if err := gitSource.Init(ctx, "trufflehog - git", jobID, sourceID, true, &conn, runtime.NumCPU()); err != nil {
+				return nil, err
+			}
+			gitSource.WithScanOptions(scanOptions)
+			// Don't try to clean up the provided directory. That's
+			// handled by the caller of ScanGit.
+			gitSource.WithPreserveTempDirs(true)
+			return gitSource, nil
+		})
+	if err != nil {
+		return err
+	}
+	return e.sourceManager.ScheduleRun(ctx, handle)
 }

pkg/engine/github.go

@@ -1,13 +1,10 @@
 package engine
 
 import (
-	"fmt"
-
 	gogit "github.com/go-git/go-git/v5"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/anypb"
 
-	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
@@ -17,8 +14,6 @@ import (
 
 // ScanGitHub scans Github with the provided options.
 func (e *Engine) ScanGitHub(ctx context.Context, c sources.GithubConfig) error {
-	source := github.Source{}
-
 	connection := sourcespb.GitHub{
 		Endpoint:      c.Endpoint,
 		Organizations: c.Orgs,
@@ -42,31 +37,24 @@ func (e *Engine) ScanGitHub(ctx context.Context, c sources.GithubConfig) error {
 		return err
 	}
 
-	ctx = context.WithValues(ctx,
-		"source_type", source.Type().String(),
-		"source_name", "github",
-	)
-	err = source.Init(ctx, "trufflehog - github", 0, 0, false, &conn, c.Concurrency)
-	if err != nil {
-		ctx.Logger().Error(err, "failed to initialize github source")
-		return err
-	}
-
 	logOptions := &gogit.LogOptions{}
 	opts := []git.ScanOption{
 		git.ScanOptionFilter(c.Filter),
 		git.ScanOptionLogOptions(logOptions),
 	}
 	scanOptions := git.NewScanOptions(opts...)
-	source.WithScanOptions(scanOptions)
 
-	e.sourcesWg.Go(func() error {
-		defer common.RecoverWithExit(ctx)
-		err := source.Chunks(ctx, e.ChunksChan())
-		if err != nil {
-			return fmt.Errorf("could not scan github: %w", err)
-		}
-		return nil
-	})
-	return nil
+	handle, err := e.sourceManager.Enroll(ctx, "trufflehog - GitHub", (&github.Source{}).Type(),
+		func(ctx context.Context, jobID, sourceID int64) (sources.Source, error) {
+			source := &github.Source{}
+			source.WithScanOptions(scanOptions)
+			if err := source.Init(ctx, "trufflehog - GitHub", jobID, sourceID, false, &conn, c.Concurrency); err != nil {
+				return nil, err
+			}
+			return source, nil
+		})
+	if err != nil {
+		return err
+	}
+	return e.sourceManager.Run(ctx, handle)
 }

pkg/engine/gitlab.go

@@ -4,12 +4,10 @@ import (
 	"fmt"
 	"runtime"
 
-	"github.com/go-errors/errors"
 	gogit "github.com/go-git/go-git/v5"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/anypb"
 
-	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
@@ -52,24 +50,17 @@ func (e *Engine) ScanGitLab(ctx context.Context, c sources.GitlabConfig) error {
 		return err
 	}
 
-	gitlabSource := gitlab.Source{}
-	ctx = context.WithValues(ctx,
-		"source_type", gitlabSource.Type().String(),
-		"source_name", "gitlab",
-	)
-	err = gitlabSource.Init(ctx, "trufflehog - gitlab", 0, int64(sourcespb.SourceType_SOURCE_TYPE_GITLAB), true, &conn, runtime.NumCPU())
+	handle, err := e.sourceManager.Enroll(ctx, "trufflehog - GitLab", (&gitlab.Source{}).Type(),
+		func(ctx context.Context, jobID, sourceID int64) (sources.Source, error) {
+			gitlabSource := &gitlab.Source{}
+			gitlabSource.WithScanOptions(scanOptions)
+			if err := gitlabSource.Init(ctx, "trufflehog - GitLab", jobID, sourceID, true, &conn, runtime.NumCPU()); err != nil {
+				return nil, err
+			}
+			return gitlabSource, nil
+		})
 	if err != nil {
-		return errors.WrapPrefix(err, "could not init GitLab source", 0)
+		return err
 	}
-	gitlabSource.WithScanOptions(scanOptions)
-
-	e.sourcesWg.Go(func() error {
-		defer common.RecoverWithExit(ctx)
-		err := gitlabSource.Chunks(ctx, e.ChunksChan())
-		if err != nil {
-			return fmt.Errorf("error scanning GitLab: %w", err)
-		}
-		return nil
-	})
-	return nil
+	return e.sourceManager.ScheduleRun(ctx, handle)
 }

pkg/engine/s3.go

@@ -4,11 +4,9 @@ import (
 	"fmt"
 	"runtime"
 
-	"github.com/go-errors/errors"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/anypb"
 
-	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/credentialspb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
@@ -55,23 +53,16 @@ func (e *Engine) ScanS3(ctx context.Context, c sources.S3Config) error {
 		return err
 	}
 
-	s3Source := s3.Source{}
-	ctx = context.WithValues(ctx,
-		"source_type", s3Source.Type().String(),
-		"source_name", "s3",
-	)
-	err = s3Source.Init(ctx, "trufflehog - s3", 0, int64(sourcespb.SourceType_SOURCE_TYPE_S3), true, &conn, runtime.NumCPU())
+	handle, err := e.sourceManager.Enroll(ctx, "trufflehog - S3", (&s3.Source{}).Type(),
+		func(ctx context.Context, jobID, sourceID int64) (sources.Source, error) {
+			s3Source := &s3.Source{}
+			if err := s3Source.Init(ctx, "trufflehog - S3", jobID, sourceID, true, &conn, runtime.NumCPU()); err != nil {
+				return nil, err
+			}
+			return s3Source, nil
+		})
 	if err != nil {
-		return errors.WrapPrefix(err, "failed to init S3 source", 0)
+		return err
 	}
-
-	e.sourcesWg.Go(func() error {
-		defer common.RecoverWithExit(ctx)
-		err := s3Source.Chunks(ctx, e.ChunksChan())
-		if err != nil {
-			return fmt.Errorf("error scanning S3: %w", err)
-		}
-		return nil
-	})
-	return nil
+	return e.sourceManager.ScheduleRun(ctx, handle)
 }

pkg/engine/syslog.go

@@ -1,14 +1,12 @@
 package engine
 
 import (
-	"fmt"
 	"os"
 
 	"github.com/go-errors/errors"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/anypb"
 
-	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
@@ -42,25 +40,17 @@ func (e *Engine) ScanSyslog(ctx context.Context, c sources.SyslogConfig) error {
 	if err != nil {
 		return errors.WrapPrefix(err, "error unmarshalling connection", 0)
 	}
-	source := syslog.Source{}
-	ctx = context.WithValues(ctx,
-		"source_type", source.Type().String(),
-		"source_name", "syslog",
-	)
-	err = source.Init(ctx, "trufflehog - syslog", 0, 0, false, &conn, c.Concurrency)
-	source.InjectConnection(connection)
+
+	handle, err := e.sourceManager.Enroll(ctx, "trufflehog - syslog", (&syslog.Source{}).Type(),
+		func(ctx context.Context, jobID, sourceID int64) (sources.Source, error) {
+			source := &syslog.Source{}
+			if err := source.Init(ctx, "trufflehog - syslog", jobID, sourceID, false, &conn, c.Concurrency); err != nil {
+				return nil, err
+			}
+			return source, nil
+		})
 	if err != nil {
-		ctx.Logger().Error(err, "failed to initialize syslog source")
 		return err
 	}
-
-	e.sourcesWg.Go(func() error {
-		defer common.RecoverWithExit(ctx)
-		err := source.Chunks(ctx, e.ChunksChan())
-		if err != nil {
-			return fmt.Errorf("could not scan syslog: %w", err)
-		}
-		return nil
-	})
-	return nil
+	return e.sourceManager.ScheduleRun(ctx, handle)
 }

pkg/sources/git/git.go

@@ -41,7 +41,11 @@ type Source struct {
 	verify   bool
 	git      *Git
 	sources.Progress
-	conn *sourcespb.Git
+	conn        *sourcespb.Git
+	scanOptions *ScanOptions
+	// Kludge to preserve engine.ScanGit functionality which doesn't expect
+	// the scanning to clean up the directory.
+	preserveTempDirs bool
 }
 
 type Git struct {
@@ -86,6 +90,19 @@ func (s *Source) JobID() int64 {
 	return s.jobId
 }
 
+// WithScanOptions sets the scan options.
+func (s *Source) WithScanOptions(scanOptions *ScanOptions) {
+	s.scanOptions = scanOptions
+}
+
+// WithPreserveTempDirs sets whether to preserve temp directories when scanning
+// the provided list of s.conn.Directories. NOTE: This is *only* for
+// s.conn.Directories, not all temp directories created. This is also a kludge
+// and should be refactored away.
+func (s *Source) WithPreserveTempDirs(preserve bool) {
+	s.preserveTempDirs = preserve
+}
+
 // Init returns an initialized GitHub source.
 func (s *Source) Init(aCtx context.Context, name string, jobId, sourceId int64, verify bool, connection *anypb.Any, concurrency int) error {
 	s.name = name
@@ -129,8 +146,29 @@ func (s *Source) Init(aCtx context.Context, name string, jobId, sourceId int64,
 
 // Chunks emits chunks of bytes over a channel.
 func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) error {
-	// TODO: refactor to remove duplicate code
+	if err := s.scanRepos(ctx, chunksChan); err != nil {
+		return err
+	}
+	if err := s.scanDirs(ctx, chunksChan); err != nil {
+		return err
+	}
+
 	totalRepos := len(s.conn.Repositories) + len(s.conn.Directories)
+	ctx.Logger().V(1).Info("Git source finished scanning", "repo_count", totalRepos)
+	s.SetProgressComplete(
+		totalRepos, totalRepos,
+		fmt.Sprintf("Completed scanning source %s", s.name), "",
+	)
+	return nil
+}
+
+// scanRepos scans the configured repositories in s.conn.Repositories.
+func (s *Source) scanRepos(ctx context.Context, chunksChan chan *sources.Chunk) error {
+	if len(s.conn.Repositories) == 0 {
+		return nil
+	}
+	totalRepos := len(s.conn.Repositories) + len(s.conn.Directories)
+	// TODO: refactor to remove duplicate code
 	switch cred := s.conn.GetCredential().(type) {
 	case *sourcespb.Git_BasicAuth:
 		user := cred.BasicAuth.Username
@@ -147,7 +185,7 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
 				if err != nil {
 					return err
 				}
-				return s.git.ScanRepo(ctx, repo, path, NewScanOptions(), chunksChan)
+				return s.git.ScanRepo(ctx, repo, path, s.scanOptions, chunksChan)
 			}(repoURI)
 			if err != nil {
 				ctx.Logger().Info("error scanning repository", "repo", repoURI, "error", err)
@@ -166,7 +204,7 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
 				if err != nil {
 					return err
 				}
-				return s.git.ScanRepo(ctx, repo, path, NewScanOptions(), chunksChan)
+				return s.git.ScanRepo(ctx, repo, path, s.scanOptions, chunksChan)
 			}(repoURI)
 			if err != nil {
 				ctx.Logger().Info("error scanning repository", "repo", repoURI, "error", err)
@@ -185,7 +223,7 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
 				if err != nil {
 					return err
 				}
-				return s.git.ScanRepo(ctx, repo, path, NewScanOptions(), chunksChan)
+				return s.git.ScanRepo(ctx, repo, path, s.scanOptions, chunksChan)
 			}(repoURI)
 			if err != nil {
 				ctx.Logger().Info("error scanning repository", "repo", repoURI, "error", err)
@@ -195,41 +233,42 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
 	default:
 		return errors.New("invalid connection type for git source")
 	}
+	return nil
+}
 
+// scanDirs scans the configured directories in s.conn.Directories.
+func (s *Source) scanDirs(ctx context.Context, chunksChan chan *sources.Chunk) error {
+	totalRepos := len(s.conn.Repositories) + len(s.conn.Directories)
 	for i, gitDir := range s.conn.Directories {
 		s.SetProgressComplete(len(s.conn.Repositories)+i, totalRepos, fmt.Sprintf("Repo: %s", gitDir), "")
 
 		if len(gitDir) == 0 {
 			continue
 		}
-		if !strings.HasSuffix(gitDir, "git") {
-			// try paths instead of url
-			repo, err := RepoFromPath(gitDir)
-			if err != nil {
-				ctx.Logger().Info("error scanning repository", "repo", gitDir, "error", err)
-				continue
+		if strings.HasSuffix(gitDir, "git") {
+			// TODO: Figure out why we skip directories ending in "git".
+			continue
+		}
+		// try paths instead of url
+		repo, err := RepoFromPath(gitDir)
+		if err != nil {
+			ctx.Logger().Info("error scanning repository", "repo", gitDir, "error", err)
+			continue
+		}
+
+		err = func(repoPath string) error {
+			if !s.preserveTempDirs && strings.HasPrefix(repoPath, filepath.Join(os.TempDir(), "trufflehog")) {
+				defer os.RemoveAll(repoPath)
 			}
 
-			err = func(repoPath string) error {
-				if strings.HasPrefix(repoPath, filepath.Join(os.TempDir(), "trufflehog")) {
-					defer os.RemoveAll(repoPath)
-				}
-
-				return s.git.ScanRepo(ctx, repo, repoPath, NewScanOptions(), chunksChan)
-			}(gitDir)
-			if err != nil {
-				ctx.Logger().Info("error scanning repository", "repo", gitDir, "error", err)
-				continue
-			}
+			return s.git.ScanRepo(ctx, repo, repoPath, s.scanOptions, chunksChan)
+		}(gitDir)
+		if err != nil {
+			ctx.Logger().Info("error scanning repository", "repo", gitDir, "error", err)
+			continue
 		}
 
 	}
-
-	ctx.Logger().V(1).Info("Git source finished scanning", "repo-count", totalRepos)
-	s.SetProgressComplete(
-		totalRepos, totalRepos,
-		fmt.Sprintf("Completed scanning source %s", s.name), "",
-	)
 	return nil
 }
 

pkg/sources/source_manager.go

@@ -18,7 +18,7 @@ type handle int64
 
 // SourceInitFunc is a function that takes a source and job ID and returns an
 // initialized Source.
-type SourceInitFunc func(sourceID int64, jobID int64) (Source, error)
+type SourceInitFunc func(ctx context.Context, sourceID int64, jobID int64) (Source, error)
 
 type SourceManager struct {
 	api apiClient
@@ -50,12 +50,12 @@ func WithConcurrency(concurrency int) func(*SourceManager) {
 }
 
 // NewManager creates a new manager with the provided options.
-func NewManager(outputChunks chan *Chunk, opts ...func(*SourceManager)) *SourceManager {
+func NewManager(opts ...func(*SourceManager)) *SourceManager {
 	man := SourceManager{
 		// Default to the headless API. Can be overwritten by the WithAPI option.
 		api:          &headlessAPI{},
 		handles:      make(map[handle]SourceInitFunc),
-		outputChunks: outputChunks,
+		outputChunks: make(chan *Chunk),
 	}
 	for _, opt := range opts {
 		opt(&man)
@@ -77,7 +77,7 @@ func (s *SourceManager) Enroll(ctx context.Context, name string, kind sourcespb.
 		return 0, fmt.Errorf("handle ID '%d' already in use", handleID)
 	}
 	s.handles[handleID] = f
-	return 0, nil
+	return handleID, nil
 }
 
 // Run blocks until a resource is available to run the source, then synchronously runs it.
@@ -119,6 +119,21 @@ func (s *SourceManager) ScheduleRun(ctx context.Context, handle handle) error {
 	return 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 {
+	return s.outputChunks
+}
+
+// Wait blocks until all running sources are completed and returns the first
+// error encountered if any. It also closes the channel returned by Chunks().
+// The manager should not be reused after calling this method.
+func (s *SourceManager) Wait() error {
+	// TODO: Aggregate all errors from all sources.
+	defer close(s.outputChunks)
+	return s.pool.Wait()
+}
+
 // run is a helper method to sychronously run the source. It does not check for
 // acquired resources.
 func (s *SourceManager) run(ctx context.Context, handle handle) error {
@@ -130,7 +145,7 @@ func (s *SourceManager) run(ctx context.Context, handle handle) error {
 	if !ok {
 		return fmt.Errorf("unrecognized handle")
 	}
-	source, err := initFunc(jobID, int64(handle))
+	source, err := initFunc(ctx, jobID, int64(handle))
 	if err != nil {
 		return err
 	}

proto/detectors.proto

@@ -932,7 +932,6 @@ enum DetectorType {
   Dockerhub = 921;
   TrufflehogEnterprise = 922;
   EnvoyApiKey = 923;
-  Postgres = 925;
 }
 
 message Result {