From b5c9cb87825a02b4a74f36d212d08cf8fb4c71bf Mon Sep 17 00:00:00 2001
From: Tiago Silva <tiago.silva@goteleport.com>
Date: Wed, 24 Jul 2024 09:48:08 +0100
Subject: [PATCH] handle comments

---
 .../authorizedkeys/authorized_keys.go         | 17 +++++--------
 .../authorizedkeys/supervisor.go              | 24 ++++++++++++-------
 .../authorizedkeys/supervisor_test.go         |  2 +-
 3 files changed, 23 insertions(+), 20 deletions(-)

diff --git a/lib/secretsscanner/authorizedkeys/authorized_keys.go b/lib/secretsscanner/authorizedkeys/authorized_keys.go
index ebb457704e9dd..3890cec4b5be2 100644
--- a/lib/secretsscanner/authorizedkeys/authorized_keys.go
+++ b/lib/secretsscanner/authorizedkeys/authorized_keys.go
@@ -89,9 +89,7 @@ type WatcherConfig struct {
 // Returns [ErrUnsupportedPlatform] if the operating system is not supported.
 func NewWatcher(ctx context.Context, config WatcherConfig) (*Watcher, error) {
 
-	switch getOS(config) {
-	case constants.LinuxOS:
-	default:
+	if getOS(config) != constants.LinuxOS {
 		return nil, trace.Wrap(ErrUnsupportedPlatform)
 	}
 
@@ -258,7 +256,9 @@ func (w *Watcher) fetchAndReportAuthorizedKeys(
 		}
 
 		hostKeys, err := w.parseAuthorizedKeysFile(u, authorizedKeysPath)
-		if err != nil {
+		if errors.Is(err, os.ErrNotExist) {
+			continue
+		} else if err != nil {
 			w.logger.Warn("Failed to parse authorized_keys file", "error", err)
 			continue
 		}
@@ -273,10 +273,7 @@ func (w *Watcher) fetchAndReportAuthorizedKeys(
 	const maxKeysPerReport = 500
 	for i := 0; i < len(keys); i += maxKeysPerReport {
 		start := i
-		end := i + maxKeysPerReport
-		if end > len(keys) {
-			end = len(keys)
-		}
+		end := min(i+maxKeysPerReport, len(keys))
 		if err := stream.Send(
 			&accessgraphsecretsv1pb.ReportAuthorizedKeysRequest{
 				Keys:      keys[start:end],
@@ -338,9 +335,7 @@ func userList(ctx context.Context, log *slog.Logger, filePath string) ([]user.Us
 
 func (w *Watcher) parseAuthorizedKeysFile(u user.User, authorizedKeysPath string) ([]*accessgraphsecretsv1pb.AuthorizedKey, error) {
 	file, err := os.Open(authorizedKeysPath)
-	if errors.Is(err, os.ErrNotExist) {
-		return nil, nil
-	} else if err != nil {
+	if err != nil {
 		return nil, trace.Wrap(err)
 	}
 	defer func() {
diff --git a/lib/secretsscanner/authorizedkeys/supervisor.go b/lib/secretsscanner/authorizedkeys/supervisor.go
index 29a64a84c775f..c0048abb07877 100644
--- a/lib/secretsscanner/authorizedkeys/supervisor.go
+++ b/lib/secretsscanner/authorizedkeys/supervisor.go
@@ -56,13 +56,23 @@ func supervisorRunner(parentCtx context.Context, cfg supervisorRunnerConfig) err
 		mu           sync.Mutex
 	)
 
+	getIsRunning := func() bool {
+		mu.Lock()
+		defer mu.Unlock()
+		return isRunning
+	}
+
+	setIsRunning := func(s bool) {
+		mu.Lock()
+		defer mu.Unlock()
+		isRunning = s
+	}
+
 	runRoutine := func(ctx context.Context, cancel context.CancelCauseFunc) {
 		defer func() {
 			wg.Done()
 			cancel(errShutdown)
-			mu.Lock()
-			isRunning = false
-			mu.Unlock()
+			setIsRunning(false)
 		}()
 		if err := cfg.runner(ctx); err != nil && !errors.Is(err, errShutdown) {
 			cfg.logger.WarnContext(ctx, "Runner failed", "error", err)
@@ -75,14 +85,12 @@ func supervisorRunner(parentCtx context.Context, cfg supervisorRunnerConfig) err
 		switch enabled, err := cfg.checkIfMonitorEnabled(parentCtx); {
 		case err != nil:
 			cfg.logger.WarnContext(parentCtx, "Failed to check if authorized keys report is enabled", "error", err)
-		case enabled && !isRunning:
+		case enabled && !getIsRunning():
 			runCtx, runCtxCancel = context.WithCancelCause(parentCtx)
-			mu.Lock()
-			isRunning = true
-			mu.Unlock()
+			setIsRunning(true)
 			wg.Add(1)
 			go runRoutine(runCtx, runCtxCancel)
-		case !enabled && isRunning:
+		case !enabled && getIsRunning():
 			runCtxCancel(errShutdown)
 			// Wait for the runner to stop before checking if the monitor is enabled again.
 			wg.Wait()
diff --git a/lib/secretsscanner/authorizedkeys/supervisor_test.go b/lib/secretsscanner/authorizedkeys/supervisor_test.go
index 4becb464bb567..93399ec1fdc5b 100644
--- a/lib/secretsscanner/authorizedkeys/supervisor_test.go
+++ b/lib/secretsscanner/authorizedkeys/supervisor_test.go
@@ -103,7 +103,7 @@ func TestSupervisorRunner(t *testing.T) {
 			mu.Lock()
 			defer mu.Unlock()
 			return !running
-		}, 1000*time.Millisecond, 10*time.Millisecond, "expected runner to re-stop, but it did not")
+		}, 100*time.Millisecond, 10*time.Millisecond, "expected runner to re-stop, but it did not")
 
 		// Cancel the context to stop the supervisor
 		cancel()