trufflesecurity/trufflehog
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add Display method to SourceUnit and Kind member to the CommonSourceUnit (#2450)

Browse Source 
* 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.
This commit is contained in:
Miccah
2024-02-20 11:24:13 -08:00
committed by GitHub
parent 7fea969333
commit c60443891b
15 changed files with 141 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/sources/filesystem/filesystem.go
+1 -1
View File
@@ -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)
pkg/sources/filesystem/filesystem_test.go
+2 -1
View File
@@ -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})
pkg/sources/git/git.go
+5 -8
View File
@@ -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)
}
}
pkg/sources/git/git_test.go
+2 -1
View File
@@ -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})
pkg/sources/git/unit.go
+31 -6
View File
@@ -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.
pkg/sources/git/unit_test.go
+42
View File
@@ -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())
}
pkg/sources/gitlab/gitlab.go
+7 -5
View File
@@ -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
pkg/sources/job_progress.go
+3 -3
View File
@@ -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 }
pkg/sources/job_progress_hook.go
+2 -1
View File
@@ -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)
}
pkg/sources/job_progress_test.go
+1 -1
View File
@@ -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)
pkg/sources/source_manager.go
+2 -1
View File
@@ -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}})
pkg/sources/source_manager_test.go
+19 -7
View File
@@ -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)
pkg/sources/source_unit.go
+11 -2
View File
@@ -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
}
pkg/sources/sources.go
+9 -2
View File
@@ -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.
pkg/sources/travisci/travisci.go
+4 -2
View File
@@ -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)
}