chore: remove time.sleep at the end of run (#171)

Fixes #59 First I notices we're not waiting for the `PreRun` routine that handles all the channels. Then, to `Done` the routine that listen to the channel, I need to close the channel.
Checkmarx · Aug 24, 2023 · d4fada8 · d4fada8
1 parent 8156dfa
commit d4fada8
Show file tree

Hide file tree

Showing 8 changed files with 74 additions and 41 deletions.
diff --git a/cmd/main.go b/cmd/main.go
@@ -11,7 +11,6 @@ import (
 	"github.com/checkmarx/2ms/lib"
 
 	"sync"
-	"time"
 
 	"github.com/checkmarx/2ms/plugins"
 	"github.com/checkmarx/2ms/reporting"
@@ -129,7 +128,7 @@ func Execute() {
 	rootCmd.AddGroup(&cobra.Group{Title: group, ID: group})
 
 	for _, plugin := range allPlugins {
-		subCommand, err := plugin.DefineCommand(channels)
+		subCommand, err := plugin.DefineCommand(channels.Items, channels.Errors)
 		if err != nil {
 			log.Fatal().Msg(fmt.Sprintf("error while defining command for plugin %s: %s", plugin.GetName(), err.Error()))
 		}
@@ -170,22 +169,32 @@ func preRun(cmd *cobra.Command, args []string) {
 		log.Fatal().Msg(err.Error())
 	}
 
+	channels.WaitGroup.Add(1)
 	go func() {
-		for {
-			select {
-			case item := <-channels.Items:
-				report.TotalItemsScanned++
-				channels.WaitGroup.Add(1)
-				go secrets.Detect(item, secretsChan, channels.WaitGroup, ignoreVar)
-			case secret := <-secretsChan:
-				report.TotalSecretsFound++
-				report.Results[secret.ID] = append(report.Results[secret.ID], secret)
-			case err, ok := <-channels.Errors:
-				if !ok {
-					return
-				}
-				log.Fatal().Msg(err.Error())
-			}
+		defer channels.WaitGroup.Done()
+
+		wgItems := &sync.WaitGroup{}
+		for item := range channels.Items {
+			report.TotalItemsScanned++
+			wgItems.Add(1)
+			go secrets.Detect(item, secretsChan, wgItems, ignoreVar)
+		}
+		wgItems.Wait()
+		close(secretsChan)
+	}()
+
+	channels.WaitGroup.Add(1)
+	go func() {
+		defer channels.WaitGroup.Done()
+		for secret := range secretsChan {
+			report.TotalSecretsFound++
+			report.Results[secret.ID] = append(report.Results[secret.ID], secret)
+		}
+	}()
+
+	go func() {
+		for err := range channels.Errors {
+			log.Fatal().Msg(err.Error())
 		}
 	}()
 }
@@ -195,9 +204,6 @@ func postRun(cmd *cobra.Command, args []string) {
 
 	cfg := config.LoadConfig("2ms", Version)
 
-	// Wait for last secret to be added to report
-	time.Sleep(time.Millisecond * timeSleepInterval)
-
 	// -------------------------------------
 	// Show Report
 	if report.TotalItemsScanned > 0 {

diff --git a/plugins/confluence.go b/plugins/confluence.go
@@ -46,7 +46,7 @@ func (p *ConfluencePlugin) GetAuthorizationHeader() string {
 	return lib.CreateBasicAuthCredentials(p)
 }
 
-func (p *ConfluencePlugin) DefineCommand(channels Channels) (*cobra.Command, error) {
+func (p *ConfluencePlugin) DefineCommand(items chan Item, errors chan error) (*cobra.Command, error) {
 	var confluenceCmd = &cobra.Command{
 		Use:   fmt.Sprintf("%s --%s URL", p.GetName(), argUrl),
 		Short: "Scan Confluence server",
@@ -67,11 +67,14 @@ func (p *ConfluencePlugin) DefineCommand(channels Channels) (*cobra.Command, err
 	confluenceCmd.Run = func(cmd *cobra.Command, args []string) {
 		err := p.initialize(cmd)
 		if err != nil {
-			channels.Errors <- fmt.Errorf("error while initializing confluence plugin: %w", err)
+			errors <- fmt.Errorf("error while initializing confluence plugin: %w", err)
 			return
 		}
 
-		p.getItems(channels.Items, channels.Errors, channels.WaitGroup)
+		wg := &sync.WaitGroup{}
+		p.getItems(items, errors, wg)
+		wg.Wait()
+		close(items)
 	}
 
 	return confluenceCmd, nil
@@ -90,18 +93,14 @@ func (p *ConfluencePlugin) initialize(cmd *cobra.Command) error {
 }
 
 func (p *ConfluencePlugin) getItems(items chan Item, errs chan error, wg *sync.WaitGroup) {
-	p.getSpacesItems(items, errs, wg)
-}
-
-func (p *ConfluencePlugin) getSpacesItems(items chan Item, errs chan error, wg *sync.WaitGroup) {
 	spaces, err := p.getSpaces()
 	if err != nil {
 		errs <- err
 	}
 
 	for _, space := range spaces {
-		go p.getSpaceItems(items, errs, wg, space)
 		wg.Add(1)
+		go p.getSpaceItems(items, errs, wg, space)
 	}
 }
 

diff --git a/plugins/discord.go b/plugins/discord.go
@@ -38,7 +38,7 @@ func (p *DiscordPlugin) GetName() string {
 	return "discord"
 }
 
-func (p *DiscordPlugin) DefineCommand(channels Channels) (*cobra.Command, error) {
+func (p *DiscordPlugin) DefineCommand(items chan Item, errors chan error) (*cobra.Command, error) {
 	var discordCmd = &cobra.Command{
 		Use:   fmt.Sprintf("%s --%s TOKEN --%s SERVER", p.GetName(), tokenFlag, serversFlag),
 		Short: "Scan Discord server",
@@ -63,11 +63,14 @@ func (p *DiscordPlugin) DefineCommand(channels Channels) (*cobra.Command, error)
 	discordCmd.Run = func(cmd *cobra.Command, args []string) {
 		err := p.initialize(cmd)
 		if err != nil {
-			channels.Errors <- fmt.Errorf("discord plugin initialization failed: %w", err)
+			errors <- fmt.Errorf("discord plugin initialization failed: %w", err)
 			return
 		}
 
-		p.getItems(channels.Items, channels.Errors, channels.WaitGroup)
+		wg := &sync.WaitGroup{}
+		p.getItems(items, errors, wg)
+		wg.Wait()
+		close(items)
 	}
 
 	return discordCmd, nil

diff --git a/plugins/filesystem.go b/plugins/filesystem.go
@@ -29,14 +29,18 @@ func (p *FileSystemPlugin) GetName() string {
 	return "filesystem"
 }
 
-func (p *FileSystemPlugin) DefineCommand(channels Channels) (*cobra.Command, error) {
+func (p *FileSystemPlugin) DefineCommand(items chan Item, errors chan error) (*cobra.Command, error) {
 	var cmd = &cobra.Command{
 		Use:   fmt.Sprintf("%s --%s PATH", p.GetName(), flagFolder),
 		Short: "Scan local folder",
 		Long:  "Scan local folder for sensitive information",
 		Run: func(cmd *cobra.Command, args []string) {
 			log.Info().Msg("Folder plugin started")
-			p.getFiles(channels.Items, channels.Errors, channels.WaitGroup)
+
+			wg := &sync.WaitGroup{}
+			p.getFiles(items, errors, wg)
+			wg.Wait()
+			close(items)
 		},
 	}
 

diff --git a/plugins/git.go b/plugins/git.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"strings"
+	"sync"
 
 	"github.com/gitleaks/go-gitdiff/gitdiff"
 	"github.com/rs/zerolog/log"
@@ -29,8 +30,12 @@ func (p *GitPlugin) GetName() string {
 	return "git"
 }
 
-func (p *GitPlugin) DefineCommand(channels Channels) (*cobra.Command, error) {
-	p.Channels = channels
+func (p *GitPlugin) DefineCommand(items chan Item, errors chan error) (*cobra.Command, error) {
+	p.Channels = Channels{
+		Items:     items,
+		Errors:    errors,
+		WaitGroup: &sync.WaitGroup{},
+	}
 
 	command := &cobra.Command{
 		Use:   fmt.Sprintf("%s <CLONED_REPO>", p.GetName()),
@@ -39,7 +44,9 @@ func (p *GitPlugin) DefineCommand(channels Channels) (*cobra.Command, error) {
 		Args:  cobra.MatchAll(cobra.ExactArgs(1), validGitRepoArgs),
 		Run: func(cmd *cobra.Command, args []string) {
 			log.Info().Msg("Git plugin started")
-			p.scanGit(args[0], p.buildScanOptions(), channels.Items, channels.Errors)
+			p.scanGit(args[0], p.buildScanOptions(), p.Channels.Items, p.Channels.Errors)
+			p.WaitGroup.Wait()
+			close(items)
 		},
 	}
 	flags := command.Flags()

diff --git a/plugins/paligo.go b/plugins/paligo.go
@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/checkmarx/2ms/lib"
@@ -54,8 +55,12 @@ func (p *PaligoPlugin) GetName() string {
 	return "paligo"
 }
 
-func (p *PaligoPlugin) DefineCommand(channels Channels) (*cobra.Command, error) {
-	p.Channels = channels
+func (p *PaligoPlugin) DefineCommand(items chan Item, errors chan error) (*cobra.Command, error) {
+	p.Channels = Channels{
+		Items:     items,
+		Errors:    errors,
+		WaitGroup: &sync.WaitGroup{},
+	}
 
 	command := &cobra.Command{
 		Use: fmt.Sprintf("%s --%s %s --%s %s --%s %s",
@@ -73,6 +78,8 @@ func (p *PaligoPlugin) DefineCommand(channels Channels) (*cobra.Command, error)
 			}
 			log.Info().Msg("Paligo plugin started")
 			p.getItems()
+			p.WaitGroup.Wait()
+			close(items)
 		},
 	}
 

diff --git a/plugins/plugins.go b/plugins/plugins.go
@@ -27,5 +27,5 @@ type Channels struct {
 
 type IPlugin interface {
 	GetName() string
-	DefineCommand(channels Channels) (*cobra.Command, error)
+	DefineCommand(items chan Item, errors chan error) (*cobra.Command, error)
 }
diff --git a/plugins/slack.go b/plugins/slack.go
@@ -3,6 +3,7 @@ package plugins
 import (
 	"fmt"
 	"strconv"
+	"sync"
 	"time"
 
 	"github.com/rs/zerolog/log"
@@ -38,15 +39,21 @@ var (
 	messagesCountArg    int
 )
 
-func (p *SlackPlugin) DefineCommand(channels Channels) (*cobra.Command, error) {
-	p.Channels = channels
+func (p *SlackPlugin) DefineCommand(items chan Item, errors chan error) (*cobra.Command, error) {
+	p.Channels = Channels{
+		Items:     items,
+		Errors:    errors,
+		WaitGroup: &sync.WaitGroup{},
+	}
 
 	command := &cobra.Command{
 		Use:   fmt.Sprintf("%s --%s TOKEN --%s TEAM", p.GetName(), slackTokenFlag, slackTeamFlag),
 		Short: "Scan Slack team",
 		Long:  "Scan Slack team for sensitive information.",
 		Run: func(cmd *cobra.Command, args []string) {
 			p.getItems()
+			p.Channels.WaitGroup.Wait()
+			close(items)
 		},
 	}