Implement SourceUnitUnmarshaller for all sources (#1416)

* Implement CommonSourceUnitUnmarshaller

* Add SourceUnitUnmarshaller to all sources using

All sources, with the exception of git, will use the CommonSourceUnit as
they only contain a single type of unit to scan.

* Fix method comments to adhere to Go's style guide

pkg/sources/circleci/circleci.go

@@ -31,10 +31,12 @@ type Source struct {
 	jobPool  *errgroup.Group
 	sources.Progress
 	client *http.Client
+	sources.CommonSourceUnitUnmarshaller
 }
 
-// Ensure the Source satisfies the interface at compile time.
+// Ensure the Source satisfies the interfaces at compile time.
 var _ sources.Source = (*Source)(nil)
+var _ sources.SourceUnitUnmarshaller = (*Source)(nil)
 
 // Type returns the type of source.
 // It is used for matching source types in configuration and job input.

pkg/sources/docker/docker.go

@@ -30,12 +30,14 @@ type Source struct {
 	verify   bool
 	conn     sourcespb.Docker
 	sources.Progress
+	sources.CommonSourceUnitUnmarshaller
 }
 
 var FilesizeLimitBytes int64 = 50 * 1024 * 1024 // 50MB
 
-// Ensure the Source satisfies the interface at compile time.
+// Ensure the Source satisfies the interfaces at compile time.
 var _ sources.Source = (*Source)(nil)
+var _ sources.SourceUnitUnmarshaller = (*Source)(nil)
 
 // Type returns the type of source.
 // It is used for matching source types in configuration and job input.

pkg/sources/filesystem/filesystem.go

@@ -39,10 +39,12 @@ type Source struct {
 	log      logr.Logger
 	filter   *common.Filter
 	sources.Progress
+	sources.CommonSourceUnitUnmarshaller
 }
 
-// Ensure the Source satisfies the interface at compile time
+// Ensure the Source satisfies the interfaces at compile time
 var _ sources.Source = (*Source)(nil)
+var _ sources.SourceUnitUnmarshaller = (*Source)(nil)
 
 // Type returns the type of source.
 // It is used for matching source types in configuration and job input.

pkg/sources/gcs/gcs.go

@@ -30,8 +30,9 @@ import (
 
 const defaultCachePersistIncrement = 2500
 
-// Ensure the Source satisfies the interface at compile time.
+// Ensure the Source satisfies the interfaces at compile time.
 var _ sources.Source = (*Source)(nil)
+var _ sources.SourceUnitUnmarshaller = (*Source)(nil)
 
 // Type returns the type of source.
 // It is used for matching source types in configuration and job input.
@@ -69,6 +70,7 @@ type Source struct {
 
 	mu               sync.Mutex
 	sources.Progress // progress is not thread safe
+	sources.CommonSourceUnitUnmarshaller
 }
 
 // persistableCache is a wrapper around cache.Cache that allows

pkg/sources/git/git.go

@@ -67,8 +67,9 @@ func NewGit(sourceType sourcespb.SourceType, jobID, sourceID int64, sourceName s
 	}
 }
 
-// Ensure the Source satisfies the interface at compile time.
+// Ensure the Source satisfies the interfaces at compile time.
 var _ sources.Source = (*Source)(nil)
+var _ sources.SourceUnitUnmarshaller = (*Source)(nil)
 
 // Type returns the type of source.
 // It is used for matching source types in configuration and job input.
@@ -875,3 +876,7 @@ func handleBinary(ctx context.Context, repo *git.Repository, chunksChan chan *so
 
 	return nil
 }
+
+func (s *Source) UnmarshalSourceUnit(data []byte) (sources.SourceUnit, error) {
+	return UnmarshalUnit(data)
+}

pkg/sources/git/unit.go

@@ -0,0 +1,40 @@
+package git
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
+)
+
+const (
+	UnitRepo string = "repo"
+	UnitDir  string = "dir"
+)
+
+// Ensure SourceUnit implements the interface at compile time.
+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"`
+}
+
+// Implement sources.SourceUnit interface.
+func (u SourceUnit) SourceUnitID() string {
+	return u.ID
+}
+
+// Helper function to unmarshal raw bytes into our SourceUnit struct.
+func UnmarshalUnit(data []byte) (sources.SourceUnit, error) {
+	var unit SourceUnit
+	if err := json.Unmarshal(data, &unit); err != nil {
+		return nil, err
+	}
+	if unit.ID == "" || (unit.Kind != UnitRepo && unit.Kind != UnitDir) {
+		return nil, fmt.Errorf("not a git.SourceUnit")
+	}
+	return unit, nil
+}

pkg/sources/github/github.go

@@ -69,14 +69,17 @@ type Source struct {
 	mu                sync.Mutex
 	publicMap         map[string]source_metadatapb.Visibility
 	sources.Progress
+	sources.CommonSourceUnitUnmarshaller
 }
 
 func (s *Source) WithScanOptions(scanOptions *git.ScanOptions) {
 	s.scanOptions = scanOptions
 }
 
-// Ensure the Source satisfies the interface at compile time
+// Ensure the Source satisfies the interfaces at compile time
 var _ sources.Source = (*Source)(nil)
+var _ sources.SourceUnitUnmarshaller = (*Source)(nil)
+
 var endsWithGithub = regexp.MustCompile(`github\.com/?$`)
 
 // Type returns the type of source.

pkg/sources/gitlab/gitlab.go

@@ -47,10 +47,12 @@ type Source struct {
 	resumeInfoMutex sync.Mutex
 	sources.Progress
 	jobPool *errgroup.Group
+	sources.CommonSourceUnitUnmarshaller
 }
 
-// Ensure the Source satisfies the interface at compile time.
+// Ensure the Source satisfies the interfaces at compile time.
 var _ sources.Source = (*Source)(nil)
+var _ sources.SourceUnitUnmarshaller = (*Source)(nil)
 
 // Type returns the type of source.
 // It is used for matching source types in configuration and job input.

pkg/sources/s3/s3.go

@@ -46,10 +46,12 @@ type Source struct {
 	conn          *sourcespb.S3
 	jobPool       *errgroup.Group
 	maxObjectSize int64
+	sources.CommonSourceUnitUnmarshaller
 }
 
-// Ensure the Source satisfies the interface at compile time
+// Ensure the Source satisfies the interfaces at compile time
 var _ sources.Source = (*Source)(nil)
+var _ sources.SourceUnitUnmarshaller = (*Source)(nil)
 
 // Type returns the type of source
 func (s *Source) Type() sourcespb.SourceType {

pkg/sources/source_unit.go

@@ -1,12 +1,37 @@
 package sources
 
+import (
+	"encoding/json"
+	"fmt"
+)
+
+// Ensure CommonSourceUnit implements SourceUnit at compile time.
+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"`
 }
 
-// Implement the SourceUnit interface.
+// SourceUnitID implements the SourceUnit interface.
 func (c CommonSourceUnit) SourceUnitID() string {
 	return c.ID
 }
+
+// CommonSourceUnitUnmarshaller is an implementation of SourceUnitUnmarshaller
+// for the CommonSourceUnit. A source can embed this struct to gain the
+// functionality of converting []byte to a CommonSourceUnit.
+type CommonSourceUnitUnmarshaller struct{}
+
+// UnmarshalSourceUnit implements the SourceUnitUnmarshaller interface.
+func (c CommonSourceUnitUnmarshaller) UnmarshalSourceUnit(data []byte) (SourceUnit, error) {
+	var unit CommonSourceUnit
+	if err := json.Unmarshal(data, &unit); err != nil {
+		return nil, err
+	}
+	if unit.ID == "" {
+		return nil, fmt.Errorf("not a CommonSourceUnit")
+	}
+	return unit, nil
+}