Fix channel issues

cmd/main.go

@@ -76,7 +76,7 @@ var report = reporting.Init()
 var secretsChan = make(chan *secrets.Secret)
 var secretsExtrasChan = make(chan *secrets.Secret)
 var validationChan = make(chan *secrets.Secret)
-var cvssScoreChan = make(chan *secrets.Secret)
+var cvssScoreWithoutValidationChan = make(chan *secrets.Secret)
 
 func Execute() (int, error) {
 	vConfig.SetEnvPrefix(envPrefix)
@@ -150,11 +150,11 @@ func preRun(pluginName string, cmd *cobra.Command, args []string) error {
 
 	if validateVar {
 		channels.WaitGroup.Add(1)
-		go processValidation(engine)
+		go processValidationAndScoreWithValidation(engine)
 	}
 
 	channels.WaitGroup.Add(1)
-	go processScore(engine)
+	go processScoreWithoutValidation(engine)
 
 	return nil
 }

cmd/workers.go

@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"github.com/checkmarx/2ms/lib/secrets"
 	"sync"
 
 	"github.com/checkmarx/2ms/engine"
@@ -29,13 +30,13 @@ func processSecrets() {
 		if validateVar {
 			validationChan <- secret
 		} else {
-			cvssScoreChan <- secret
+			cvssScoreWithoutValidationChan <- secret
 		}
 		report.Results[secret.ID] = append(report.Results[secret.ID], secret)
 	}
 	close(secretsExtrasChan)
 	close(validationChan)
-	close(cvssScoreChan)
+	close(cvssScoreWithoutValidationChan)
 }
 
 func processSecretsExtras() {
@@ -49,30 +50,29 @@ func processSecretsExtras() {
 	wgExtras.Wait()
 }
 
-func processValidation(engine *engine.Engine) {
+func processValidationAndScoreWithValidation(engine *engine.Engine) {
 	defer channels.WaitGroup.Done()
 
 	wgValidation := &sync.WaitGroup{}
 	for secret := range validationChan {
-		wgValidation.Add(1)
-		go func() {
-			wgValidation.Done()
-			engine.RegisterForValidation(secret)
-			cvssScoreChan <- secret
-		}()
+		wgValidation.Add(2)
+		go func(secret *secrets.Secret, wg *sync.WaitGroup) {
+			engine.RegisterForValidation(secret, wg)
+			engine.Score(secret, true, wg)
+		}(secret, wgValidation)
 	}
 	wgValidation.Wait()
 
 	engine.Validate()
 }
 
-func processScore(engine *engine.Engine) {
+func processScoreWithoutValidation(engine *engine.Engine) {
 	defer channels.WaitGroup.Done()
 
 	wgScore := &sync.WaitGroup{}
-	for secret := range cvssScoreChan {
+	for secret := range cvssScoreWithoutValidationChan {
 		wgScore.Add(1)
-		go engine.Score(secret, validateVar, wgScore)
+		go engine.Score(secret, false, wgScore)
 	}
 	wgScore.Wait()
 }

engine/engine.go

@@ -131,7 +131,8 @@ func (e *Engine) AddRegexRules(patterns []string) error {
 	return nil
 }
 
-func (s *Engine) RegisterForValidation(secret *secrets.Secret) {
+func (s *Engine) RegisterForValidation(secret *secrets.Secret, wg *sync.WaitGroup) {
+	defer wg.Done()
 	s.validator.RegisterForValidation(secret)
 }