Compare commits
27 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 30008b46a5 | |||
| cf2e56e483 | |||
| a076410b53 | |||
| 98edac9dfb | |||
| 9beee3cbba | |||
| 4a3c8e852c | |||
| 53813884b5 | |||
| caeef5c5ab | |||
| b55013cad3 | |||
| 9dfd307ff4 | |||
| 5c1aa826b1 | |||
| e7313cabf9 | |||
| 948f3ef375 | |||
| 0d755d82eb | |||
| 185a7635f3 | |||
| d0d521123a | |||
| bcab69b8ee | |||
| c84cf53ed6 | |||
| d6d6940cf9 | |||
| e0e29f5b5f | |||
| aa94718ba4 | |||
| 7e84313366 | |||
| f1a9dcff3f | |||
| c5662d76a2 | |||
| 9e8b24ae6b | |||
| 8ed8d6ab04 | |||
| d9f296e7d3 |
@@ -156,7 +156,9 @@ func main() {
|
|||||||
}
|
}
|
||||||
|
|
||||||
for _, repo := range strings.Split(*scanCmdRepo, ",") {
|
for _, repo := range strings.Split(*scanCmdRepo, ",") {
|
||||||
sem.Acquire(ctx, 1)
|
if err := sem.Acquire(ctx, 1); err != nil {
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
wgChunkers.Add(1)
|
wgChunkers.Add(1)
|
||||||
go func(r string) {
|
go func(r string) {
|
||||||
defer sem.Release(1)
|
defer sem.Release(1)
|
||||||
|
|||||||
@@ -79,8 +79,8 @@ func (s *Source) Init(aCtx context.Context, name string, jobId, sourceId int64,
|
|||||||
|
|
||||||
// Chunks emits chunks of bytes over a channel.
|
// Chunks emits chunks of bytes over a channel.
|
||||||
func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) error {
|
func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) error {
|
||||||
for i, path := range s.paths {
|
for _, path := range s.paths {
|
||||||
s.SetProgressComplete(i, len(s.paths), fmt.Sprintf("Path: %s", path), "")
|
s.Update(len(s.paths), fmt.Sprintf("Path: %s", path), "")
|
||||||
|
|
||||||
cleanPath := filepath.Clean(path)
|
cleanPath := filepath.Clean(path)
|
||||||
done := false
|
done := false
|
||||||
|
|||||||
+13
-9
@@ -132,8 +132,8 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
|
|||||||
user := cred.BasicAuth.Username
|
user := cred.BasicAuth.Username
|
||||||
token := cred.BasicAuth.Password
|
token := cred.BasicAuth.Password
|
||||||
|
|
||||||
for i, repoURI := range s.conn.Repositories {
|
for _, repoURI := range s.conn.Repositories {
|
||||||
s.SetProgressComplete(i, len(s.conn.Repositories), fmt.Sprintf("Repo: %s", repoURI), "")
|
s.Update(len(s.conn.Repositories), fmt.Sprintf("Repo: %s", repoURI), "")
|
||||||
if len(repoURI) == 0 {
|
if len(repoURI) == 0 {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
@@ -150,8 +150,8 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
case *sourcespb.Git_Unauthenticated:
|
case *sourcespb.Git_Unauthenticated:
|
||||||
for i, repoURI := range s.conn.Repositories {
|
for _, repoURI := range s.conn.Repositories {
|
||||||
s.SetProgressComplete(i, len(s.conn.Repositories), fmt.Sprintf("Repo: %s", repoURI), "")
|
s.Update(len(s.conn.Repositories), fmt.Sprintf("Repo: %s", repoURI), "")
|
||||||
if len(repoURI) == 0 {
|
if len(repoURI) == 0 {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
@@ -168,8 +168,8 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
case *sourcespb.Git_SshAuth:
|
case *sourcespb.Git_SshAuth:
|
||||||
for i, repoURI := range s.conn.Repositories {
|
for _, repoURI := range s.conn.Repositories {
|
||||||
s.SetProgressComplete(i, len(s.conn.Repositories), fmt.Sprintf("Repo: %s", repoURI), "")
|
s.Update(len(s.conn.Repositories), fmt.Sprintf("Repo: %s", repoURI), "")
|
||||||
if len(repoURI) == 0 {
|
if len(repoURI) == 0 {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
@@ -189,8 +189,12 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
|
|||||||
return errors.New("invalid connection type for git source")
|
return errors.New("invalid connection type for git source")
|
||||||
}
|
}
|
||||||
|
|
||||||
for i, u := range s.conn.Directories {
|
for _, u := range s.conn.Directories {
|
||||||
s.SetProgressComplete(i, len(s.conn.Repositories), fmt.Sprintf("Repo: %s", u), "")
|
s.Counter.IncTotal()
|
||||||
|
go func() {
|
||||||
|
defer s.Counter.IncSuccess()
|
||||||
|
}()
|
||||||
|
s.Update(len(s.conn.Repositories), fmt.Sprintf("Repo: %s", u), "")
|
||||||
|
|
||||||
if len(u) == 0 {
|
if len(u) == 0 {
|
||||||
continue
|
continue
|
||||||
@@ -215,7 +219,7 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
|
|||||||
}
|
}
|
||||||
|
|
||||||
}
|
}
|
||||||
s.SetProgressComplete(len(s.conn.Repositories), len(s.conn.Repositories), fmt.Sprintf("Completed scanning source %s", s.name), "")
|
s.Complete(fmt.Sprintf("Completed scanning source %s", s.name))
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -496,15 +496,17 @@ func (s *Source) scan(ctx context.Context, installationClient *github.Client, ch
|
|||||||
s.repos = reposToScan
|
s.repos = reposToScan
|
||||||
|
|
||||||
var scanErrs []error
|
var scanErrs []error
|
||||||
for i, repoURL := range s.repos {
|
repoCnt := len(s.repos)
|
||||||
i, repoURL := i, repoURL
|
for _, repoURL := range s.repos {
|
||||||
|
repoURL := repoURL
|
||||||
s.jobPool.Go(func() error {
|
s.jobPool.Go(func() error {
|
||||||
if common.IsDone(ctx) {
|
if common.IsDone(ctx) {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
cnt := s.Counter.IncTotal()
|
||||||
|
|
||||||
// TODO: set progress complete is being called concurrently with i
|
defer s.Counter.IncSuccess()
|
||||||
s.setProgressCompleteWithRepo(i, progressIndexOffset, repoURL)
|
s.setProgressCompleteWithRepo(progressIndexOffset, repoURL)
|
||||||
// Ensure the repo is removed from the resume info after being scanned.
|
// Ensure the repo is removed from the resume info after being scanned.
|
||||||
defer func(s *Source, repoURL string) {
|
defer func(s *Source, repoURL string) {
|
||||||
s.resumeInfoMutex.Lock()
|
s.resumeInfoMutex.Lock()
|
||||||
@@ -512,15 +514,15 @@ func (s *Source) scan(ctx context.Context, installationClient *github.Client, ch
|
|||||||
s.resumeInfoSlice = sources.RemoveRepoFromResumeInfo(s.resumeInfoSlice, repoURL)
|
s.resumeInfoSlice = sources.RemoveRepoFromResumeInfo(s.resumeInfoSlice, repoURL)
|
||||||
}(s, repoURL)
|
}(s, repoURL)
|
||||||
|
|
||||||
|
var err error
|
||||||
if !strings.HasSuffix(repoURL, ".git") {
|
if !strings.HasSuffix(repoURL, ".git") {
|
||||||
scanErrs = append(scanErrs, fmt.Errorf("repo %s does not end in .git", repoURL))
|
scanErrs = append(scanErrs, fmt.Errorf("repo %s does not end in .git", repoURL))
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
s.log.WithField("repo", repoURL).Debugf("attempting to clone repo %d/%d", i+1, len(s.repos))
|
s.log.WithField("repo", repoURL).Debugf("attempting to clone repo #%d of %d", cnt, repoCnt)
|
||||||
var path string
|
var path string
|
||||||
var repo *gogit.Repository
|
var repo *gogit.Repository
|
||||||
var err error
|
|
||||||
|
|
||||||
path, repo, err = s.cloneRepo(ctx, repoURL, installationClient)
|
path, repo, err = s.cloneRepo(ctx, repoURL, installationClient)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -537,21 +539,20 @@ func (s *Source) scan(ctx context.Context, installationClient *github.Client, ch
|
|||||||
git.ScanOptionHeadCommit(s.conn.Head),
|
git.ScanOptionHeadCommit(s.conn.Head),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
log.Debugf("Starting to scan repo #%d of %d: %s", cnt, repoCnt, repo)
|
||||||
if err = s.git.ScanRepo(ctx, repo, path, scanOptions, chunksChan); err != nil {
|
if err = s.git.ScanRepo(ctx, repo, path, scanOptions, chunksChan); err != nil {
|
||||||
log.WithError(err).Errorf("unable to scan repo, continuing")
|
log.WithError(err).Errorf("unable to scan repo, continuing")
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
atomic.AddUint64(&scanned, 1)
|
atomic.AddUint64(&scanned, 1)
|
||||||
log.Debugf("scanned %d/%d repos", scanned, len(s.repos))
|
log.Debugf("scanned %d/%d repos", scanned, repoCnt)
|
||||||
|
|
||||||
return nil
|
return nil
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
_ = s.jobPool.Wait()
|
_ = s.jobPool.Wait()
|
||||||
if len(scanErrs) == 0 {
|
s.Complete(fmt.Sprintf("Completed scanning source %s", s.name))
|
||||||
s.SetProgressComplete(len(s.repos), len(s.repos), "Completed Github scan", "")
|
|
||||||
}
|
|
||||||
|
|
||||||
return scanErrs
|
return scanErrs
|
||||||
}
|
}
|
||||||
@@ -1006,8 +1007,8 @@ func (s *Source) normalizeRepos(ctx context.Context) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// setProgressCompleteWithRepo calls the s.SetProgressComplete after safely setting up the encoded resume info string.
|
// setProgressCompleteWithRepo calls the s.Update after safely setting up the encoded resume info string.
|
||||||
func (s *Source) setProgressCompleteWithRepo(index int, offset int, repoURL string) {
|
func (s *Source) setProgressCompleteWithRepo(offset int, repoURL string) {
|
||||||
s.resumeInfoMutex.Lock()
|
s.resumeInfoMutex.Lock()
|
||||||
defer s.resumeInfoMutex.Unlock()
|
defer s.resumeInfoMutex.Unlock()
|
||||||
|
|
||||||
@@ -1018,5 +1019,5 @@ func (s *Source) setProgressCompleteWithRepo(index int, offset int, repoURL stri
|
|||||||
// Make the resume info string from the slice.
|
// Make the resume info string from the slice.
|
||||||
encodedResumeInfo := sources.EncodeResumeInfo(s.resumeInfoSlice)
|
encodedResumeInfo := sources.EncodeResumeInfo(s.resumeInfoSlice)
|
||||||
|
|
||||||
s.SetProgressComplete(index+offset, len(s.repos)+offset, fmt.Sprintf("Repo: %s", repoURL), encodedResumeInfo)
|
s.Update(len(s.repos)+offset, fmt.Sprintf("Repo: %s", repoURL), encodedResumeInfo)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -383,7 +383,7 @@ func Test_setProgressCompleteWithRepo_resumeInfo(t *testing.T) {
|
|||||||
|
|
||||||
for _, tt := range tests {
|
for _, tt := range tests {
|
||||||
s.resumeInfoSlice = tt.startingResumeInfoSlice
|
s.resumeInfoSlice = tt.startingResumeInfoSlice
|
||||||
s.setProgressCompleteWithRepo(0, 0, tt.repoURL)
|
s.setProgressCompleteWithRepo(0, tt.repoURL)
|
||||||
if !reflect.DeepEqual(s.resumeInfoSlice, tt.wantResumeInfoSlice) {
|
if !reflect.DeepEqual(s.resumeInfoSlice, tt.wantResumeInfoSlice) {
|
||||||
t.Errorf("s.setProgressCompleteWithRepo() got: %v, want: %v", s.resumeInfoSlice, tt.wantResumeInfoSlice)
|
t.Errorf("s.setProgressCompleteWithRepo() got: %v, want: %v", s.resumeInfoSlice, tt.wantResumeInfoSlice)
|
||||||
}
|
}
|
||||||
@@ -394,7 +394,6 @@ func Test_setProgressCompleteWithRepo_Progress(t *testing.T) {
|
|||||||
repos := []string{"a", "b", "c", "d", "e"}
|
repos := []string{"a", "b", "c", "d", "e"}
|
||||||
tests := map[string]struct {
|
tests := map[string]struct {
|
||||||
repos []string
|
repos []string
|
||||||
index int
|
|
||||||
offset int
|
offset int
|
||||||
wantPercentComplete int64
|
wantPercentComplete int64
|
||||||
wantSectionsCompleted int32
|
wantSectionsCompleted int32
|
||||||
@@ -402,7 +401,6 @@ func Test_setProgressCompleteWithRepo_Progress(t *testing.T) {
|
|||||||
}{
|
}{
|
||||||
"starting from the beginning, no offset": {
|
"starting from the beginning, no offset": {
|
||||||
repos: repos,
|
repos: repos,
|
||||||
index: 0,
|
|
||||||
offset: 0,
|
offset: 0,
|
||||||
wantPercentComplete: 0,
|
wantPercentComplete: 0,
|
||||||
wantSectionsCompleted: 0,
|
wantSectionsCompleted: 0,
|
||||||
@@ -410,18 +408,16 @@ func Test_setProgressCompleteWithRepo_Progress(t *testing.T) {
|
|||||||
},
|
},
|
||||||
"resume from the third, offset 2": {
|
"resume from the third, offset 2": {
|
||||||
repos: repos[2:],
|
repos: repos[2:],
|
||||||
index: 0,
|
|
||||||
offset: 2,
|
offset: 2,
|
||||||
wantPercentComplete: 40,
|
wantPercentComplete: 40,
|
||||||
wantSectionsCompleted: 2,
|
wantSectionsCompleted: 2,
|
||||||
wantSectionsRemaining: 5,
|
wantSectionsRemaining: 5,
|
||||||
},
|
},
|
||||||
"resume from the third, on last repo, offset 2": {
|
"resume from the third, on last repo, offset 3": {
|
||||||
repos: repos[2:],
|
repos: repos[3:],
|
||||||
index: 2,
|
offset: 3,
|
||||||
offset: 2,
|
wantPercentComplete: 60,
|
||||||
wantPercentComplete: 80,
|
wantSectionsCompleted: 3,
|
||||||
wantSectionsCompleted: 4,
|
|
||||||
wantSectionsRemaining: 5,
|
wantSectionsRemaining: 5,
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
@@ -430,22 +426,29 @@ func Test_setProgressCompleteWithRepo_Progress(t *testing.T) {
|
|||||||
logger.Out = io.Discard
|
logger.Out = io.Discard
|
||||||
|
|
||||||
for _, tt := range tests {
|
for _, tt := range tests {
|
||||||
s := &Source{
|
t.Run("", func(t *testing.T) {
|
||||||
repos: tt.repos,
|
s := &Source{
|
||||||
log: logger.WithField("no", "output"),
|
repos: tt.repos,
|
||||||
}
|
log: logger.WithField("no", "output"),
|
||||||
|
}
|
||||||
|
|
||||||
s.setProgressCompleteWithRepo(tt.index, tt.offset, "")
|
// Increment the progress counter by number of repos already processed.
|
||||||
gotProgress := s.GetProgress()
|
for i := 1; i <= tt.offset; i++ {
|
||||||
if gotProgress.PercentComplete != tt.wantPercentComplete {
|
s.Counter.IncSuccess()
|
||||||
t.Errorf("s.setProgressCompleteWithRepo() PercentComplete got: %v want: %v", gotProgress.PercentComplete, tt.wantPercentComplete)
|
}
|
||||||
}
|
|
||||||
if gotProgress.SectionsCompleted != tt.wantSectionsCompleted {
|
s.setProgressCompleteWithRepo(tt.offset, "")
|
||||||
t.Errorf("s.setProgressCompleteWithRepo() PercentComplete got: %v want: %v", gotProgress.SectionsCompleted, tt.wantSectionsCompleted)
|
gotProgress := s.GetProgress()
|
||||||
}
|
if gotProgress.PercentComplete != tt.wantPercentComplete {
|
||||||
if gotProgress.SectionsRemaining != tt.wantSectionsRemaining {
|
t.Errorf("s.setProgressCompleteWithRepo() PercentComplete got: %v want: %v", gotProgress.PercentComplete, tt.wantPercentComplete)
|
||||||
t.Errorf("s.setProgressCompleteWithRepo() PercentComplete got: %v want: %v", gotProgress.SectionsRemaining, tt.wantSectionsRemaining)
|
}
|
||||||
}
|
if gotProgress.SectionsCompleted != tt.wantSectionsCompleted {
|
||||||
|
t.Errorf("s.setProgressCompleteWithRepo() SectionsCompleted got: %v want: %v", gotProgress.SectionsCompleted, tt.wantSectionsCompleted)
|
||||||
|
}
|
||||||
|
if gotProgress.SectionsRemaining != tt.wantSectionsRemaining {
|
||||||
|
t.Errorf("s.setProgressCompleteWithRepo() SectionsRemaining got: %v want: %v", gotProgress.SectionsRemaining, tt.wantSectionsRemaining)
|
||||||
|
}
|
||||||
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -454,16 +457,15 @@ func Test_scan_SetProgressComplete(t *testing.T) {
|
|||||||
name string
|
name string
|
||||||
repos []string
|
repos []string
|
||||||
wantComplete bool
|
wantComplete bool
|
||||||
wantErr bool
|
|
||||||
}{
|
}{
|
||||||
{
|
{
|
||||||
name: "no repos",
|
name: "no repos",
|
||||||
wantComplete: true,
|
wantComplete: true,
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
name: "one valid repo",
|
name: "one valid repo",
|
||||||
repos: []string{"a"},
|
repos: []string{"a"},
|
||||||
wantErr: true,
|
wantComplete: true,
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -475,9 +477,6 @@ func Test_scan_SetProgressComplete(t *testing.T) {
|
|||||||
src.jobPool = &errgroup.Group{}
|
src.jobPool = &errgroup.Group{}
|
||||||
|
|
||||||
_ = src.scan(context.Background(), nil, nil)
|
_ = src.scan(context.Background(), nil, nil)
|
||||||
if !tc.wantErr {
|
|
||||||
assert.Equal(t, "", src.GetProgress().EncodedResumeInfo)
|
|
||||||
}
|
|
||||||
|
|
||||||
gotComplete := src.GetProgress().PercentComplete == 100
|
gotComplete := src.GetProgress().PercentComplete == 100
|
||||||
if gotComplete != tc.wantComplete {
|
if gotComplete != tc.wantComplete {
|
||||||
|
|||||||
@@ -8,6 +8,9 @@ import (
|
|||||||
"sort"
|
"sort"
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
|
|
||||||
|
"golang.org/x/sync/errgroup"
|
||||||
|
|
||||||
"github.com/trufflesecurity/trufflehog/v3/pkg/common"
|
"github.com/trufflesecurity/trufflehog/v3/pkg/common"
|
||||||
"github.com/trufflesecurity/trufflehog/v3/pkg/context"
|
"github.com/trufflesecurity/trufflehog/v3/pkg/context"
|
||||||
@@ -46,6 +49,7 @@ type Source struct {
|
|||||||
resumeInfoMutex sync.Mutex
|
resumeInfoMutex sync.Mutex
|
||||||
jobSem *semaphore.Weighted
|
jobSem *semaphore.Weighted
|
||||||
sources.Progress
|
sources.Progress
|
||||||
|
jobPool *errgroup.Group
|
||||||
}
|
}
|
||||||
|
|
||||||
// Ensure the Source satisfies the interface at compile time.
|
// Ensure the Source satisfies the interface at compile time.
|
||||||
@@ -73,7 +77,8 @@ func (s *Source) Init(aCtx context.Context, name string, jobId, sourceId int64,
|
|||||||
s.sourceId = sourceId
|
s.sourceId = sourceId
|
||||||
s.jobId = jobId
|
s.jobId = jobId
|
||||||
s.verify = verify
|
s.verify = verify
|
||||||
s.jobSem = semaphore.NewWeighted(int64(concurrency))
|
s.jobPool = &errgroup.Group{}
|
||||||
|
s.jobPool.SetLimit(concurrency)
|
||||||
|
|
||||||
var conn sourcespb.GitLab
|
var conn sourcespb.GitLab
|
||||||
err := anypb.UnmarshalTo(connection, &conn, proto.UnmarshalOptions{})
|
err := anypb.UnmarshalTo(connection, &conn, proto.UnmarshalOptions{})
|
||||||
@@ -269,34 +274,32 @@ func (s *Source) getRepos() ([]string, []error) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (s *Source) scanRepos(ctx context.Context, chunksChan chan *sources.Chunk) []error {
|
func (s *Source) scanRepos(ctx context.Context, chunksChan chan *sources.Chunk) []error {
|
||||||
wg := sync.WaitGroup{}
|
var scanned uint64
|
||||||
var errs []error
|
|
||||||
var errsMut sync.Mutex
|
|
||||||
|
|
||||||
// If there is resume information available, limit this scan to only the repos that still need scanning.
|
// If there is resume information available, limit this scan to only the repos that still need scanning.
|
||||||
reposToScan, progressIndexOffset := sources.FilterReposToResume(s.repos, s.GetProgress().EncodedResumeInfo)
|
reposToScan, progressIndexOffset := sources.FilterReposToResume(s.repos, s.GetProgress().EncodedResumeInfo)
|
||||||
s.repos = reposToScan
|
s.repos = reposToScan
|
||||||
|
log.Debugf("Found %v total repos to scan", len(s.repos))
|
||||||
|
|
||||||
for i, repo := range s.repos {
|
var errMut sync.Mutex
|
||||||
if common.IsDone(ctx) {
|
var scanErrs []error
|
||||||
// We are returning nil instead of the scanErrors slice here because
|
repoCnt := len(s.repos)
|
||||||
// we don't want to mark this scan as errored if we cancelled it.
|
for _, repoURL := range s.repos {
|
||||||
return nil
|
repoURL := repoURL
|
||||||
}
|
s.jobPool.Go(func() error {
|
||||||
if err := s.jobSem.Acquire(ctx, 1); err != nil {
|
if common.IsDone(ctx) {
|
||||||
log.WithError(err).Debug("could not acquire semaphore")
|
// We are returning nil instead of the scanErrors slice here because
|
||||||
continue
|
// we don't want to mark this scan as errored if we cancelled it.
|
||||||
}
|
return nil
|
||||||
wg.Add(1)
|
}
|
||||||
go func(ctx context.Context, repoURL string, i int) {
|
|
||||||
defer s.jobSem.Release(1)
|
|
||||||
defer wg.Done()
|
|
||||||
if len(repoURL) == 0 {
|
if len(repoURL) == 0 {
|
||||||
return
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
s.setProgressCompleteWithRepo(i, progressIndexOffset, repoURL)
|
cnt := s.Counter.IncTotal()
|
||||||
// Ensure the repo is removed from the resume info after being scanned.
|
defer s.Counter.IncSuccess()
|
||||||
|
s.setProgressCompleteWithRepo(progressIndexOffset, repoURL)
|
||||||
|
// Ensure the repoURL is removed from the resume info after being scanned.
|
||||||
defer func(s *Source) {
|
defer func(s *Source) {
|
||||||
s.resumeInfoMutex.Lock()
|
s.resumeInfoMutex.Lock()
|
||||||
defer s.resumeInfoMutex.Unlock()
|
defer s.resumeInfoMutex.Unlock()
|
||||||
@@ -309,7 +312,7 @@ func (s *Source) scanRepos(ctx context.Context, chunksChan chan *sources.Chunk)
|
|||||||
if s.authMethod == "UNAUTHENTICATED" {
|
if s.authMethod == "UNAUTHENTICATED" {
|
||||||
path, repo, err = git.CloneRepoUsingUnauthenticated(repoURL)
|
path, repo, err = git.CloneRepoUsingUnauthenticated(repoURL)
|
||||||
} else {
|
} else {
|
||||||
// If a username is not provided we need to use a default one in order to clone a private repo.
|
// If a username is not provided we need to use a default one in order to clone a private repoURL.
|
||||||
// Not setting "placeholder" as s.user on purpose in case any downstream services rely on a "" value for s.user.
|
// Not setting "placeholder" as s.user on purpose in case any downstream services rely on a "" value for s.user.
|
||||||
user := s.user
|
user := s.user
|
||||||
if user == "" {
|
if user == "" {
|
||||||
@@ -319,28 +322,32 @@ func (s *Source) scanRepos(ctx context.Context, chunksChan chan *sources.Chunk)
|
|||||||
}
|
}
|
||||||
defer os.RemoveAll(path)
|
defer os.RemoveAll(path)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
errsMut.Lock()
|
errMut.Lock()
|
||||||
errs = append(errs, err)
|
scanErrs = append(scanErrs, err)
|
||||||
errsMut.Unlock()
|
errMut.Unlock()
|
||||||
return
|
return nil
|
||||||
}
|
}
|
||||||
log.Debugf("Starting to scan repo %d/%d: %s", i+1, len(s.repos), repoURL)
|
|
||||||
|
log.Debugf("Starting to scan repoURL #%d of %d: %s", cnt, repoCnt, repo)
|
||||||
err = s.git.ScanRepo(ctx, repo, path, git.NewScanOptions(), chunksChan)
|
err = s.git.ScanRepo(ctx, repo, path, git.NewScanOptions(), chunksChan)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
errsMut.Lock()
|
errMut.Lock()
|
||||||
errs = append(errs, err)
|
scanErrs = append(scanErrs, err)
|
||||||
errsMut.Unlock()
|
errMut.Unlock()
|
||||||
return
|
return nil
|
||||||
}
|
}
|
||||||
log.Debugf("Completed scanning repo %d/%d: %s", i+1, len(s.repos), repoURL)
|
atomic.AddUint64(&scanned, 1)
|
||||||
}(ctx, repo, i)
|
log.Debugf("scanned %d/%d repos", scanned, repoCnt)
|
||||||
}
|
log.Debugf("Completed scanning repoURL # %d of %d: %s", cnt, repoCnt, repo)
|
||||||
wg.Wait()
|
|
||||||
if len(errs) == 0 {
|
return nil
|
||||||
s.SetProgressComplete(len(s.repos), len(s.repos), fmt.Sprintf("Completed scanning source %s", s.name), "")
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
return errs
|
_ = s.jobPool.Wait()
|
||||||
|
s.Complete(fmt.Sprintf("Completed scanning source %s", s.name))
|
||||||
|
|
||||||
|
return scanErrs
|
||||||
}
|
}
|
||||||
|
|
||||||
// Chunks emits chunks of bytes over a channel.
|
// Chunks emits chunks of bytes over a channel.
|
||||||
@@ -417,8 +424,8 @@ func (s *Source) basicAuthSuccessful(apiClient *gitlab.Client) bool {
|
|||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
||||||
// setProgressCompleteWithRepo calls the s.SetProgressComplete after safely setting up the encoded resume info string.
|
// setProgressCompleteWithRepo calls the s.Update after safely setting up the encoded resume info string.
|
||||||
func (s *Source) setProgressCompleteWithRepo(index int, offset int, repoURL string) {
|
func (s *Source) setProgressCompleteWithRepo(offset int, repoURL string) {
|
||||||
s.resumeInfoMutex.Lock()
|
s.resumeInfoMutex.Lock()
|
||||||
defer s.resumeInfoMutex.Unlock()
|
defer s.resumeInfoMutex.Unlock()
|
||||||
|
|
||||||
@@ -430,5 +437,5 @@ func (s *Source) setProgressCompleteWithRepo(index int, offset int, repoURL stri
|
|||||||
encodedResumeInfo := sources.EncodeResumeInfo(s.resumeInfoSlice)
|
encodedResumeInfo := sources.EncodeResumeInfo(s.resumeInfoSlice)
|
||||||
|
|
||||||
// Add the offset to both the index and the repos to give the proper place and proper repo count.
|
// Add the offset to both the index and the repos to give the proper place and proper repo count.
|
||||||
s.SetProgressComplete(index+offset, len(s.repos)+offset, fmt.Sprintf("Repo: %s", repoURL), encodedResumeInfo)
|
s.Update(len(s.repos)+offset, fmt.Sprintf("Repo: %s", repoURL), encodedResumeInfo)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -9,7 +9,7 @@ import (
|
|||||||
"github.com/kylelemons/godebug/pretty"
|
"github.com/kylelemons/godebug/pretty"
|
||||||
log "github.com/sirupsen/logrus"
|
log "github.com/sirupsen/logrus"
|
||||||
"github.com/stretchr/testify/assert"
|
"github.com/stretchr/testify/assert"
|
||||||
"golang.org/x/sync/semaphore"
|
"golang.org/x/sync/errgroup"
|
||||||
"google.golang.org/protobuf/types/known/anypb"
|
"google.golang.org/protobuf/types/known/anypb"
|
||||||
|
|
||||||
"github.com/trufflesecurity/trufflehog/v3/pkg/common"
|
"github.com/trufflesecurity/trufflehog/v3/pkg/common"
|
||||||
@@ -187,7 +187,7 @@ func Test_setProgressCompleteWithRepo_resumeInfo(t *testing.T) {
|
|||||||
|
|
||||||
for _, tt := range tests {
|
for _, tt := range tests {
|
||||||
s.resumeInfoSlice = tt.startingResumeInfoSlice
|
s.resumeInfoSlice = tt.startingResumeInfoSlice
|
||||||
s.setProgressCompleteWithRepo(0, 0, tt.repoURL)
|
s.setProgressCompleteWithRepo(0, tt.repoURL)
|
||||||
if !reflect.DeepEqual(s.resumeInfoSlice, tt.wantResumeInfoSlice) {
|
if !reflect.DeepEqual(s.resumeInfoSlice, tt.wantResumeInfoSlice) {
|
||||||
t.Errorf("s.setProgressCompleteWithRepo() got: %v, want: %v", s.resumeInfoSlice, tt.wantResumeInfoSlice)
|
t.Errorf("s.setProgressCompleteWithRepo() got: %v, want: %v", s.resumeInfoSlice, tt.wantResumeInfoSlice)
|
||||||
}
|
}
|
||||||
@@ -198,7 +198,6 @@ func Test_setProgressCompleteWithRepo_Progress(t *testing.T) {
|
|||||||
repos := []string{"a", "b", "c", "d", "e"}
|
repos := []string{"a", "b", "c", "d", "e"}
|
||||||
tests := map[string]struct {
|
tests := map[string]struct {
|
||||||
repos []string
|
repos []string
|
||||||
index int
|
|
||||||
offset int
|
offset int
|
||||||
wantPercentComplete int64
|
wantPercentComplete int64
|
||||||
wantSectionsCompleted int32
|
wantSectionsCompleted int32
|
||||||
@@ -206,7 +205,6 @@ func Test_setProgressCompleteWithRepo_Progress(t *testing.T) {
|
|||||||
}{
|
}{
|
||||||
"starting from the beginning, no offset": {
|
"starting from the beginning, no offset": {
|
||||||
repos: repos,
|
repos: repos,
|
||||||
index: 0,
|
|
||||||
offset: 0,
|
offset: 0,
|
||||||
wantPercentComplete: 0,
|
wantPercentComplete: 0,
|
||||||
wantSectionsCompleted: 0,
|
wantSectionsCompleted: 0,
|
||||||
@@ -214,18 +212,16 @@ func Test_setProgressCompleteWithRepo_Progress(t *testing.T) {
|
|||||||
},
|
},
|
||||||
"resume from the third, offset 2": {
|
"resume from the third, offset 2": {
|
||||||
repos: repos[2:],
|
repos: repos[2:],
|
||||||
index: 0,
|
|
||||||
offset: 2,
|
offset: 2,
|
||||||
wantPercentComplete: 40,
|
wantPercentComplete: 40,
|
||||||
wantSectionsCompleted: 2,
|
wantSectionsCompleted: 2,
|
||||||
wantSectionsRemaining: 5,
|
wantSectionsRemaining: 5,
|
||||||
},
|
},
|
||||||
"resume from the third, on last repo, offset 2": {
|
"resume from the third, on last repo, offset 2": {
|
||||||
repos: repos[2:],
|
repos: repos[3:],
|
||||||
index: 2,
|
offset: 3,
|
||||||
offset: 2,
|
wantPercentComplete: 60,
|
||||||
wantPercentComplete: 80,
|
wantSectionsCompleted: 3,
|
||||||
wantSectionsCompleted: 4,
|
|
||||||
wantSectionsRemaining: 5,
|
wantSectionsRemaining: 5,
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
@@ -233,21 +229,28 @@ func Test_setProgressCompleteWithRepo_Progress(t *testing.T) {
|
|||||||
log.SetOutput(io.Discard)
|
log.SetOutput(io.Discard)
|
||||||
|
|
||||||
for _, tt := range tests {
|
for _, tt := range tests {
|
||||||
s := &Source{
|
t.Run("", func(t *testing.T) {
|
||||||
repos: tt.repos,
|
s := &Source{
|
||||||
}
|
repos: tt.repos,
|
||||||
|
}
|
||||||
|
|
||||||
s.setProgressCompleteWithRepo(tt.index, tt.offset, "")
|
// Increment the progress counter by number of repos already processed.
|
||||||
gotProgress := s.GetProgress()
|
for i := 1; i <= tt.offset; i++ {
|
||||||
if gotProgress.PercentComplete != tt.wantPercentComplete {
|
s.Counter.IncSuccess()
|
||||||
t.Errorf("s.setProgressCompleteWithRepo() PercentComplete got: %v want: %v", gotProgress.PercentComplete, tt.wantPercentComplete)
|
}
|
||||||
}
|
|
||||||
if gotProgress.SectionsCompleted != tt.wantSectionsCompleted {
|
s.setProgressCompleteWithRepo(tt.offset, "")
|
||||||
t.Errorf("s.setProgressCompleteWithRepo() PercentComplete got: %v want: %v", gotProgress.SectionsCompleted, tt.wantSectionsCompleted)
|
gotProgress := s.GetProgress()
|
||||||
}
|
if gotProgress.PercentComplete != tt.wantPercentComplete {
|
||||||
if gotProgress.SectionsRemaining != tt.wantSectionsRemaining {
|
t.Errorf("s.setProgressCompleteWithRepo() PercentComplete got: %v want: %v", gotProgress.PercentComplete, tt.wantPercentComplete)
|
||||||
t.Errorf("s.setProgressCompleteWithRepo() PercentComplete got: %v want: %v", gotProgress.SectionsRemaining, tt.wantSectionsRemaining)
|
}
|
||||||
}
|
if gotProgress.SectionsCompleted != tt.wantSectionsCompleted {
|
||||||
|
t.Errorf("s.setProgressCompleteWithRepo() SectionsCompleted got: %v want: %v", gotProgress.SectionsCompleted, tt.wantSectionsCompleted)
|
||||||
|
}
|
||||||
|
if gotProgress.SectionsRemaining != tt.wantSectionsRemaining {
|
||||||
|
t.Errorf("s.setProgressCompleteWithRepo() SectionsRemaining got: %v want: %v", gotProgress.SectionsRemaining, tt.wantSectionsRemaining)
|
||||||
|
}
|
||||||
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -256,16 +259,15 @@ func Test_scanRepos_SetProgressComplete(t *testing.T) {
|
|||||||
name string
|
name string
|
||||||
repos []string
|
repos []string
|
||||||
wantComplete bool
|
wantComplete bool
|
||||||
wantErr bool
|
|
||||||
}{
|
}{
|
||||||
{
|
{
|
||||||
name: "no repos",
|
name: "no repos",
|
||||||
wantComplete: true,
|
wantComplete: true,
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
name: "one valid repo",
|
name: "one valid repo",
|
||||||
repos: []string{"a"},
|
repos: []string{"a"},
|
||||||
wantErr: true,
|
wantComplete: true,
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -274,12 +276,9 @@ func Test_scanRepos_SetProgressComplete(t *testing.T) {
|
|||||||
src := &Source{
|
src := &Source{
|
||||||
repos: tc.repos,
|
repos: tc.repos,
|
||||||
}
|
}
|
||||||
src.jobSem = semaphore.NewWeighted(1)
|
src.jobPool = &errgroup.Group{}
|
||||||
|
|
||||||
_ = src.scanRepos(context.Background(), nil)
|
_ = src.scanRepos(context.Background(), nil)
|
||||||
if !tc.wantErr {
|
|
||||||
assert.Equal(t, "", src.GetProgress().EncodedResumeInfo)
|
|
||||||
}
|
|
||||||
|
|
||||||
gotComplete := src.GetProgress().PercentComplete == 100
|
gotComplete := src.GetProgress().PercentComplete == 100
|
||||||
if gotComplete != tc.wantComplete {
|
if gotComplete != tc.wantComplete {
|
||||||
|
|||||||
@@ -135,12 +135,15 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
|
|||||||
return errors.Errorf("invalid configuration given for %s source", s.name)
|
return errors.Errorf("invalid configuration given for %s source", s.name)
|
||||||
}
|
}
|
||||||
|
|
||||||
for i, bucket := range bucketsToScan {
|
numBuckets := len(bucketsToScan)
|
||||||
|
log.Debugf("Found %v total buckets to scan", numBuckets)
|
||||||
|
for _, bucket := range bucketsToScan {
|
||||||
if common.IsDone(ctx) {
|
if common.IsDone(ctx) {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
s.SetProgressComplete(i, len(bucketsToScan), fmt.Sprintf("Bucket: %s", bucket), "")
|
cnt := s.Counter.IncSuccess()
|
||||||
|
s.Update(len(bucketsToScan), fmt.Sprintf("Bucket: %s", bucket), "")
|
||||||
|
|
||||||
s.log.Debugf("Scanning bucket: %s", bucket)
|
s.log.Debugf("Scanning bucket: %s", bucket)
|
||||||
region, err := s3manager.GetBucketRegionWithClient(context.Background(), client, bucket)
|
region, err := s3manager.GetBucketRegionWithClient(context.Background(), client, bucket)
|
||||||
@@ -161,6 +164,7 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
|
|||||||
// pf := "public"
|
// pf := "public"
|
||||||
errorCount := sync.Map{}
|
errorCount := sync.Map{}
|
||||||
|
|
||||||
|
log.Debugf("Starting to scan bucket #%d of %d: %s", cnt, numBuckets, bucket)
|
||||||
err = regionalClient.ListObjectsV2PagesWithContext(
|
err = regionalClient.ListObjectsV2PagesWithContext(
|
||||||
ctx, &s3.ListObjectsV2Input{Bucket: &bucket},
|
ctx, &s3.ListObjectsV2Input{Bucket: &bucket},
|
||||||
func(page *s3.ListObjectsV2Output, last bool) bool {
|
func(page *s3.ListObjectsV2Output, last bool) bool {
|
||||||
@@ -172,8 +176,9 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk) err
|
|||||||
s.log.WithError(err).Errorf("could not list objects in s3 bucket: %s", bucket)
|
s.log.WithError(err).Errorf("could not list objects in s3 bucket: %s", bucket)
|
||||||
return errors.WrapPrefix(err, fmt.Sprintf("could not list objects in s3 bucket: %s", bucket), 0)
|
return errors.WrapPrefix(err, fmt.Sprintf("could not list objects in s3 bucket: %s", bucket), 0)
|
||||||
}
|
}
|
||||||
|
log.Debugf("scanned %d/%d buckets", cnt, numBuckets)
|
||||||
}
|
}
|
||||||
s.SetProgressComplete(len(bucketsToScan), len(bucketsToScan), fmt.Sprintf("Completed scanning source %s", s.name), "")
|
s.Complete(fmt.Sprintf("Completed scanning source %s", s.name))
|
||||||
|
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|||||||
+45
-12
@@ -2,6 +2,7 @@ package sources
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"sync"
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
|
|
||||||
"google.golang.org/protobuf/types/known/anypb"
|
"google.golang.org/protobuf/types/known/anypb"
|
||||||
|
|
||||||
@@ -104,39 +105,71 @@ func NewConfig(opts ...func(*Config)) Config {
|
|||||||
return *c
|
return *c
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Counter is a simple monotonically-increasing counter for tracking progress.
|
||||||
|
// It includes a successCnt that records the number of successful scans.
|
||||||
|
// It also includes a totalCnt that records the total number of scans.
|
||||||
|
// This is necessary in order to display the number of concurrent scans
|
||||||
|
// regardless of whether they are successful or not.
|
||||||
|
type Counter struct {
|
||||||
|
successCnt, totalCnt uint32
|
||||||
|
}
|
||||||
|
|
||||||
|
// IncTotal increments the total field of counter and returns the new value.
|
||||||
|
func (c *Counter) IncTotal() int {
|
||||||
|
return int(atomic.AddUint32(&c.totalCnt, 1))
|
||||||
|
}
|
||||||
|
|
||||||
|
// IncSuccess increments the success field of counter and returns the new value.
|
||||||
|
func (c *Counter) IncSuccess() int {
|
||||||
|
return int(atomic.AddUint32(&c.successCnt, 1))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get returns the current successCnt.
|
||||||
|
func (c *Counter) Get() uint32 {
|
||||||
|
return atomic.LoadUint32(&c.successCnt)
|
||||||
|
}
|
||||||
|
|
||||||
// Progress is used to update job completion progress across sources.
|
// Progress is used to update job completion progress across sources.
|
||||||
type Progress struct {
|
type Progress struct {
|
||||||
mut sync.Mutex
|
SectionsCompleted int32
|
||||||
|
SectionsRemaining int32
|
||||||
PercentComplete int64
|
PercentComplete int64
|
||||||
Message string
|
Message string
|
||||||
EncodedResumeInfo string
|
EncodedResumeInfo string
|
||||||
SectionsCompleted int32
|
mut sync.Mutex
|
||||||
SectionsRemaining int32
|
Counter Counter
|
||||||
}
|
}
|
||||||
|
|
||||||
// SetProgressComplete sets job progress information for a running job based on the highest level objects in the source.
|
// Update sets job progress information for a running job based on the highest level objects in the source.
|
||||||
// i is the current iteration in the loop of target scope
|
|
||||||
// scope should be the len(scopedItems)
|
// scope should be the len(scopedItems)
|
||||||
// message is the public facing user information about the current progress
|
// message is the public facing user information about the current progress
|
||||||
// encodedResumeInfo is an optional string representing any information necessary to resume the job if interrupted
|
// encodedResumeInfo is an optional string representing any information necessary to resume the job if interrupted
|
||||||
func (p *Progress) SetProgressComplete(i, scope int, message, encodedResumeInfo string) {
|
func (p *Progress) Update(scope int, message, encodedResumeInfo string) {
|
||||||
p.mut.Lock()
|
p.mut.Lock()
|
||||||
defer p.mut.Unlock()
|
defer p.mut.Unlock()
|
||||||
|
|
||||||
|
i := p.Counter.Get()
|
||||||
|
|
||||||
p.Message = message
|
p.Message = message
|
||||||
p.EncodedResumeInfo = encodedResumeInfo
|
p.EncodedResumeInfo = encodedResumeInfo
|
||||||
p.SectionsCompleted = int32(i)
|
p.SectionsCompleted = int32(i)
|
||||||
p.SectionsRemaining = int32(scope)
|
p.SectionsRemaining = int32(scope)
|
||||||
|
|
||||||
// If the iteration and scope are both 0, completion is 100%.
|
|
||||||
if i == 0 && scope == 0 {
|
|
||||||
p.PercentComplete = 100
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
p.PercentComplete = int64((float64(i) / float64(scope)) * 100)
|
p.PercentComplete = int64((float64(i) / float64(scope)) * 100)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Complete sets job progress information for a completed job.
|
||||||
|
func (p *Progress) Complete(msg string) {
|
||||||
|
p.mut.Lock()
|
||||||
|
defer p.mut.Unlock()
|
||||||
|
|
||||||
|
p.SectionsRemaining = 0
|
||||||
|
p.SectionsCompleted = int32(p.Counter.Get())
|
||||||
|
p.Message = msg
|
||||||
|
p.EncodedResumeInfo = ""
|
||||||
|
p.PercentComplete = 100
|
||||||
|
}
|
||||||
|
|
||||||
// GetProgress gets job completion percentage for metrics reporting.
|
// GetProgress gets job completion percentage for metrics reporting.
|
||||||
func (p *Progress) GetProgress() *Progress {
|
func (p *Progress) GetProgress() *Progress {
|
||||||
p.mut.Lock()
|
p.mut.Lock()
|
||||||
|
|||||||
@@ -0,0 +1,192 @@
|
|||||||
|
package sources
|
||||||
|
|
||||||
|
import (
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/google/go-cmp/cmp"
|
||||||
|
"github.com/google/go-cmp/cmp/cmpopts"
|
||||||
|
"google.golang.org/protobuf/testing/protocmp"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestCounter_IncTotal(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
want int
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "increment total 50 times",
|
||||||
|
want: 50,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "increment total 0 times",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "increment total 1 times",
|
||||||
|
want: 1,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
c := &Counter{}
|
||||||
|
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
wg.Add(tt.want)
|
||||||
|
for i := 0; i < tt.want; i++ {
|
||||||
|
go func(c *Counter) {
|
||||||
|
defer wg.Done()
|
||||||
|
c.IncTotal()
|
||||||
|
}(c)
|
||||||
|
}
|
||||||
|
|
||||||
|
wg.Wait()
|
||||||
|
if got := c.totalCnt; int(got) != tt.want {
|
||||||
|
t.Errorf("Counter.IncTotal() = %v, want %v", got, tt.want)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestCounter_IncSuccess(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
want int
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "increment total 50 times",
|
||||||
|
want: 50,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "increment total 0 times",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "increment total 1 times",
|
||||||
|
want: 1,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
c := &Counter{}
|
||||||
|
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
wg.Add(tt.want)
|
||||||
|
for i := 0; i < tt.want; i++ {
|
||||||
|
go func(c *Counter) {
|
||||||
|
defer wg.Done()
|
||||||
|
c.IncSuccess()
|
||||||
|
}(c)
|
||||||
|
}
|
||||||
|
|
||||||
|
wg.Wait()
|
||||||
|
if got := c.Get(); int(got) != tt.want {
|
||||||
|
t.Errorf("Counter.IncTotal() = %v, want %v", got, tt.want)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProgress_Update(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name, msg, encodeResumeInfo string
|
||||||
|
counter uint32
|
||||||
|
scope, routines, want int
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "update progress, no counter",
|
||||||
|
scope: 100,
|
||||||
|
routines: 10,
|
||||||
|
want: 10,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "update progress, counter 5",
|
||||||
|
scope: 100,
|
||||||
|
routines: 10,
|
||||||
|
counter: 5,
|
||||||
|
want: 15,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "update progress, counter 50",
|
||||||
|
scope: 100,
|
||||||
|
routines: 10,
|
||||||
|
counter: 50,
|
||||||
|
want: 60,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
p := &Progress{}
|
||||||
|
p.Counter.successCnt = tt.counter
|
||||||
|
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
wg.Add(tt.routines)
|
||||||
|
for i := 0; i < tt.routines; i++ {
|
||||||
|
p.Counter.IncSuccess()
|
||||||
|
go func(p *Progress) {
|
||||||
|
defer wg.Done()
|
||||||
|
p.Update(tt.scope, tt.msg, tt.encodeResumeInfo)
|
||||||
|
}(p)
|
||||||
|
}
|
||||||
|
|
||||||
|
wg.Wait()
|
||||||
|
if got := p.PercentComplete; int(got) != tt.want {
|
||||||
|
t.Errorf("Progress.Update(%v, %v, %v): got %v, want %v", tt.scope, tt.msg, tt.encodeResumeInfo, got, tt.want)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProgress_Complete(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name,
|
||||||
|
msg string
|
||||||
|
curProgress, wantProgress *Progress
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "complete progress",
|
||||||
|
msg: "completed scanning thinger",
|
||||||
|
curProgress: &Progress{
|
||||||
|
PercentComplete: 42,
|
||||||
|
Message: "scanning thinger",
|
||||||
|
EncodedResumeInfo: "resume info",
|
||||||
|
SectionsRemaining: 12,
|
||||||
|
SectionsCompleted: 2,
|
||||||
|
},
|
||||||
|
wantProgress: &Progress{
|
||||||
|
PercentComplete: 100,
|
||||||
|
Message: "completed scanning thinger",
|
||||||
|
EncodedResumeInfo: "",
|
||||||
|
SectionsRemaining: 0,
|
||||||
|
SectionsCompleted: 14,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "complete progress with no message",
|
||||||
|
curProgress: &Progress{
|
||||||
|
PercentComplete: 80,
|
||||||
|
Message: "scanning thinger2",
|
||||||
|
EncodedResumeInfo: "resume info2",
|
||||||
|
SectionsRemaining: 2,
|
||||||
|
SectionsCompleted: 42,
|
||||||
|
},
|
||||||
|
wantProgress: &Progress{
|
||||||
|
PercentComplete: 100,
|
||||||
|
EncodedResumeInfo: "",
|
||||||
|
SectionsRemaining: 0,
|
||||||
|
SectionsCompleted: 44,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
tt.curProgress.Counter.successCnt = uint32(tt.wantProgress.SectionsCompleted)
|
||||||
|
tt.curProgress.Complete(tt.msg)
|
||||||
|
|
||||||
|
ignoreOpts := cmpopts.IgnoreFields(Progress{}, "mut", "Counter")
|
||||||
|
if !cmp.Equal(tt.curProgress, tt.wantProgress, protocmp.Transform(), ignoreOpts) {
|
||||||
|
t.Errorf("Progress.Complete(%v): got %v, want %v", tt.msg, tt.curProgress, tt.wantProgress)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user