From f3152b68850f838c64d1ed67f1bc2869e4ce23c3 Mon Sep 17 00:00:00 2001 From: Miccah Date: Fri, 23 Jun 2023 11:15:51 -0500 Subject: [PATCH] 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 | 4 ++- pkg/sources/docker/docker.go | 4 ++- pkg/sources/filesystem/filesystem.go | 4 ++- pkg/sources/gcs/gcs.go | 4 ++- pkg/sources/git/git.go | 7 ++++- pkg/sources/git/unit.go | 40 ++++++++++++++++++++++++++++ pkg/sources/github/github.go | 5 +++- pkg/sources/gitlab/gitlab.go | 4 ++- pkg/sources/s3/s3.go | 4 ++- pkg/sources/source_unit.go | 27 ++++++++++++++++++- 10 files changed, 94 insertions(+), 9 deletions(-) create mode 100644 pkg/sources/git/unit.go diff --git a/pkg/sources/circleci/circleci.go b/pkg/sources/circleci/circleci.go index f6d95e50c9af..ebf7b6cf212d 100644 --- a/pkg/sources/circleci/circleci.go +++ b/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. diff --git a/pkg/sources/docker/docker.go b/pkg/sources/docker/docker.go index 17534da8c8e3..102fd55f1ac7 100644 --- a/pkg/sources/docker/docker.go +++ b/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. diff --git a/pkg/sources/filesystem/filesystem.go b/pkg/sources/filesystem/filesystem.go index 58e756192d7b..eb6dfcf33ea0 100644 --- a/pkg/sources/filesystem/filesystem.go +++ b/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. diff --git a/pkg/sources/gcs/gcs.go b/pkg/sources/gcs/gcs.go index 9f6f0103c4cd..f2a0bedbf883 100644 --- a/pkg/sources/gcs/gcs.go +++ b/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 diff --git a/pkg/sources/git/git.go b/pkg/sources/git/git.go index 14ddc1ea9141..e269c530919c 100644 --- a/pkg/sources/git/git.go +++ b/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) +} diff --git a/pkg/sources/git/unit.go b/pkg/sources/git/unit.go new file mode 100644 index 000000000000..9457e23a6ef5 --- /dev/null +++ b/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 +} diff --git a/pkg/sources/github/github.go b/pkg/sources/github/github.go index 9f413b3c824c..365c0ead346d 100644 --- a/pkg/sources/github/github.go +++ b/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. diff --git a/pkg/sources/gitlab/gitlab.go b/pkg/sources/gitlab/gitlab.go index 7fb96fdeaebb..a32495925fed 100644 --- a/pkg/sources/gitlab/gitlab.go +++ b/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. diff --git a/pkg/sources/s3/s3.go b/pkg/sources/s3/s3.go index 4fef9be10019..5481201ca5a8 100644 --- a/pkg/sources/s3/s3.go +++ b/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 { diff --git a/pkg/sources/source_unit.go b/pkg/sources/source_unit.go index 0c05858a2e6d..8eac88430431 100644 --- a/pkg/sources/source_unit.go +++ b/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 +}