Parallelization: fix wrong length for error handling (again)

DataDog · Feb 10, 2022 · f2ce722 · f2ce722
1 parent 02b8784
commit f2ce722
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 19 deletions.
diff --git a/cmd/stratus/cleanup_cmd.go b/cmd/stratus/cleanup_cmd.go
@@ -2,13 +2,11 @@ package main
 
 import (
 	"errors"
-	"github.com/datadog/stratus-red-team/internal/utils"
-	"log"
-	"os"
-
 	"github.com/datadog/stratus-red-team/pkg/stratus"
 	"github.com/datadog/stratus-red-team/pkg/stratus/runner"
 	"github.com/spf13/cobra"
+	"log"
+	"os"
 )
 
 var flagForceCleanup bool
@@ -54,7 +52,7 @@ func buildCleanupCmd() *cobra.Command {
 }
 
 func doCleanupCmd(techniques []*stratus.AttackTechnique) {
-	workerCount := utils.Min(len(techniques), maxWorkerCount)
+	workerCount := len(techniques)
 	techniquesChan := make(chan *stratus.AttackTechnique, workerCount)
 	errorsChan := make(chan error, workerCount)
 	for i := 0; i < workerCount; i++ {
@@ -65,7 +63,7 @@ func doCleanupCmd(techniques []*stratus.AttackTechnique) {
 	}
 	close(techniquesChan)
 
-	hadError := handleErrorsChannel(errorsChan, len(techniques))
+	hadError := handleErrorsChannel(errorsChan, workerCount)
 	doStatusCmd(techniques)
 	if hadError {
 		os.Exit(1)

diff --git a/cmd/stratus/detonate_cmd.go b/cmd/stratus/detonate_cmd.go
@@ -49,7 +49,7 @@ func buildDetonateCmd() *cobra.Command {
 	return detonateCmd
 }
 func doDetonateCmd(techniques []*stratus.AttackTechnique, cleanup bool) {
-	workerCount := utils.Min(len(techniques), maxWorkerCount)
+	workerCount := len(techniques)
 	techniquesChan := make(chan *stratus.AttackTechnique, workerCount)
 	errorsChan := make(chan error, workerCount)
 
@@ -64,7 +64,7 @@ func doDetonateCmd(techniques []*stratus.AttackTechnique, cleanup bool) {
 	}
 	close(techniquesChan)
 
-	if hadError := handleErrorsChannel(errorsChan, len(techniques)); hadError {
+	if hadError := handleErrorsChannel(errorsChan, workerCount); hadError {
 		os.Exit(1)
 	}
 }

diff --git a/cmd/stratus/main.go b/cmd/stratus/main.go
@@ -1,8 +1,6 @@
 package main
 
 import (
-	"runtime"
-
 	_ "github.com/datadog/stratus-red-team/internal/attacktechniques"
 	"github.com/spf13/cobra"
 )
@@ -11,10 +9,6 @@ var rootCmd = &cobra.Command{
 	Use: "stratus",
 }
 
-var (
-	maxWorkerCount = runtime.GOMAXPROCS(0)
-)
-
 func init() {
 	listCmd := buildListCmd()
 	showCmd := buildShowCmd()

diff --git a/cmd/stratus/revert_cmd.go b/cmd/stratus/revert_cmd.go
@@ -2,7 +2,6 @@ package main
 
 import (
 	"errors"
-	"github.com/datadog/stratus-red-team/internal/utils"
 	"log"
 	"os"
 
@@ -43,7 +42,7 @@ func buildRevertCmd() *cobra.Command {
 }
 
 func doRevertCmd(techniques []*stratus.AttackTechnique) {
-	workerCount := utils.Min(len(techniques), maxWorkerCount)
+	workerCount := len(techniques)
 	techniquesChan := make(chan *stratus.AttackTechnique, workerCount)
 	errorsChan := make(chan error, workerCount)
 
@@ -58,7 +57,7 @@ func doRevertCmd(techniques []*stratus.AttackTechnique) {
 	}
 	close(techniquesChan)
 
-	hadError := handleErrorsChannel(errorsChan, len(techniques))
+	hadError := handleErrorsChannel(errorsChan, workerCount)
 	doStatusCmd(techniques)
 	if hadError {
 		os.Exit(1)

diff --git a/cmd/stratus/warmup_cmd.go b/cmd/stratus/warmup_cmd.go
@@ -2,7 +2,6 @@ package main
 
 import (
 	"errors"
-	"github.com/datadog/stratus-red-team/internal/utils"
 	"os"
 
 	"github.com/datadog/stratus-red-team/pkg/stratus"
@@ -42,7 +41,7 @@ func buildWarmupCmd() *cobra.Command {
 }
 
 func doWarmupCmd(techniques []*stratus.AttackTechnique) {
-	workerCount := utils.Min(len(techniques), maxWorkerCount)
+	workerCount := len(techniques)
 	techniquesChan := make(chan *stratus.AttackTechnique, workerCount)
 	errorsChan := make(chan error, workerCount)
 	for i := 0; i < workerCount; i++ {