From 150c15d815bde4946359bd6fd581a83207361535 Mon Sep 17 00:00:00 2001
From: Evan Jones <evan.jones@datadoghq.com>
Date: Tue, 8 Aug 2023 13:10:37 -0400
Subject: [PATCH] v2/pkg/stratus/runner: Allow concurrent NewRunner()

The tests in github.com/datadog/threatest failed when I ran them on
my machine, because they call NewRunner() from multiple goroutines at
the same time. If you have never run them before, they will attempt
to create the directories at the same time, causing all but one to
panic(). This adds a test to reproduce this, and fixes it by ignoring
the "already exists" errors.
---
 v2/internal/state/state.go           | 13 +++++++---
 v2/pkg/stratus/runner/runner_test.go | 38 +++++++++++++++++++++++++++-
 2 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/v2/internal/state/state.go b/v2/internal/state/state.go
index cb16d0a8..77437ca7 100644
--- a/v2/internal/state/state.go
+++ b/v2/internal/state/state.go
@@ -2,11 +2,14 @@ package state
 
 import (
 	"encoding/json"
-	"github.com/datadog/stratus-red-team/v2/internal/utils"
-	"github.com/datadog/stratus-red-team/v2/pkg/stratus"
+	"errors"
+	"io/fs"
 	"log"
 	"os"
 	"path/filepath"
+
+	"github.com/datadog/stratus-red-team/v2/internal/utils"
+	"github.com/datadog/stratus-red-team/v2/pkg/stratus"
 )
 
 const StratusStateDirectoryName = ".stratus-red-team"
@@ -76,14 +79,16 @@ func (m *FileSystemStateManager) Initialize() {
 	if !m.FileSystem.FileExists(m.RootDirectory) {
 		log.Println("Creating " + m.RootDirectory + " as it doesn't exist yet")
 		err := m.FileSystem.CreateDirectory(m.RootDirectory, 0744)
-		if err != nil {
+		// ignore "already exists" errors caused by concurrent instances
+		if err != nil && !errors.Is(err, fs.ErrExist) {
 			panic("Unable to create persistent directory: " + err.Error())
 		}
 	}
 
 	if !m.FileSystem.FileExists(m.getTechniqueStateDirectory()) {
 		err := m.FileSystem.CreateDirectory(m.getTechniqueStateDirectory(), 0744)
-		if err != nil {
+		// ignore "already exists" errors caused by concurrent instances
+		if err != nil && !errors.Is(err, fs.ErrExist) {
 			panic("Unable to create persistent directory: " + err.Error())
 		}
 	}
diff --git a/v2/pkg/stratus/runner/runner_test.go b/v2/pkg/stratus/runner/runner_test.go
index ee4ba0fc..9a220bd4 100644
--- a/v2/pkg/stratus/runner/runner_test.go
+++ b/v2/pkg/stratus/runner/runner_test.go
@@ -2,12 +2,17 @@ package runner
 
 import (
 	"errors"
+	"os"
+	"sync"
+	"testing"
+
+	_ "github.com/datadog/stratus-red-team/v2/internal/attacktechniques/aws/credential-access/ec2-get-password-data"
 	statemocks "github.com/datadog/stratus-red-team/v2/internal/state/mocks"
 	"github.com/datadog/stratus-red-team/v2/pkg/stratus"
 	"github.com/datadog/stratus-red-team/v2/pkg/stratus/runner/mocks"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
-	"testing"
+	"github.com/stretchr/testify/require"
 )
 
 func TestRunnerWarmUp(t *testing.T) {
@@ -457,3 +462,34 @@ func TestRunnerCleanup(t *testing.T) {
 		t.Run(scenario[i].Name, func(t *testing.T) { scenario[i].CheckExpectations(t, terraform, state, err) })
 	}
 }
+
+func TestNewRunnerParallel(t *testing.T) {
+	// check that we can create the runner in multiple goroutines without hitting errors
+	// set the user's home dir to a temporary dir during this test
+	tempDir := t.TempDir()
+	const homeEnvVar = "HOME"
+	previousHome := os.Getenv(homeEnvVar)
+	defer func() {
+		os.Setenv(homeEnvVar, previousHome)
+	}()
+	os.Setenv(homeEnvVar, tempDir)
+	homeDir, err := os.UserHomeDir()
+	require.NoError(t, err)
+	require.Equal(t, homeDir, tempDir)
+
+	technique := stratus.GetRegistry().GetAttackTechniqueByName(
+		"aws.credential-access.ec2-get-password-data")
+	require.NotNil(t, technique)
+
+	const numGoroutines = 3
+	var wg sync.WaitGroup
+	wg.Add(numGoroutines)
+	for i := 0; i < numGoroutines; i++ {
+		go func() {
+			defer wg.Done()
+			NewRunner(technique, StratusRunnerNoForce)
+		}()
+	}
+
+	wg.Wait()
+}