Refactor channel to use an LRU cache instead

An LRU cache has a more favorable failure mode than the channel. With the channel, if the consumer stopped consuming metrics, scanning would block. With the LRU cache, metrics will be dropped when space runs out and a log message emitted.
Add UnitHook and NoopHook implementations
2023-10-20 22:53:05 -07:00 · 2023-10-19 23:53:20 -07:00
11 changed files with 360 additions and 276 deletions
@@ -47,6 +47,7 @@ require (
 	github.com/h2non/filetype v1.1.3
 	github.com/hashicorp/go-retryablehttp v0.7.4
 	github.com/hashicorp/golang-lru v0.5.1
+	github.com/hashicorp/golang-lru/v2 v2.0.7
 	github.com/jlaffaye/ftp v0.2.0
 	github.com/joho/godotenv v1.5.1
 	github.com/jpillora/overseer v1.1.6
@@ -418,6 +418,8 @@ github.com/hashicorp/go-retryablehttp v0.7.4/go.mod h1:Jy/gPYAdjqffZ/yFGCFV2doI5
 github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hashicorp/golang-lru v0.5.1 h1:0hERBMJE1eitiLkihrMvRVBYAkpHzc/J3QdDN+dAcgU=
 github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
+github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
+github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
 github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
 github.com/ianlancetaylor/demangle v0.0.0-20210905161508-09a460cdf81d/go.mod h1:aYm2/VgdVmcIU8iMfdMvDMsRAQjcfZSKFby6HOFvi/w=
@@ -1,4 +1,4 @@
-package glob
+package common

 import (
 	"fmt"
@@ -1,4 +1,4 @@
-package glob
+package common

 import (
 	"testing"
@@ -1,89 +0,0 @@
-package privacy
-
-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{"privacy"}) + `\b([a-z0-9]{8}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{12})\b`)
-)
-
-// 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{"privacy"}
-}
-
-// FromData will find and optionally verify Privacy 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_Privacy,
-			Raw:          []byte(resMatch),
-		}
-
-		if verify {
-			client := s.client
-			if client == nil {
-				client = defaultClient
-			}
-			req, err := http.NewRequestWithContext(ctx, "GET", "https://api.privacy.com/v1/card?page=1&page_size=50", nil)
-			if err != nil {
-				continue
-			}
-			req.Header.Set("Authorization", "api-key "+resMatch)
-			res, err := client.Do(req)
-			if err == nil {
-				defer res.Body.Close()
-				if res.StatusCode >= 200 && res.StatusCode < 300 {
-					s1.Verified = true
-				} else if res.StatusCode == 401 {
-					// The secret is determinately not verified (nothing to do)
-				} else {
-					s1.VerificationError = fmt.Errorf("unexpected HTTP response status %d", res.StatusCode)
-				}
-			} else {
-				s1.VerificationError = err
-			}
-		}
-
-		// 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 !s1.Verified && detectors.IsKnownFalsePositive(resMatch, detectors.DefaultFalsePositives, true) {
-			continue
-		}
-
-		results = append(results, s1)
-	}
-
-	return results, nil
-}
-
-func (s Scanner) Type() detectorspb.DetectorType {
-	return detectorspb.DetectorType_Privacy
-}
@@ -1,161 +0,0 @@
-//go:build detectors
-// +build detectors
-
-package privacy
-
-import (
-	"context"
-	"fmt"
-	"github.com/google/go-cmp/cmp"
-	"github.com/google/go-cmp/cmp/cmpopts"
-	"testing"
-	"time"
-
-	"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
-
-	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
-	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"
-)
-
-func TestPrivacy_FromChunk(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
-	defer cancel()
-	testSecrets, err := common.GetSecret(ctx, "trufflehog-testing", "detectors5")
-	if err != nil {
-		t.Fatalf("could not get test secrets from GCP: %s", err)
-	}
-	secret := testSecrets.MustGetField("PRIVACY")
-	inactiveSecret := testSecrets.MustGetField("PRIVACY_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 privacy secret %s within", secret)),
-				verify: true,
-			},
-			want: []detectors.Result{
-				{
-					DetectorType: detectorspb.DetectorType_Privacy,
-					Verified:     true,
-				},
-			},
-			wantErr:             false,
-			wantVerificationErr: false,
-		},
-		{
-			name: "found, unverified",
-			s:    Scanner{},
-			args: args{
-				ctx:    context.Background(),
-				data:   []byte(fmt.Sprintf("You can find a privacy 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_Privacy,
-					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 privacy secret %s within", secret)),
-				verify: true,
-			},
-			want: []detectors.Result{
-				{
-					DetectorType: detectorspb.DetectorType_Privacy,
-					Verified:     false,
-				},
-			},
-			wantErr:             false,
-			wantVerificationErr: true,
-		},
-		{
-			name: "found, verified but unexpected api surface",
-			s:    Scanner{client: common.ConstantResponseHttpClient(404, "")},
-			args: args{
-				ctx:    context.Background(),
-				data:   []byte(fmt.Sprintf("You can find a privacy secret %s within", secret)),
-				verify: true,
-			},
-			want: []detectors.Result{
-				{
-					DetectorType: detectorspb.DetectorType_Privacy,
-					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("Privacy.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])
-				}
-				if (got[i].VerificationError != nil) != tt.wantVerificationErr {
-					t.Fatalf("wantVerificationError = %v, verification error = %v", tt.wantVerificationErr, got[i].VerificationError)
-				}
-			}
-			ignoreOpts := cmpopts.IgnoreFields(detectors.Result{}, "Raw", "VerificationError")
-			if diff := cmp.Diff(got, tt.want, ignoreOpts); diff != "" {
-				t.Errorf("Privacy.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) {
-			b.ResetTimer()
-			for n := 0; n < b.N; n++ {
-				_, err := s.FromData(ctx, false, data)
-				if err != nil {
-					b.Fatal(err)
-				}
-			}
-		})
-	}
-}
@@ -1009,7 +1009,6 @@ const (
 	DetectorType_Klaviyo                       DetectorType = 935
 	DetectorType_SourcegraphCody               DetectorType = 936
 	DetectorType_Voiceflow                     DetectorType = 937
-	DetectorType_Privacy                       DetectorType = 938
 )

 // Enum value maps for DetectorType.
@@ -1949,7 +1948,6 @@ var (
 		935: "Klaviyo",
 		936: "SourcegraphCody",
 		937: "Voiceflow",
-		938: "Privacy",
 	}
 	DetectorType_value = map[string]int32{
 		"Alibaba":                       0,
@@ -2886,7 +2884,6 @@ var (
 		"Klaviyo":                       935,
 		"SourcegraphCody":               936,
 		"Voiceflow":                     937,
-		"Privacy":                       938,
 	}
 )

@@ -3265,7 +3262,7 @@ var file_detectors_proto_rawDesc = []byte{
 	0x44, 0x65, 0x63, 0x6f, 0x64, 0x65, 0x72, 0x54, 0x79, 0x70, 0x65, 0x12, 0x0b, 0x0a, 0x07, 0x55,
 	0x4e, 0x4b, 0x4e, 0x4f, 0x57, 0x4e, 0x10, 0x00, 0x12, 0x09, 0x0a, 0x05, 0x50, 0x4c, 0x41, 0x49,
 	0x4e, 0x10, 0x01, 0x12, 0x0a, 0x0a, 0x06, 0x42, 0x41, 0x53, 0x45, 0x36, 0x34, 0x10, 0x02, 0x12,
-	0x09, 0x0a, 0x05, 0x55, 0x54, 0x46, 0x31, 0x36, 0x10, 0x03, 0x2a, 0xc4, 0x75, 0x0a, 0x0c, 0x44,
+	0x09, 0x0a, 0x05, 0x55, 0x54, 0x46, 0x31, 0x36, 0x10, 0x03, 0x2a, 0xb6, 0x75, 0x0a, 0x0c, 0x44,
 	0x65, 0x74, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x54, 0x79, 0x70, 0x65, 0x12, 0x0b, 0x0a, 0x07, 0x41,
 	0x6c, 0x69, 0x62, 0x61, 0x62, 0x61, 0x10, 0x00, 0x12, 0x08, 0x0a, 0x04, 0x41, 0x4d, 0x51, 0x50,
 	0x10, 0x01, 0x12, 0x07, 0x0a, 0x03, 0x41, 0x57, 0x53, 0x10, 0x02, 0x12, 0x09, 0x0a, 0x05, 0x41,
@@ -4205,12 +4202,11 @@ var file_detectors_proto_rawDesc = []byte{
 	0x0c, 0x0a, 0x07, 0x4b, 0x6c, 0x61, 0x76, 0x69, 0x79, 0x6f, 0x10, 0xa7, 0x07, 0x12, 0x14, 0x0a,
 	0x0f, 0x53, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x67, 0x72, 0x61, 0x70, 0x68, 0x43, 0x6f, 0x64, 0x79,
 	0x10, 0xa8, 0x07, 0x12, 0x0e, 0x0a, 0x09, 0x56, 0x6f, 0x69, 0x63, 0x65, 0x66, 0x6c, 0x6f, 0x77,
-	0x10, 0xa9, 0x07, 0x12, 0x0c, 0x0a, 0x07, 0x50, 0x72, 0x69, 0x76, 0x61, 0x63, 0x79, 0x10, 0xaa,
-	0x07, 0x42, 0x3d, 0x5a, 0x3b, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f,
-	0x74, 0x72, 0x75, 0x66, 0x66, 0x6c, 0x65, 0x73, 0x65, 0x63, 0x75, 0x72, 0x69, 0x74, 0x79, 0x2f,
-	0x74, 0x72, 0x75, 0x66, 0x66, 0x6c, 0x65, 0x68, 0x6f, 0x67, 0x2f, 0x76, 0x33, 0x2f, 0x70, 0x6b,
-	0x67, 0x2f, 0x70, 0x62, 0x2f, 0x64, 0x65, 0x74, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73, 0x70, 0x62,
-	0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
+	0x10, 0xa9, 0x07, 0x42, 0x3d, 0x5a, 0x3b, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f,
+	0x6d, 0x2f, 0x74, 0x72, 0x75, 0x66, 0x66, 0x6c, 0x65, 0x73, 0x65, 0x63, 0x75, 0x72, 0x69, 0x74,
+	0x79, 0x2f, 0x74, 0x72, 0x75, 0x66, 0x66, 0x6c, 0x65, 0x68, 0x6f, 0x67, 0x2f, 0x76, 0x33, 0x2f,
+	0x70, 0x6b, 0x67, 0x2f, 0x70, 0x62, 0x2f, 0x64, 0x65, 0x74, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x73,
+	0x70, 0x62, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
 }

 var (
@@ -222,10 +222,7 @@ func (s *Source) Init(aCtx context.Context, name string, jobID sources.JobID, so
 	}
 	s.conn = &conn

-	s.filteredRepoCache = s.newFilteredRepoCache(memory.New(),
-		append(s.conn.GetRepositories(), s.conn.GetIncludeRepos()...),
-		s.conn.GetIgnoreRepos(),
-	)
+	s.filteredRepoCache = s.newFilteredRepoCache(memory.New(), s.conn.GetRepositories(), s.conn.GetIgnoreRepos())
 	s.memberCache = make(map[string]struct{})

 	s.repoSizes = newRepoSize()
@@ -511,12 +508,12 @@ func (s *Source) enumerateUnauthenticated(ctx context.Context, apiEndpoint strin

 	for _, org := range s.orgsCache.Keys() {
 		if err := s.getReposByOrg(ctx, org); err != nil {
-			s.log.Error(err, "error fetching repos for org")
+			s.log.Error(err, "error fetching repos for org or user")
 		}

 		// We probably don't need to do this, since getting repos by org makes more sense?
 		if err := s.getReposByUser(ctx, org); err != nil {
-			s.log.Error(err, "error fetching repos for user")
+			s.log.Error(err, "error fetching repos for org or user")
 		}

 		if s.conn.ScanUsers {
@@ -1092,7 +1089,7 @@ func (s *Source) scanComments(ctx context.Context, repoPath string, chunksChan c
 }

 func (s *Source) processGistComments(ctx context.Context, repoPath string, trimmedURL []string, repoURL *url.URL, chunksChan chan *sources.Chunk) error {
-	ctx.Logger().V(2).Info("scanning github gist comments", "repository", repoPath)
+	s.log.Info("scanning github gist comments", "repository", repoPath)
 	// GitHub Gist URL.
 	gistID, err := extractGistID(trimmedURL)
 	if err != nil {
@@ -1173,21 +1170,19 @@ func (s *Source) processRepoComments(ctx context.Context, repoPath string, trimm
 	}

 	if s.includeIssueComments {
-		ctx.Logger().V(2).Info("scanning github issues", "repository", repoInfo.repoPath)
-		if err := s.processIssues(ctx, repoInfo, chunksChan); err != nil {
+		if err := s.processIssueComments(ctx, repoInfo, chunksChan); err != nil {
 			return err
 		}
-		if err := s.processIssueComments(ctx, repoInfo, chunksChan); err != nil {
+		if err := s.processIssues(ctx, repoInfo, chunksChan); err != nil {
 			return err
 		}
 	}

 	if s.includePRComments {
-		ctx.Logger().V(2).Info("scanning github pull requests", "repository", repoInfo.repoPath)
-		if err := s.processPRs(ctx, repoInfo, chunksChan); err != nil {
+		if err := s.processPRComments(ctx, repoInfo, chunksChan); err != nil {
 			return err
 		}
-		if err := s.processPRComments(ctx, repoInfo, chunksChan); err != nil {
+		if err := s.processPRs(ctx, repoInfo, chunksChan); err != nil {
 			return err
 		}
 	}
@@ -1197,6 +1192,8 @@ func (s *Source) processRepoComments(ctx context.Context, repoPath string, trimm
 }

 func (s *Source) processIssues(ctx context.Context, info repoInfo, chunksChan chan *sources.Chunk) error {
+	s.log.Info("scanning github issue descriptions", "repository", info.repoPath)
+
 	bodyTextsOpts := &github.IssueListByRepoOptions{
 		Sort:      sortType,
 		Direction: directionType,
@@ -1231,6 +1228,8 @@ func (s *Source) processIssues(ctx context.Context, info repoInfo, chunksChan ch
 }

 func (s *Source) processIssueComments(ctx context.Context, info repoInfo, chunksChan chan *sources.Chunk) error {
+	s.log.Info("scanning github issue comments", "repository", info.repoPath)
+
 	issueOpts := &github.IssueListCommentsOptions{
 		Sort:      &sortType,
 		Direction: &directionType,
@@ -1264,6 +1263,8 @@ func (s *Source) processIssueComments(ctx context.Context, info repoInfo, chunks
 }

 func (s *Source) processPRs(ctx context.Context, info repoInfo, chunksChan chan *sources.Chunk) error {
+	s.log.Info("scanning github pull request descriptions", "repository", info.repoPath)
+
 	prOpts := &github.PullRequestListOptions{
 		Sort:      sortType,
 		Direction: directionType,
@@ -1298,6 +1299,8 @@ func (s *Source) processPRs(ctx context.Context, info repoInfo, chunksChan chan
 }

 func (s *Source) processPRComments(ctx context.Context, info repoInfo, chunksChan chan *sources.Chunk) error {
+	s.log.Info("scanning github pull request comments", "repository", info.repoPath)
+
 	prOpts := &github.PullRequestListCommentsOptions{
 		Sort:      sortType,
 		Direction: directionType,
@@ -0,0 +1,218 @@
+package sources
+
+import (
+	"errors"
+	"fmt"
+	"strings"
+	"sync"
+	"time"
+
+	lru "github.com/hashicorp/golang-lru/v2"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
+)
+
+// UnitHook implements JobProgressHook for tracking the progress of each
+// individual unit.
+type UnitHook struct {
+	metrics *lru.Cache[string, *UnitMetrics]
+	mu      sync.Mutex
+	NoopHook
+}
+
+type UnitHookOpt func(*UnitHook)
+
+func WithUnitHookCache(cache *lru.Cache[string, *UnitMetrics]) UnitHookOpt {
+	return func(hook *UnitHook) { hook.metrics = cache }
+}
+
+func NewUnitHook(ctx context.Context, opts ...UnitHookOpt) *UnitHook {
+	// lru.NewWithEvict can only fail if the size is < 0.
+	cache, _ := lru.NewWithEvict(1024, func(key string, value *UnitMetrics) {
+		if value.handled {
+			return
+		}
+		ctx.Logger().Error(fmt.Errorf("eviction"), "dropping unit metric",
+			"id", key,
+			"metric", value,
+		)
+	})
+	hook := UnitHook{metrics: cache}
+	for _, opt := range opts {
+		opt(&hook)
+	}
+	return &hook
+}
+
+// id is a helper method to generate an ID for the given job and unit.
+func (u *UnitHook) id(ref JobProgressRef, unit SourceUnit) string {
+	unitID := ""
+	if unit != nil {
+		unitID = unit.SourceUnitID()
+	}
+	return fmt.Sprintf("%d/%d/%s", ref.SourceID, ref.JobID, unitID)
+}
+
+func (u *UnitHook) StartUnitChunking(ref JobProgressRef, unit SourceUnit, start time.Time) {
+	id := u.id(ref, unit)
+	u.mu.Lock()
+	defer u.mu.Unlock()
+
+	u.metrics.Add(id, &UnitMetrics{
+		Unit:      unit,
+		Parent:    ref,
+		StartTime: start,
+	})
+}
+
+func (u *UnitHook) EndUnitChunking(ref JobProgressRef, unit SourceUnit, end time.Time) {
+	id := u.id(ref, unit)
+	u.mu.Lock()
+	defer u.mu.Unlock()
+
+	metrics, ok := u.metrics.Get(id)
+	if !ok {
+		return
+	}
+	metrics.EndTime = end
+}
+
+func (u *UnitHook) ReportChunk(ref JobProgressRef, unit SourceUnit, chunk *Chunk) {
+	id := u.id(ref, unit)
+	u.mu.Lock()
+	defer u.mu.Unlock()
+
+	metrics, ok := u.metrics.Get(id)
+	if !ok && unit != nil {
+		// The unit has been evicted.
+		return
+	} else if !ok && unit == nil {
+		// This is a chunk from a non-unit source.
+		metrics = &UnitMetrics{
+			Unit:      nil,
+			Parent:    ref,
+			StartTime: ref.Snapshot().StartTime,
+		}
+		u.metrics.Add(id, metrics)
+	}
+	metrics.TotalChunks++
+	metrics.TotalBytes += uint64(len(chunk.Data))
+}
+
+func (u *UnitHook) ReportError(ref JobProgressRef, err error) {
+	u.mu.Lock()
+	defer u.mu.Unlock()
+
+	// Always add the error to the nil unit if it exists.
+	if metrics, ok := u.metrics.Get(u.id(ref, nil)); ok {
+		metrics.Errors = append(metrics.Errors, err)
+	}
+
+	// Check if it's a ChunkError for a specific unit.
+	var chunkErr ChunkError
+	if !errors.As(err, &chunkErr) {
+		return
+	}
+	id := u.id(ref, chunkErr.Unit)
+
+	metrics, ok := u.metrics.Get(id)
+	if !ok {
+		return
+	}
+	metrics.Errors = append(metrics.Errors, err)
+}
+
+func (u *UnitHook) Finish(ref JobProgressRef) {
+	u.mu.Lock()
+	defer u.mu.Unlock()
+	// Clear out any metrics on this job. This covers the case for the
+	// source running without unit support.
+	prefix := u.id(ref, nil)
+	for _, id := range u.metrics.Keys() {
+		if !strings.HasPrefix(id, prefix) {
+			continue
+		}
+		metric, ok := u.metrics.Get(id)
+		if !ok {
+			continue
+		}
+		// If the unit is nil, the source does not support units.
+		// Use the overall job metrics instead.
+		if metric.Unit == nil {
+			snap := ref.Snapshot()
+			metric.StartTime = snap.StartTime
+			metric.EndTime = snap.EndTime
+			metric.Errors = snap.Errors
+		}
+	}
+}
+
+// UnitMetrics gets all the currently active or newly finished metrics for this
+// job. If a unit returned from this method has finished, it will be removed
+// from the cache and no longer returned in successive calls to UnitMetrics().
+func (u *UnitHook) UnitMetrics() []UnitMetrics {
+	u.mu.Lock()
+	defer u.mu.Unlock()
+	output := make([]UnitMetrics, 0, u.metrics.Len())
+	for _, id := range u.metrics.Keys() {
+		metric, ok := u.metrics.Get(id)
+		if !ok {
+			continue
+		}
+		output = append(output, *metric)
+		if metric.IsFinished() {
+			metric.handled = true
+			u.metrics.Remove(id)
+		}
+	}
+	return output
+}
+
+type UnitMetrics struct {
+	Unit   SourceUnit
+	Parent JobProgressRef
+	// Start and end time for chunking this unit.
+	StartTime time.Time
+	EndTime   time.Time
+	// Total number of chunks produced from this unit.
+	TotalChunks uint64
+	// Total number of bytes produced from this unit.
+	TotalBytes uint64
+	// All errors encountered by this unit.
+	Errors []error
+	// Flag to mark that these metrics were intentionally evicted from
+	// the cache.
+	handled bool
+}
+
+func (u UnitMetrics) IsFinished() bool {
+	return !u.EndTime.IsZero()
+}
+
+// ElapsedTime is a convenience method that provides the elapsed time the job
+// has been running. If it hasn't started yet, 0 is returned. If it has
+// finished, the total time is returned.
+func (u UnitMetrics) ElapsedTime() time.Duration {
+	if u.StartTime.IsZero() {
+		return 0
+	}
+	if u.EndTime.IsZero() {
+		return time.Since(u.StartTime)
+	}
+	return u.EndTime.Sub(u.StartTime)
+}
+
+// NoopHook implements JobProgressHook by doing nothing. This is useful for
+// embedding in other structs to overwrite only the methods of the interface
+// that you care about.
+type NoopHook struct{}
+
+func (NoopHook) Start(JobProgressRef, time.Time)                         {}
+func (NoopHook) End(JobProgressRef, time.Time)                           {}
+func (NoopHook) StartEnumerating(JobProgressRef, time.Time)              {}
+func (NoopHook) EndEnumerating(JobProgressRef, time.Time)                {}
+func (NoopHook) StartUnitChunking(JobProgressRef, SourceUnit, time.Time) {}
+func (NoopHook) EndUnitChunking(JobProgressRef, SourceUnit, time.Time)   {}
+func (NoopHook) ReportError(JobProgressRef, error)                       {}
+func (NoopHook) ReportUnit(JobProgressRef, SourceUnit)                   {}
+func (NoopHook) ReportChunk(JobProgressRef, SourceUnit, *Chunk)          {}
+func (NoopHook) Finish(JobProgressRef)                                   {}
@@ -3,8 +3,10 @@ package sources
 import (
 	"errors"
 	"fmt"
+	"sort"
 	"testing"

+	lru "github.com/hashicorp/golang-lru/v2"
 	"github.com/stretchr/testify/assert"
 	"google.golang.org/protobuf/types/known/anypb"

@@ -200,6 +202,8 @@ func (c *unitChunker) ChunkUnit(ctx context.Context, unit SourceUnit, rep ChunkR
 			if err := rep.ChunkErr(ctx, fmt.Errorf(step.err)); err != nil {
 				return err
 			}
+		}
+		if step.output == "" {
 			continue
 		}
 		if err := rep.ChunkOk(ctx, Chunk{Data: []byte(step.output)}); err != nil {
@@ -310,3 +314,114 @@ func TestSourceManagerAvailableCapacity(t *testing.T) {
 	<-ref.Done()      // Wait for the job to finish.
 	assert.Equal(t, 1337, mgr.AvailableCapacity())
 }
+
+func TestSourceManagerUnitHook(t *testing.T) {
+	hook := NewUnitHook(context.TODO())
+
+	input := []unitChunk{
+		{unit: "one", output: "bar"},
+		{unit: "two", err: "oh no"},
+		{unit: "three", err: "not again"},
+	}
+	mgr := NewManager(
+		WithBufferedOutput(8),
+		WithSourceUnits(), WithConcurrentUnits(1),
+		WithReportHook(hook),
+	)
+	source, err := buildDummy(&unitChunker{input})
+	assert.NoError(t, err)
+	ref, err := mgr.Run(context.Background(), "dummy", source)
+	assert.NoError(t, err)
+	<-ref.Done()
+
+	metrics := hook.UnitMetrics()
+	assert.Equal(t, 3, len(metrics))
+	sort.Slice(metrics, func(i, j int) bool {
+		return metrics[i].EndTime.Before(metrics[j].EndTime)
+	})
+	m0, m1, m2 := metrics[0], metrics[1], metrics[2]
+
+	assert.Equal(t, "one", m0.Unit.SourceUnitID())
+	assert.Equal(t, uint64(1), m0.TotalChunks)
+	assert.Equal(t, uint64(3), m0.TotalBytes)
+	assert.NotZero(t, m0.StartTime)
+	assert.NotZero(t, m0.EndTime)
+	assert.NotZero(t, m0.ElapsedTime())
+	assert.Equal(t, 0, len(m0.Errors))
+
+	assert.Equal(t, "two", m1.Unit.SourceUnitID())
+	assert.Equal(t, uint64(0), m1.TotalChunks)
+	assert.Equal(t, uint64(0), m1.TotalBytes)
+	assert.NotZero(t, m1.StartTime)
+	assert.NotZero(t, m1.EndTime)
+	assert.NotZero(t, m1.ElapsedTime())
+	assert.Equal(t, 1, len(m1.Errors))
+
+	assert.Equal(t, "three", m2.Unit.SourceUnitID())
+	assert.Equal(t, uint64(0), m2.TotalChunks)
+	assert.Equal(t, uint64(0), m2.TotalBytes)
+	assert.NotZero(t, m2.StartTime)
+	assert.NotZero(t, m2.EndTime)
+	assert.NotZero(t, m2.ElapsedTime())
+	assert.Equal(t, 1, len(m2.Errors))
+}
+
+// TestSourceManagerUnitHookNoBlock tests that the UnitHook drops metrics if
+// they aren't handled fast enough.
+func TestSourceManagerUnitHookNoBlock(t *testing.T) {
+	var evictedKeys []string
+	cache, _ := lru.NewWithEvict(1, func(key string, _ *UnitMetrics) {
+		evictedKeys = append(evictedKeys, key)
+	})
+	hook := NewUnitHook(context.TODO(), WithUnitHookCache(cache))
+
+	input := []unitChunk{
+		{unit: "one", output: "bar"},
+		{unit: "two", err: "oh no"},
+		{unit: "three", err: "not again"},
+	}
+	mgr := NewManager(
+		WithBufferedOutput(8),
+		WithSourceUnits(), WithConcurrentUnits(1),
+		WithReportHook(hook),
+	)
+	source, err := buildDummy(&unitChunker{input})
+	assert.NoError(t, err)
+	ref, err := mgr.Run(context.Background(), "dummy", source)
+	assert.NoError(t, err)
+	<-ref.Done()
+
+	assert.Equal(t, 2, len(evictedKeys))
+	metrics := hook.UnitMetrics()
+	assert.Equal(t, 1, len(metrics))
+	assert.Equal(t, "three", metrics[0].Unit.SourceUnitID())
+}
+
+// TestSourceManagerUnitHookNoUnits tests whether the UnitHook works for
+// sources that don't support units.
+func TestSourceManagerUnitHookNoUnits(t *testing.T) {
+	hook := NewUnitHook(context.TODO())
+
+	mgr := NewManager(
+		WithBufferedOutput(8),
+		WithReportHook(hook),
+	)
+	source, err := buildDummy(&counterChunker{count: 5})
+	assert.NoError(t, err)
+
+	ref, err := mgr.Run(context.Background(), "dummy", source)
+	assert.NoError(t, err)
+	<-ref.Done()
+
+	metrics := hook.UnitMetrics()
+	assert.Equal(t, 1, len(metrics))
+
+	m := metrics[0]
+	assert.Equal(t, nil, m.Unit)
+	assert.Equal(t, uint64(5), m.TotalChunks)
+	assert.Equal(t, uint64(5), m.TotalBytes)
+	assert.NotZero(t, m.StartTime)
+	assert.NotZero(t, m.EndTime)
+	assert.NotZero(t, m.ElapsedTime())
+	assert.Equal(t, 0, len(m.Errors))
+}
@@ -946,7 +946,6 @@ enum DetectorType {
  Klaviyo = 935;
  SourcegraphCody = 936;
  Voiceflow = 937;
-  Privacy = 938;
 }

 message Result {