From c60443891b4383505facfcda0fe1595a60ccdd4b Mon Sep 17 00:00:00 2001
From: Miccah <m.castorina93@gmail.com>
Date: Tue, 20 Feb 2024 11:24:13 -0800
Subject: [PATCH] Add Display method to SourceUnit and Kind member to the
 CommonSourceUnit (#2450)

* Add Display method to SourceUnit and Kind member to the CommonSourceUnit

* Make SourceUnitID return the ID and a kind

These two values together uniquely represent a unit.
---
 pkg/sources/filesystem/filesystem.go      |  2 +-
 pkg/sources/filesystem/filesystem_test.go |  3 +-
 pkg/sources/git/git.go                    | 13 +++----
 pkg/sources/git/git_test.go               |  3 +-
 pkg/sources/git/unit.go                   | 37 ++++++++++++++++----
 pkg/sources/git/unit_test.go              | 42 +++++++++++++++++++++++
 pkg/sources/gitlab/gitlab.go              | 12 ++++---
 pkg/sources/job_progress.go               |  6 ++--
 pkg/sources/job_progress_hook.go          |  3 +-
 pkg/sources/job_progress_test.go          |  2 +-
 pkg/sources/source_manager.go             |  3 +-
 pkg/sources/source_manager_test.go        | 26 ++++++++++----
 pkg/sources/source_unit.go                | 13 +++++--
 pkg/sources/sources.go                    | 11 ++++--
 pkg/sources/travisci/travisci.go          |  6 ++--
 15 files changed, 141 insertions(+), 41 deletions(-)
 create mode 100644 pkg/sources/git/unit_test.go

diff --git a/pkg/sources/filesystem/filesystem.go b/pkg/sources/filesystem/filesystem.go
index 1d3d0cc8b..63aaf9bda 100644
--- a/pkg/sources/filesystem/filesystem.go
+++ b/pkg/sources/filesystem/filesystem.go
@@ -270,7 +270,7 @@ func (s *Source) Enumerate(ctx context.Context, reporter sources.UnitReporter) e
 
 // ChunkUnit implements SourceUnitChunker interface.
 func (s *Source) ChunkUnit(ctx context.Context, unit sources.SourceUnit, reporter sources.ChunkReporter) error {
-	path := unit.SourceUnitID()
+	path, _ := unit.SourceUnitID()
 	logger := ctx.Logger().WithValues("path", path)
 
 	cleanPath := filepath.Clean(path)
diff --git a/pkg/sources/filesystem/filesystem_test.go b/pkg/sources/filesystem/filesystem_test.go
index 4a846a673..a5ce36122 100644
--- a/pkg/sources/filesystem/filesystem_test.go
+++ b/pkg/sources/filesystem/filesystem_test.go
@@ -172,7 +172,8 @@ func TestEnumerate(t *testing.T) {
 	assert.Equal(t, len(units), len(reporter.Units))
 	assert.Equal(t, 0, len(reporter.UnitErrs))
 	for _, unit := range reporter.Units {
-		assert.Contains(t, units, unit.SourceUnitID())
+		path, _ := unit.SourceUnitID()
+		assert.Contains(t, units, path)
 	}
 	for _, unit := range units {
 		assert.Contains(t, reporter.Units, sources.CommonSourceUnit{ID: unit})
diff --git a/pkg/sources/git/git.go b/pkg/sources/git/git.go
index af3d93eb9..c60dc4032 100644
--- a/pkg/sources/git/git.go
+++ b/pkg/sources/git/git.go
@@ -1264,18 +1264,15 @@ func (s *Source) Enumerate(ctx context.Context, reporter sources.UnitReporter) e
 }
 
 func (s *Source) ChunkUnit(ctx context.Context, unit sources.SourceUnit, reporter sources.ChunkReporter) error {
-	gitUnit, ok := unit.(SourceUnit)
-	if !ok {
-		return fmt.Errorf("unsupported unit type: %T", unit)
-	}
+	unitID, kind := unit.SourceUnitID()
 
-	switch gitUnit.Kind {
+	switch kind {
 	case UnitRepo:
-		return s.scanRepo(ctx, gitUnit.ID, reporter)
+		return s.scanRepo(ctx, unitID, reporter)
 	case UnitDir:
-		return s.scanDir(ctx, gitUnit.ID, reporter)
+		return s.scanDir(ctx, unitID, reporter)
 	default:
-		return fmt.Errorf("unexpected git unit kind: %q", gitUnit.Kind)
+		return fmt.Errorf("unexpected git unit kind: %q", kind)
 	}
 }
 
diff --git a/pkg/sources/git/git_test.go b/pkg/sources/git/git_test.go
index cc9d31756..324dac695 100644
--- a/pkg/sources/git/git_test.go
+++ b/pkg/sources/git/git_test.go
@@ -533,7 +533,8 @@ func TestEnumerate(t *testing.T) {
 	assert.Equal(t, len(units), len(reporter.Units))
 	assert.Equal(t, 0, len(reporter.UnitErrs))
 	for _, unit := range reporter.Units {
-		assert.Contains(t, units, unit.SourceUnitID())
+		id, _ := unit.SourceUnitID()
+		assert.Contains(t, units, id)
 	}
 	for _, unit := range units[:3] {
 		assert.Contains(t, reporter.Units, SourceUnit{ID: unit, Kind: UnitRepo})
diff --git a/pkg/sources/git/unit.go b/pkg/sources/git/unit.go
index 9457e23a6..990866b40 100644
--- a/pkg/sources/git/unit.go
+++ b/pkg/sources/git/unit.go
@@ -3,13 +3,16 @@ package git
 import (
 	"encoding/json"
 	"fmt"
+	"net/url"
+	"path/filepath"
+	"strings"
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
 )
 
 const (
-	UnitRepo string = "repo"
-	UnitDir  string = "dir"
+	UnitRepo sources.SourceUnitKind = "repo"
+	UnitDir  sources.SourceUnitKind = "dir"
 )
 
 // Ensure SourceUnit implements the interface at compile time.
@@ -18,13 +21,35 @@ var _ sources.SourceUnit = SourceUnit{}
 // A git source unit can be two kinds of units: either a local directory path
 // or a remote repository.
 type SourceUnit struct {
-	Kind string `json:"kind"`
-	ID   string `json:"id"`
+	Kind sources.SourceUnitKind `json:"kind"`
+	ID   string                 `json:"id"`
 }
 
 // Implement sources.SourceUnit interface.
-func (u SourceUnit) SourceUnitID() string {
-	return u.ID
+func (u SourceUnit) SourceUnitID() (string, sources.SourceUnitKind) {
+	return u.ID, u.Kind
+}
+
+// Provide a custom Display method.
+func (u SourceUnit) Display() string {
+	switch u.Kind {
+	case UnitRepo:
+		repo := u.ID
+		if parsedURL, err := url.Parse(u.ID); err == nil {
+			// scheme://host/owner/repo
+			repo = strings.TrimPrefix(parsedURL.Path, "/")
+		} else if _, path, found := strings.Cut(u.ID, ":"); found {
+			// git@host:owner/repo
+			// TODO: Is this possible? We should maybe canonicalize
+			// the URL before getting here.
+			repo = path
+		}
+		return strings.TrimSuffix(repo, ".git")
+	case UnitDir:
+		return filepath.Base(u.ID)
+	default:
+		return "mysterious git unit"
+	}
 }
 
 // Helper function to unmarshal raw bytes into our SourceUnit struct.
diff --git a/pkg/sources/git/unit_test.go b/pkg/sources/git/unit_test.go
new file mode 100644
index 000000000..e5ac45843
--- /dev/null
+++ b/pkg/sources/git/unit_test.go
@@ -0,0 +1,42 @@
+package git
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestUnmarshalUnit(t *testing.T) {
+	s := `{"kind":"repo","id":"https://github.com/trufflesecurity/test_keys.git"}`
+	expectedUnit := SourceUnit{ID: "https://github.com/trufflesecurity/test_keys.git", Kind: UnitRepo}
+	gotUnit, err := UnmarshalUnit([]byte(s))
+	assert.NoError(t, err)
+
+	assert.Equal(t, expectedUnit, gotUnit)
+
+	_, err = UnmarshalUnit(nil)
+	assert.Error(t, err)
+
+	_, err = UnmarshalUnit([]byte(`{"kind":"idk","id":"id"}`))
+	assert.Error(t, err)
+}
+
+func TestMarshalUnit(t *testing.T) {
+	unit := SourceUnit{ID: "https://github.com/trufflesecurity/test_keys.git", Kind: UnitRepo}
+	b, err := json.Marshal(unit)
+	assert.NoError(t, err)
+
+	assert.Equal(t, `{"kind":"repo","id":"https://github.com/trufflesecurity/test_keys.git"}`, string(b))
+}
+
+func TestDisplayUnit(t *testing.T) {
+	unit := SourceUnit{ID: "https://github.com/trufflesecurity/test_keys.git", Kind: UnitRepo}
+	assert.Equal(t, "trufflesecurity/test_keys", unit.Display())
+
+	unit = SourceUnit{ID: "/path/to/repo", Kind: UnitDir}
+	assert.Equal(t, "repo", unit.Display())
+
+	unit = SourceUnit{ID: "ssh://github.com/trufflesecurity/test_keys", Kind: UnitRepo}
+	assert.Equal(t, "trufflesecurity/test_keys", unit.Display())
+}
diff --git a/pkg/sources/gitlab/gitlab.go b/pkg/sources/gitlab/gitlab.go
index 850d153b3..cce768a34 100644
--- a/pkg/sources/gitlab/gitlab.go
+++ b/pkg/sources/gitlab/gitlab.go
@@ -186,7 +186,8 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk, _ .
 		})
 		reporter := sources.VisitorReporter{
 			VisitUnit: func(ctx context.Context, unit sources.SourceUnit) error {
-				repos = append(repos, unit.SourceUnitID())
+				id, _ := unit.SourceUnitID()
+				repos = append(repos, id)
 				return ctx.Err()
 			},
 		}
@@ -255,7 +256,8 @@ func (s *Source) Validate(ctx context.Context) []error {
 	var repos []string
 	visitor := sources.VisitorReporter{
 		VisitUnit: func(ctx context.Context, unit sources.SourceUnit) error {
-			repos = append(repos, unit.SourceUnitID())
+			id, _ := unit.SourceUnitID()
+			repos = append(repos, id)
 			return nil
 		},
 	}
@@ -362,7 +364,7 @@ func (s *Source) getAllProjectRepos(
 				continue
 			}
 			// Report the unit.
-			unit := sources.CommonSourceUnit{ID: proj.HTTPURLToRepo}
+			unit := git.SourceUnit{Kind: git.UnitRepo, ID: proj.HTTPURLToRepo}
 			if err := reporter.UnitOk(ctx, unit); err != nil {
 				return err
 			}
@@ -616,7 +618,7 @@ func (s *Source) Enumerate(ctx context.Context, reporter sources.UnitReporter) e
 	// Report all repos if specified.
 	if len(repos) > 0 {
 		for _, repo := range repos {
-			unit := sources.CommonSourceUnit{ID: repo}
+			unit := git.SourceUnit{Kind: git.UnitRepo, ID: repo}
 			if err := reporter.UnitOk(ctx, unit); err != nil {
 				return err
 			}
@@ -635,7 +637,7 @@ func (s *Source) Enumerate(ctx context.Context, reporter sources.UnitReporter) e
 
 // ChunkUnit downloads and reports chunks for the given GitLab repository unit.
 func (s *Source) ChunkUnit(ctx context.Context, unit sources.SourceUnit, reporter sources.ChunkReporter) error {
-	repoURL := unit.SourceUnitID()
+	repoURL, _ := unit.SourceUnitID()
 
 	var path string
 	var repo *gogit.Repository
diff --git a/pkg/sources/job_progress.go b/pkg/sources/job_progress.go
index 430804d13..08d609c3e 100644
--- a/pkg/sources/job_progress.go
+++ b/pkg/sources/job_progress.go
@@ -94,11 +94,11 @@ type ChunkError struct {
 }
 
 func (f ChunkError) Error() string {
-	unit := "<nil>"
+	unit, kind := "<nil>", SourceUnitKind("unit")
 	if f.Unit != nil {
-		unit = f.Unit.SourceUnitID()
+		unit, kind = f.Unit.SourceUnitID()
 	}
-	return fmt.Sprintf("error chunking unit %q: %s", unit, f.Err.Error())
+	return fmt.Sprintf("error chunking %s %q: %s", kind, unit, f.Err.Error())
 }
 func (f ChunkError) Unwrap() error { return f.Err }
 
diff --git a/pkg/sources/job_progress_hook.go b/pkg/sources/job_progress_hook.go
index 0f6c45810..6f94feff2 100644
--- a/pkg/sources/job_progress_hook.go
+++ b/pkg/sources/job_progress_hook.go
@@ -64,7 +64,8 @@ func NewUnitHook(ctx context.Context, opts ...UnitHookOpt) (*UnitHook, <-chan Un
 func (u *UnitHook) id(ref JobProgressRef, unit SourceUnit) string {
 	unitID := ""
 	if unit != nil {
-		unitID = unit.SourceUnitID()
+		id, kind := unit.SourceUnitID()
+		unitID = fmt.Sprintf("%s:%s", kind, id)
 	}
 	return fmt.Sprintf("%d/%d/%s", ref.SourceID, ref.JobID, unitID)
 }
diff --git a/pkg/sources/job_progress_test.go b/pkg/sources/job_progress_test.go
index 05d4c6e6a..20ba06785 100644
--- a/pkg/sources/job_progress_test.go
+++ b/pkg/sources/job_progress_test.go
@@ -81,7 +81,7 @@ func TestJobProgressHook(t *testing.T) {
 	startChunk := time.Now().Add(40 * time.Second)
 	endChunk := time.Now().Add(50 * time.Second)
 	reportErr := fmt.Errorf("reporting error")
-	reportUnit := CommonSourceUnit{"reporting unit"}
+	reportUnit := CommonSourceUnit{ID: "reporting unit"}
 	reportChunk := &Chunk{Data: []byte("reporting chunk")}
 
 	hook.EXPECT().Start(gomock.Any(), startTime)
diff --git a/pkg/sources/source_manager.go b/pkg/sources/source_manager.go
index d3596de79..4f8430fd7 100644
--- a/pkg/sources/source_manager.go
+++ b/pkg/sources/source_manager.go
@@ -360,7 +360,8 @@ func (s *SourceManager) runWithUnits(ctx context.Context, source SourceUnitEnumC
 			report.StartUnitChunking(unit, time.Now())
 			// TODO: Catch panics and add to report.
 			defer close(chunkReporter.chunkCh)
-			ctx := context.WithValue(ctx, "unit", unit.SourceUnitID())
+			id, kind := unit.SourceUnitID()
+			ctx := context.WithValues(ctx, "unit", id, "unit_kind", kind)
 			ctx.Logger().V(3).Info("chunking unit")
 			if err := source.ChunkUnit(ctx, unit, chunkReporter); err != nil {
 				report.ReportError(Fatal{ChunkError{Unit: unit, Err: err}})
diff --git a/pkg/sources/source_manager_test.go b/pkg/sources/source_manager_test.go
index 815f0507e..e7c59dbbd 100644
--- a/pkg/sources/source_manager_test.go
+++ b/pkg/sources/source_manager_test.go
@@ -60,7 +60,14 @@ func (c *counterChunker) Chunks(ctx context.Context, ch chan *Chunk, _ ...Chunki
 // countChunk implements SourceUnit.
 type countChunk byte
 
-func (c countChunk) SourceUnitID() string { return fmt.Sprintf("countChunk(%d)", c) }
+func (c countChunk) SourceUnitID() (string, SourceUnitKind) {
+	return fmt.Sprintf("countChunk(%d)", c), "test"
+}
+
+func (c countChunk) Display() string {
+	id, _ := c.SourceUnitID()
+	return id
+}
 
 func (c *counterChunker) Enumerate(ctx context.Context, reporter UnitReporter) error {
 	for i := 0; i < c.count; i++ {
@@ -187,7 +194,7 @@ func (c *unitChunker) Chunks(ctx context.Context, ch chan *Chunk, _ ...ChunkingT
 }
 func (c *unitChunker) Enumerate(ctx context.Context, rep UnitReporter) error {
 	for _, step := range c.steps {
-		if err := rep.UnitOk(ctx, CommonSourceUnit{step.unit}); err != nil {
+		if err := rep.UnitOk(ctx, CommonSourceUnit{ID: step.unit}); err != nil {
 			return err
 		}
 	}
@@ -195,7 +202,8 @@ func (c *unitChunker) Enumerate(ctx context.Context, rep UnitReporter) error {
 }
 func (c *unitChunker) ChunkUnit(ctx context.Context, unit SourceUnit, rep ChunkReporter) error {
 	for _, step := range c.steps {
-		if unit.SourceUnitID() != step.unit {
+		id, _ := unit.SourceUnitID()
+		if id != step.unit {
 			continue
 		}
 		if step.err != "" {
@@ -340,12 +348,16 @@ func TestSourceManagerUnitHook(t *testing.T) {
 	for metric := range ch {
 		metrics = append(metrics, metric)
 	}
+	justID := func(unit SourceUnit) string {
+		id, _ := unit.SourceUnitID()
+		return id
+	}
 	sort.Slice(metrics, func(i, j int) bool {
-		return metrics[i].Unit.SourceUnitID() < metrics[j].Unit.SourceUnitID()
+		return justID(metrics[i].Unit) < justID(metrics[j].Unit)
 	})
 	m0, m1, m2 := metrics[0], metrics[1], metrics[2]
 
-	assert.Equal(t, "1 one", m0.Unit.SourceUnitID())
+	assert.Equal(t, "1 one", justID(m0.Unit))
 	assert.Equal(t, uint64(1), m0.TotalChunks)
 	assert.Equal(t, uint64(3), m0.TotalBytes)
 	assert.NotZero(t, m0.StartTime)
@@ -353,7 +365,7 @@ func TestSourceManagerUnitHook(t *testing.T) {
 	assert.NotZero(t, m0.ElapsedTime())
 	assert.Equal(t, 0, len(m0.Errors))
 
-	assert.Equal(t, "2 two", m1.Unit.SourceUnitID())
+	assert.Equal(t, "2 two", justID(m1.Unit))
 	assert.Equal(t, uint64(0), m1.TotalChunks)
 	assert.Equal(t, uint64(0), m1.TotalBytes)
 	assert.NotZero(t, m1.StartTime)
@@ -361,7 +373,7 @@ func TestSourceManagerUnitHook(t *testing.T) {
 	assert.NotZero(t, m1.ElapsedTime())
 	assert.Equal(t, 1, len(m1.Errors))
 
-	assert.Equal(t, "3 three", m2.Unit.SourceUnitID())
+	assert.Equal(t, "3 three", justID(m2.Unit))
 	assert.Equal(t, uint64(0), m2.TotalChunks)
 	assert.Equal(t, uint64(0), m2.TotalBytes)
 	assert.NotZero(t, m2.StartTime)
diff --git a/pkg/sources/source_unit.go b/pkg/sources/source_unit.go
index 8eac88430..92a9b015b 100644
--- a/pkg/sources/source_unit.go
+++ b/pkg/sources/source_unit.go
@@ -11,11 +11,20 @@ var _ SourceUnit = CommonSourceUnit{}
 // CommonSourceUnit is a common implementation of SourceUnit that Sources can
 // use instead of implementing their own types.
 type CommonSourceUnit struct {
-	ID string `json:"source_unit_id"`
+	Kind SourceUnitKind `json:"kind,omitempty"`
+	ID   string         `json:"id"`
 }
 
 // SourceUnitID implements the SourceUnit interface.
-func (c CommonSourceUnit) SourceUnitID() string {
+func (c CommonSourceUnit) SourceUnitID() (string, SourceUnitKind) {
+	kind := SourceUnitKind("unit")
+	if c.Kind != "" {
+		kind = c.Kind
+	}
+	return c.ID, kind
+}
+
+func (c CommonSourceUnit) Display() string {
 	return c.ID
 }
 
diff --git a/pkg/sources/sources.go b/pkg/sources/sources.go
index 6d8958871..f87bfd9f4 100644
--- a/pkg/sources/sources.go
+++ b/pkg/sources/sources.go
@@ -130,11 +130,18 @@ type ChunkReporter interface {
 	ChunkErr(ctx context.Context, err error) error
 }
 
+type SourceUnitKind string
+
 // SourceUnit is an object that represents a Source's unit of work. This is
 // used as the output of enumeration, progress reporting, and job distribution.
 type SourceUnit interface {
-	// SourceUnitID uniquely identifies a source unit.
-	SourceUnitID() string
+	// SourceUnitID uniquely identifies a source unit. It does not need to
+	// be human readable or two-way, however, it should be canonical and
+	// stable across runs.
+	SourceUnitID() (string, SourceUnitKind)
+
+	// Display is the human readable representation of the SourceUnit.
+	Display() string
 }
 
 // GCSConfig defines the optional configuration for a GCS source.
diff --git a/pkg/sources/travisci/travisci.go b/pkg/sources/travisci/travisci.go
index cfe457dbc..534365d93 100644
--- a/pkg/sources/travisci/travisci.go
+++ b/pkg/sources/travisci/travisci.go
@@ -111,7 +111,8 @@ func (s *Source) Enumerate(ctx context.Context, reporter sources.UnitReporter) e
 
 		for _, repo := range repositories {
 			err = reporter.UnitOk(ctx, sources.CommonSourceUnit{
-				ID: strconv.Itoa(int(*repo.Id)),
+				ID:   strconv.Itoa(int(*repo.Id)),
+				Kind: "repo",
 			})
 			if err != nil {
 				return fmt.Errorf("error reporting unit: %w", err)
@@ -125,7 +126,8 @@ func (s *Source) Enumerate(ctx context.Context, reporter sources.UnitReporter) e
 
 // ChunkUnit implements SourceUnitChunker interface.
 func (s *Source) ChunkUnit(ctx context.Context, unit sources.SourceUnit, reporter sources.ChunkReporter) error {
-	repo, _, err := s.client.Repositories.Find(ctx, unit.SourceUnitID(), nil)
+	repoURL, _ := unit.SourceUnitID()
+	repo, _, err := s.client.Repositories.Find(ctx, repoURL, nil)
 	if err != nil {
 		return fmt.Errorf("error finding repository: %w", err)
 	}