Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement SourceUnitEnumChunker for GitHub #3298

Merged
merged 8 commits into from
Sep 23, 2024
259 changes: 149 additions & 110 deletions pkg/sources/github/github.go
Original file line number Diff line number Diff line change
Expand Up @@ -84,20 +84,20 @@ var _ sources.SourceUnit = (*RepoUnit)(nil)
var _ sources.SourceUnit = (*GistUnit)(nil)

type RepoUnit struct {
name string
url string
Name string `json:"name"`
URL string `json:"url"`
}

func (r RepoUnit) SourceUnitID() (string, sources.SourceUnitKind) { return r.url, "repo" }
func (r RepoUnit) Display() string { return r.name }
func (r RepoUnit) SourceUnitID() (string, sources.SourceUnitKind) { return r.URL, "repo" }
func (r RepoUnit) Display() string { return r.Name }

type GistUnit struct {
name string
url string
Name string `json:"name"`
URL string `json:"url"`
}

func (g GistUnit) SourceUnitID() (string, sources.SourceUnitKind) { return g.url, "gist" }
func (g GistUnit) Display() string { return g.name }
func (g GistUnit) SourceUnitID() (string, sources.SourceUnitKind) { return g.URL, "gist" }
func (g GistUnit) Display() string { return g.Name }

// --------------------------------------------------------------------------------

Expand All @@ -118,6 +118,7 @@ func (s *Source) setScanOptions(base, head string) {
// Ensure the Source satisfies the interfaces at compile time
var _ sources.Source = (*Source)(nil)
var _ sources.SourceUnitUnmarshaller = (*Source)(nil)
var _ sources.SourceUnitEnumChunker = (*Source)(nil)

var endsWithGithub = regexp.MustCompile(`github\.com/?$`)

Expand Down Expand Up @@ -212,6 +213,11 @@ func (s *Source) Init(aCtx context.Context, name string, jobID sources.JobID, so
s.jobPool = &errgroup.Group{}
s.jobPool.SetLimit(concurrency)

// Setup scan options if it wasn't provided.
if s.scanOptions == nil {
s.scanOptions = &git.ScanOptions{}
}

var conn sourcespb.GitHub
err = anypb.UnmarshalTo(connection, &conn, proto.UnmarshalOptions{})
if err != nil {
Expand Down Expand Up @@ -338,19 +344,19 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk, tar
return nil
},
}
err := s.enumerate(ctx, noopReporter)
err := s.Enumerate(ctx, noopReporter)
if err != nil {
return fmt.Errorf("error enumerating: %w", err)
}

return s.scan(ctx, chunksReporter)
}

// enumerate enumerates the GitHub source based on authentication method and
// Enumerate enumerates the GitHub source based on authentication method and
// user configuration. It populates s.filteredRepoCache, s.repoInfoCache,
// s.memberCache, s.totalRepoSize, s.orgsCache, and s.repos. Additionally,
// repositories and gists are reported to the provided UnitReporter.
func (s *Source) enumerate(ctx context.Context, reporter sources.UnitReporter) error {
func (s *Source) Enumerate(ctx context.Context, reporter sources.UnitReporter) error {
seenUnits := make(map[sources.SourceUnit]struct{})
// Wrapper reporter to deduplicate and filter found units.
dedupeReporter := sources.VisitorReporter{
Expand All @@ -370,9 +376,19 @@ func (s *Source) enumerate(ctx context.Context, reporter sources.UnitReporter) e
VisitErr: reporter.UnitErr,
}
// Report any values that were already configured.
// This compensates for differences in enumeration logic between `--org` and `--repo`.
// See: https://github.com/trufflesecurity/trufflehog/pull/2379#discussion_r1487454788
for _, name := range s.filteredRepoCache.Keys() {
url, _ := s.filteredRepoCache.Get(name)
_ = dedupeReporter.UnitOk(ctx, RepoUnit{name: name, url: url})
url, err := s.ensureRepoInfoCache(ctx, url)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The day we get rid of the fourteen different caches in this source I will whistle a jaunty tune

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

14?! Ridiculous! We need to develop one universal cache that covers all the use cases!

if err != nil {
if err := dedupeReporter.UnitErr(ctx, err); err != nil {
return err
}
}
if err := dedupeReporter.UnitOk(ctx, RepoUnit{Name: name, URL: url}); err != nil {
return err
}
}

// I'm not wild about switching on the connector type here (as opposed to dispatching to the connector itself) but
Expand All @@ -395,55 +411,14 @@ func (s *Source) enumerate(ctx context.Context, reporter sources.UnitReporter) e
}
s.repos = make([]string, 0, s.filteredRepoCache.Count())

RepoLoop:
// Double make sure that all enumerated repositories in the
// filteredRepoCache have an entry in the repoInfoCache.
for _, repo := range s.filteredRepoCache.Values() {
repoCtx := context.WithValue(ctx, "repo", repo)

// Ensure that |s.repoInfoCache| contains an entry for |repo|.
// This compensates for differences in enumeration logic between `--org` and `--repo`.
// See: https://github.com/trufflesecurity/trufflehog/pull/2379#discussion_r1487454788
if _, ok := s.repoInfoCache.get(repo); !ok {
repoCtx.Logger().V(2).Info("Caching repository info")
ctx := context.WithValue(ctx, "repo", repo)

_, urlParts, err := getRepoURLParts(repo)
if err != nil {
repoCtx.Logger().Error(err, "Failed to parse repository URL")
continue
}

if isGistUrl(urlParts) {
// Cache gist info.
for {
gistID := extractGistID(urlParts)
gist, _, err := s.connector.APIClient().Gists.Get(repoCtx, gistID)
// Normalize the URL to the Gist's pull URL.
// See https://github.com/trufflesecurity/trufflehog/pull/2625#issuecomment-2025507937
repo = gist.GetGitPullURL()
if s.handleRateLimit(repoCtx, err) {
continue
}
if err != nil {
repoCtx.Logger().Error(err, "Failed to fetch gist")
continue RepoLoop
}
s.cacheGistInfo(gist)
break
}
} else {
// Cache repository info.
for {
ghRepo, _, err := s.connector.APIClient().Repositories.Get(repoCtx, urlParts[1], urlParts[2])
if s.handleRateLimit(repoCtx, err) {
continue
}
if err != nil {
repoCtx.Logger().Error(err, "Failed to fetch repository")
continue RepoLoop
}
s.cacheRepoInfo(ghRepo)
break
}
}
repo, err := s.ensureRepoInfoCache(ctx, repo)
if err != nil {
ctx.Logger().Error(err, "error caching repo info")
}
s.repos = append(s.repos, repo)
}
Expand All @@ -454,6 +429,55 @@ RepoLoop:
return nil
}

// ensureRepoInfoCache checks that s.repoInfoCache has an entry for the
// provided repository URL. If not, it fetches and stores the metadata for the
// repository. In some cases, the gist URL needs to be normalized, which is
// returned by this function.
func (s *Source) ensureRepoInfoCache(ctx context.Context, repo string) (string, error) {
if _, ok := s.repoInfoCache.get(repo); ok {
return repo, nil
}
ctx.Logger().V(2).Info("Caching repository info")

_, urlParts, err := getRepoURLParts(repo)
if err != nil {
return repo, fmt.Errorf("failed to parse repository URL: %w", err)
}

if isGistUrl(urlParts) {
// Cache gist info.
for {
gistID := extractGistID(urlParts)
gist, _, err := s.connector.APIClient().Gists.Get(ctx, gistID)
// Normalize the URL to the Gist's pull URL.
// See https://github.com/trufflesecurity/trufflehog/pull/2625#issuecomment-2025507937
repo = gist.GetGitPullURL()
if s.handleRateLimit(ctx, err) {
continue
}
if err != nil {
return repo, fmt.Errorf("failed to fetch gist")
}
s.cacheGistInfo(gist)
break
}
} else {
// Cache repository info.
for {
ghRepo, _, err := s.connector.APIClient().Repositories.Get(ctx, urlParts[1], urlParts[2])
if s.handleRateLimit(ctx, err) {
continue
}
if err != nil {
return repo, fmt.Errorf("failed to fetch repository")
}
s.cacheRepoInfo(ghRepo)
break
}
}
return repo, nil
}

func (s *Source) enumerateBasicAuth(ctx context.Context, reporter sources.UnitReporter) error {
for _, org := range s.orgsCache.Keys() {
orgCtx := context.WithValue(ctx, "account", org)
Expand Down Expand Up @@ -603,18 +627,13 @@ func (s *Source) scan(ctx context.Context, reporter sources.ChunkReporter) error
reposToScan, progressIndexOffset := sources.FilterReposToResume(s.repos, s.GetProgress().EncodedResumeInfo)
s.repos = reposToScan

scanErrs := sources.NewScanErrors()
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

❤️

// Setup scan options if it wasn't provided.
if s.scanOptions == nil {
s.scanOptions = &git.ScanOptions{}
}

for i, repoURL := range s.repos {
i, repoURL := i, repoURL
s.jobPool.Go(func() error {
if common.IsDone(ctx) {
return nil
}
ctx := context.WithValue(ctx, "repo", repoURL)

// TODO: set progress complete is being called concurrently with i
s.setProgressCompleteWithRepo(i, progressIndexOffset, repoURL)
Expand All @@ -625,64 +644,72 @@ func (s *Source) scan(ctx context.Context, reporter sources.ChunkReporter) error
s.resumeInfoSlice = sources.RemoveRepoFromResumeInfo(s.resumeInfoSlice, repoURL)
}(s, repoURL)

if !strings.HasSuffix(repoURL, ".git") {
scanErrs.Add(fmt.Errorf("repo %s does not end in .git", repoURL))
if err := s.scanRepo(ctx, repoURL, reporter); err != nil {
ctx.Logger().Error(err, "error scanning repo")
return nil
}

// Scan the repository
repoInfo, ok := s.repoInfoCache.get(repoURL)
if !ok {
// This should never happen.
err := fmt.Errorf("no repoInfo for URL: %s", repoURL)
ctx.Logger().Error(err, "failed to scan repository")
return nil
}
repoCtx := context.WithValues(ctx, "repo", repoURL)
duration, err := s.cloneAndScanRepo(repoCtx, repoURL, repoInfo, reporter)
if err != nil {
scanErrs.Add(err)
return nil
}
atomic.AddUint64(&scannedCount, 1)
return nil
})
}

// Scan the wiki, if enabled, and the repo has one.
if s.conn.IncludeWikis && repoInfo.hasWiki && s.wikiIsReachable(ctx, repoURL) {
wikiURL := strings.TrimSuffix(repoURL, ".git") + ".wiki.git"
wikiCtx := context.WithValue(ctx, "repo", wikiURL)
_ = s.jobPool.Wait()
s.SetProgressComplete(len(s.repos), len(s.repos), "Completed GitHub scan", "")

_, err := s.cloneAndScanRepo(wikiCtx, wikiURL, repoInfo, reporter)
if err != nil {
// Ignore "Repository not found" errors.
// It's common for GitHub's API to say a repo has a wiki when it doesn't.
if !strings.Contains(err.Error(), "not found") {
scanErrs.Add(fmt.Errorf("error scanning wiki: %w", err))
}
return nil
}

// Don't return, it still might be possible to scan comments.
}
}
// scanRepo attempts to scan the provided URL and any associated wiki and
// comments if configured. An error is returned if we could not find necessary
// repository metadata or clone the repo, otherwise all errors are reported to
// the ChunkReporter.
func (s *Source) scanRepo(ctx context.Context, repoURL string, reporter sources.ChunkReporter) error {
if !strings.HasSuffix(repoURL, ".git") {
return fmt.Errorf("repo does not end in .git")
}
// Scan the repository
repoInfo, ok := s.repoInfoCache.get(repoURL)
if !ok {
// This should never happen.
return fmt.Errorf("no repoInfo for URL: %s", repoURL)
}
duration, err := s.cloneAndScanRepo(ctx, repoURL, repoInfo, reporter)
if err != nil {
return err
}

// Scan the wiki, if enabled, and the repo has one.
if s.conn.IncludeWikis && repoInfo.hasWiki && s.wikiIsReachable(ctx, repoURL) {
wikiURL := strings.TrimSuffix(repoURL, ".git") + ".wiki.git"
wikiCtx := context.WithValue(ctx, "repo", wikiURL)

// Scan comments, if enabled.
if s.includeGistComments || s.includeIssueComments || s.includePRComments {
if err = s.scanComments(repoCtx, repoURL, repoInfo, reporter); err != nil {
scanErrs.Add(fmt.Errorf("error scanning comments in repo %s: %w", repoURL, err))
return nil
_, err := s.cloneAndScanRepo(wikiCtx, wikiURL, repoInfo, reporter)
if err != nil {
// Ignore "Repository not found" errors.
// It's common for GitHub's API to say a repo has a wiki when it doesn't.
if !strings.Contains(err.Error(), "not found") {
if err := reporter.ChunkErr(ctx, fmt.Errorf("error scanning wiki: %w", err)); err != nil {
return err
}
}

repoCtx.Logger().V(2).Info(fmt.Sprintf("scanned %d/%d repos", scannedCount, len(s.repos)), "duration_seconds", duration)
githubReposScanned.WithLabelValues(s.name).Inc()
atomic.AddUint64(&scannedCount, 1)
return nil
})
// Don't return, it still might be possible to scan comments.
}
}

_ = s.jobPool.Wait()
if scanErrs.Count() > 0 {
ctx.Logger().Info("failed to scan some repositories", "error_count", scanErrs.Count(), "errors", scanErrs.String())
// Scan comments, if enabled.
if s.includeGistComments || s.includeIssueComments || s.includePRComments {
if err := s.scanComments(ctx, repoURL, repoInfo, reporter); err != nil {
err := fmt.Errorf("error scanning comments: %w", err)
if err := reporter.ChunkErr(ctx, err); err != nil {
return err
}
}
}
s.SetProgressComplete(len(s.repos), len(s.repos), "Completed GitHub scan", "")

ctx.Logger().V(2).Info("finished scanning repo", "duration_seconds", duration)
githubReposScanned.WithLabelValues(s.name).Inc()
return nil
}

Expand Down Expand Up @@ -815,7 +842,7 @@ func (s *Source) addUserGistsToCache(ctx context.Context, user string, reporter
for _, gist := range gists {
s.filteredRepoCache.Set(gist.GetID(), gist.GetGitPullURL())
s.cacheGistInfo(gist)
if err := reporter.UnitOk(ctx, GistUnit{name: gist.GetID(), url: gist.GetGitPullURL()}); err != nil {
if err := reporter.UnitOk(ctx, GistUnit{Name: gist.GetID(), URL: gist.GetGitPullURL()}); err != nil {
return err
}
}
Expand Down Expand Up @@ -1478,3 +1505,15 @@ func (s *Source) scanTarget(ctx context.Context, target sources.ChunkingTarget,
Verify: s.verify}
return handlers.HandleFile(ctx, readCloser, &chunkSkel, reporter)
}

func (s *Source) ChunkUnit(ctx context.Context, unit sources.SourceUnit, reporter sources.ChunkReporter) error {
repoURL, _ := unit.SourceUnitID()
ctx = context.WithValue(ctx, "repo", repoURL)
// ChunkUnit is not guaranteed to be called from Enumerate, so we must
// check and fetch the repoInfoCache for this repo.
repoURL, err := s.ensureRepoInfoCache(ctx, repoURL)
if err != nil {
return err
}
return s.scanRepo(ctx, repoURL, reporter)
}
Loading
Loading