aquasecurity · geyslan · Jun 28, 2024 · Jun 26, 2024 · Jun 26, 2024 · Jun 27, 2024
diff --git a/pkg/logger/logger.go b/pkg/logger/logger.go
@@ -6,6 +6,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"go.uber.org/zap"
@@ -157,6 +158,9 @@ func aggregateLog(l *Logger, lvl Level, msg string, ci *callerInfo) bool {
 // It does NOT override those options.
 // For the case where innerAggregation is set to true, it will only aggregate if pkg logger's aggregation config is also set to true.
 func Log(lvl Level, innerAggregation bool, msg string, keysAndValues ...interface{}) {
+	pkgMutex.RLock()
+	defer pkgMutex.RUnlock()
+
 	if innerAggregation {
 		switch lvl {
 		case DebugLevel:
@@ -219,6 +223,9 @@ func debugw(skip int, l *Logger, msg string, keysAndValues ...interface{}) {
 }
 
 func Debugw(msg string, keysAndValues ...interface{}) {
+	pkgMutex.RLock()
+	defer pkgMutex.RUnlock()
+
 	debugw(1, pkgLogger, msg, keysAndValues...)
 }
 
@@ -241,6 +248,9 @@ func infow(skip int, l *Logger, msg string, keysAndValues ...interface{}) {
 }
 
 func Infow(msg string, keysAndValues ...interface{}) {
+	pkgMutex.RLock()
+	defer pkgMutex.RUnlock()
+
 	infow(1, pkgLogger, msg, keysAndValues...)
 }
 
@@ -263,6 +273,9 @@ func warnw(skip int, l *Logger, msg string, keysAndValues ...interface{}) {
 }
 
 func Warnw(msg string, keysAndValues ...interface{}) {
+	pkgMutex.RLock()
+	defer pkgMutex.RUnlock()
+
 	warnw(1, pkgLogger, msg, keysAndValues...)
 }
 
@@ -285,6 +298,9 @@ func errorw(skip int, l *Logger, msg string, keysAndValues ...interface{}) {
 }
 
 func Errorw(msg string, keysAndValues ...interface{}) {
+	pkgMutex.RLock()
+	defer pkgMutex.RUnlock()
+
 	errorw(1, pkgLogger, msg, keysAndValues...)
 }
 
@@ -307,6 +323,9 @@ func fatalw(skip int, l *Logger, msg string, keysAndValues ...interface{}) {
 }
 
 func Fatalw(msg string, keysAndValues ...interface{}) {
+	pkgMutex.RLock()
+	defer pkgMutex.RUnlock()
+
 	fatalw(1, pkgLogger, msg, keysAndValues...)
 }
 
@@ -322,16 +341,22 @@ func (l *Logger) Sync() error {
 var (
 	// Package-level Logger
 	pkgLogger = &Logger{}
+	pkgMutex  = &sync.RWMutex{}
 )
 
 // Current returns the package-level base logger
 func Current() *Logger {
+	pkgMutex.RLock()
+	defer pkgMutex.RUnlock()
+
 	return pkgLogger
 }
 
 // SetLogger sets package-level base logger
-// It's not thread safe so if required use it always at the beginning
 func SetLogger(l LoggerInterface) {
+	pkgMutex.Lock()
+	defer pkgMutex.Unlock()
+
 	if l == nil {
 		panic("Logger cannot be nil")
 	}
@@ -341,18 +366,27 @@ func SetLogger(l LoggerInterface) {
 
 // GetLogger gets the package-level base logger
 func GetLogger() LoggerInterface {
+	pkgMutex.RLock()
+	defer pkgMutex.RUnlock()
+
 	return pkgLogger.l
 }
 
 // SetLevel sets package-level base logger level,
 // it is threadsafe
 func SetLevel(level Level) {
+	pkgMutex.Lock()
+	defer pkgMutex.Unlock()
+
 	pkgLogger.cfg.SetLevel(level)
 }
 
 // Init sets the package-level base logger using given config
 // It's not thread safe so if required use it always at the beginning
 func Init(cfg LoggingConfig) {
+	pkgMutex.Lock()
+	defer pkgMutex.Unlock()
+
 	// set the config
 	pkgLogger.cfg = cfg
 

diff --git a/tests/integration/capture_test.go b/tests/integration/capture_test.go
@@ -12,7 +12,6 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
-	"gotest.tools/assert"
 
 	"github.com/aquasecurity/tracee/tests/testutils"
 )
@@ -91,16 +90,33 @@ func Test_TraceeCapture(t *testing.T) {
 				t.Fatal("tracee is already running")
 			}
 
+			var failed bool
+
 			captureDir := tc.directory + "/out/host/"
 			err := tc.test(t, captureDir, homeDir)
 			if err != nil {
-				t.Errorf("test %s failed: %v", tc.name, err)
-				cmdErrs := running.Stop() // stop tracee
-				require.Empty(t, cmdErrs)
-				t.Fail()
+				failed = true
+				t.Logf("test %s failed: %v", tc.name, err)
 			}
+			defer func() {
+				t.Logf("removing directory %s", tc.directory)
+				err := os.RemoveAll(tc.directory)
+				if err != nil {
+					t.Logf("failed to remove directory %s: %v", tc.directory, err)
+				}
+			}()
+
 			cmdErrs := running.Stop() // stop tracee
-			require.Empty(t, cmdErrs)
+			if len(cmdErrs) > 0 {
+				failed = true
+				t.Logf("failed to stop tracee: %v", cmdErrs)
+			} else {
+				t.Logf("  --- stopped tracee ---")
+			}
+
+			if failed {
+				t.Fail()
+			}
 		})
 	}
 }
@@ -256,10 +272,12 @@ func readWritePipe(t *testing.T, captureDir string, workingDir string) error {
 	}()
 
 	finfo, err := pipe.Stat()
-	require.NoError(t, err)
+	if err != nil {
+		return err
+	}
 	statInfo, ok := finfo.Sys().(*syscall.Stat_t)
 	if !ok {
-		t.Logf("type assertion failed: expected *syscall.Stat_t")
+		return fmt.Errorf("type assertion failed: expected *syscall.Stat_t")
 	}
 	inode := statInfo.Ino
 
@@ -323,10 +341,17 @@ func assertEntries(t *testing.T, captureDir string, input string, readOut string
 		}
 	}
 	// Found 2 entries
-	assert.Equal(t, found, 2)
+	if found != 2 {
+		return fmt.Errorf("expected 2 entries in capture dir, found %d", found)
+	}
 
 	// Compare captured data to expected data
-	assert.Equal(t, input, string(writeCaptureFile))
-	assert.Equal(t, readOut, string(readCaptureFile))
+	if string(writeCaptureFile) != input {
+		return fmt.Errorf("expected write capture file %s, got %s", input, writeCaptureFile)
+	}
+	if string(readCaptureFile) != readOut {
+		return fmt.Errorf("expected read capture file %s, got %s", readOut, readCaptureFile)
+	}
+
 	return nil
 }
diff --git a/tests/integration/dependencies_test.go b/tests/integration/dependencies_test.go
@@ -4,12 +4,11 @@ import (
 	"context"
 	"fmt"
 	"os/exec"
+	"slices"
 	"strings"
 	"testing"
 	"time"
 
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
 
 	"github.com/aquasecurity/tracee/pkg/config"
@@ -87,8 +86,6 @@ func Test_EventsDependencies(t *testing.T) {
 		},
 	}
 
-	logOutChan, restoreLogger := testutils.SetTestLogger(logger.DebugLevel)
-
 	// Each test will run a test binary that triggers the "exec_test" event.
 	// Upon its execution, which events are evicted and which not will be tested
 	createCmdEvents := func(expectedEventsIDs []events.ID, unexpectedEventsIDs []events.ID) []cmdEvents {
@@ -132,8 +129,11 @@ func Test_EventsDependencies(t *testing.T) {
 				policy.Snapshots().Store(testConfig.Policies)
 
 				ctx, cancel := context.WithCancel(context.Background())
-				done := make(chan struct{})
-				logsResultChan := testutils.TestLogs(t, testCaseInst.expectedLogs, logOutChan, done)
+
+				// set test logger
+				logsDone := make(chan struct{})
+				logOutChan, restoreLogger := testutils.SetTestLogger(t, logger.DebugLevel)
+				logsResultChan := testutils.TestLogs(t, testCaseInst.expectedLogs, logOutChan, logsDone)
 
 				// start tracee
 				trc, err := startTracee(ctx, t, testConfig, nil, nil)
@@ -142,16 +142,15 @@ func Test_EventsDependencies(t *testing.T) {
 					t.Fatal(err)
 				}
 				t.Logf("  --- started tracee ---")
-
-				stream := trc.SubscribeAll()
-				defer trc.Unsubscribe(stream)
-
 				err = waitForTraceeStart(trc)
 				if err != nil {
 					cancel()
 					t.Fatal(err)
 				}
 
+				stream := trc.SubscribeAll()
+				defer trc.Unsubscribe(stream)
+
 				// start a goroutine to read events from the channel into the buffer
 				buf := newEventBuffer()
 				go func(ctx context.Context, buf *eventBuffer) {
@@ -165,30 +164,54 @@ func Test_EventsDependencies(t *testing.T) {
 					}
 				}(ctx, buf)
 
-				testAttachedKprobes(t, testCaseInst.expectedKprobes, testCaseInst.unexpectedKprobes)
+				var failed bool
+				var testCmdEvents []cmdEvents
+
+				// test kprobes
+				err = testAttachedKprobes(testCaseInst.expectedKprobes, testCaseInst.unexpectedKprobes)
+				if err != nil {
+					t.Logf("Test %s failed: %v", t.Name(), err)
+					failed = true
+					goto cleanup
+				}
 
-				// Test events
-				testCmdEvents := createCmdEvents(testCaseInst.expectedEvents, testCaseInst.unexpectedEvents)
-				require.NoError(t, ExpectAtLeastOneForEach(t, testCmdEvents, buf, false))
+				// test events
+				testCmdEvents = createCmdEvents(testCaseInst.expectedEvents, testCaseInst.unexpectedEvents)
+				err = ExpectAtLeastOneForEach(t, testCmdEvents, buf, false)
+				if err != nil {
+					t.Logf("Test %s failed: %v", t.Name(), err)
+					failed = true
+					goto cleanup
+				}
 
+				close(logsDone)
+				if !<-logsResultChan {
+					t.Logf("Test %s failed: not all logs were found", t.Name())
+					failed = true
+				}
+			cleanup:
+				// ensure that logsDone is closed
+				select {
+				case <-logsDone:
+				default:
+					close(logsDone)
+				}
+				restoreLogger()
 				cancel()
 				errStop := waitForTraceeStop(trc)
 				if errStop != nil {
 					t.Log(errStop)
+					failed = true
 				} else {
 					t.Logf("  --- stopped tracee ---")
 				}
-				close(done)
 
-				assert.True(t, <-logsResultChan)
+				if failed {
+					t.Fail()
+				}
 			},
 		)
 	}
-	// Wait for all Tracee's goroutines to finish
-	// TODO: Remove this sleep once all Tracee's goroutines are guaranteed to complete
-	// by the time tracee.Running() returns false.
-	time.Sleep(15 * time.Second)
-	restoreLogger()
 }
 
 func createGenericEventForCmdEvents(eventId events.ID) trace.Event {
@@ -227,18 +250,26 @@ func GetAttachedKprobes() ([]string, error) {
 	return probes, nil
 }
 
-func testAttachedKprobes(t *testing.T, expectedKprobes []string, unexpectedKprobes []string) {
+func testAttachedKprobes(expectedKprobes []string, unexpectedKprobes []string) error {
 	// Get the initial list of kprobes
 	attachedKprobes, err := GetAttachedKprobes()
-	require.NoError(t, err)
+	if err != nil {
+		return err
+	}
 
 	// Check if the expected kprobes were added
 	for _, probe := range expectedKprobes {
-		assert.Contains(t, attachedKprobes, probe)
+		if !slices.Contains(attachedKprobes, probe) {
+			return fmt.Errorf("expected kprobe %s not found", probe)
+		}
 	}
 
 	// Check if the unexpected kprobes were added
 	for _, probe := range unexpectedKprobes {
-		assert.NotContains(t, attachedKprobes, probe)
+		if slices.Contains(attachedKprobes, probe) {
+			return fmt.Errorf("unexpected kprobe %s found", probe)
+		}
 	}
+
+	return nil
 }
diff --git a/tests/testutils/logger.go b/tests/testutils/logger.go
@@ -15,7 +15,9 @@ import (
 // SetTestLogger create a logger which prints the logs to the returned channel.
 // This function is meant to be used by tests to check logs, and by that test the
 // flow of Tracee from outside.
-func SetTestLogger(l logger.Level) (loggerOutput <-chan []byte, restoreLogger func()) {
+func SetTestLogger(t *testing.T, l logger.Level) (loggerOutput <-chan []byte, restoreLogger func()) {
+	t.Logf("  --- setting test logger ---")
+
 	mw, logChan := newChannelWriter()
 	chanLogger := logger.NewLogger(
 		logger.LoggerConfig{
@@ -26,6 +28,7 @@ func SetTestLogger(l logger.Level) (loggerOutput <-chan []byte, restoreLogger fu
 	)
 	currentLogger := logger.GetLogger()
 	restoreLogger = func() {
+		t.Logf("  --- restoring default logger ---")
 		err := chanLogger.Sync()
 		logger.SetLogger(currentLogger)
 		mw.Close()