From 13a5397cc8f3f16735bedc990dd91a9fe8e215df Mon Sep 17 00:00:00 2001
From: Tiago Silva <tiago.silva@goteleport.com>
Date: Thu, 17 Oct 2024 11:56:28 +0100
Subject: [PATCH] [scankeys] support context cancellation to interrupt keys
 scan (#47657)

This PR allow program cancellation propagation when CTRL+C is invoked during `tsh scan keys`.

Signed-off-by: Tiago Silva <tiago.silva@goteleport.com>
---
 lib/secretsscanner/scanner/scan.go |  4 ++
 tool/tsh/common/scan.go            | 17 ++++++--
 tool/tsh/common/scan_test.go       | 69 ++++++++++++++++++++++++++++++
 3 files changed, 87 insertions(+), 3 deletions(-)
 create mode 100644 tool/tsh/common/scan_test.go

diff --git a/lib/secretsscanner/scanner/scan.go b/lib/secretsscanner/scanner/scan.go
index 664b6fcef8e67..79de0b89f0552 100644
--- a/lib/secretsscanner/scanner/scan.go
+++ b/lib/secretsscanner/scanner/scan.go
@@ -113,6 +113,10 @@ func (s *Scanner) findPrivateKeys(ctx context.Context, root, deviceID string, pr
 	logger := s.log.With("root", root)
 
 	err := filepath.WalkDir(root, func(path string, info fs.DirEntry, err error) error {
+		// check if the context is done before processing the file.
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
 		if err != nil {
 			logger.DebugContext(ctx, "error walking directory", "path", path, "error", err)
 			return fs.SkipDir
diff --git a/tool/tsh/common/scan.go b/tool/tsh/common/scan.go
index 27bd644b80cda..606056adfb3b9 100644
--- a/tool/tsh/common/scan.go
+++ b/tool/tsh/common/scan.go
@@ -91,11 +91,12 @@ func (c *scanKeysCommand) run(cf *CLIConf) error {
 		return trace.Wrap(err, "device not enrolled")
 	}
 
-	fmt.Printf("Device trust credentials found.\nScanning %s.\n", strings.Join(c.dirs, ", "))
+	dirs := splitCommaSeparatedSlice(c.dirs)
+	fmt.Printf("Device trust credentials found.\nScanning %s.\n", strings.Join(dirs, ", "))
 
 	scanner, err := secretsscanner.New(secretsscanner.Config{
-		Dirs:      c.dirs,
-		SkipPaths: c.skipPaths,
+		Dirs:      dirs,
+		SkipPaths: splitCommaSeparatedSlice(c.skipPaths),
 		Log:       slog.Default(),
 	})
 	if err != nil {
@@ -168,3 +169,13 @@ func collectPrivateKeys(privateKeys []secretsscanner.SSHPrivateKey) []*accessgra
 	}
 	return keys
 }
+
+func splitCommaSeparatedSlice(s []string) []string {
+	var result []string
+	for _, entry := range s {
+		for _, split := range strings.Split(entry, ",") {
+			result = append(result, strings.TrimSpace(split))
+		}
+	}
+	return result
+}
diff --git a/tool/tsh/common/scan_test.go b/tool/tsh/common/scan_test.go
new file mode 100644
index 0000000000000..01fcc64c14abe
--- /dev/null
+++ b/tool/tsh/common/scan_test.go
@@ -0,0 +1,69 @@
+/*
+ * Teleport
+ * Copyright (C) 2024  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package common
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func Test_splitCommaSeparatedSlice(t *testing.T) {
+
+	tests := []struct {
+		name string
+		args []string
+		want []string
+	}{
+		{
+			name: "empty",
+		},
+		{
+			name: "single dir",
+			args: []string{"dir1"},
+			want: []string{"dir1"},
+		},
+		{
+			name: "multi dir",
+			args: []string{"dir1", "dir2"},
+			want: []string{"dir1", "dir2"},
+		},
+		{
+			name: "multi dir comma separated",
+			args: []string{"dir1,dir2"},
+			want: []string{"dir1", "dir2"},
+		},
+		{
+			name: "multi dir comma separated spaces",
+			args: []string{"dir1, dir2"},
+			want: []string{"dir1", "dir2"},
+		},
+		{
+			name: "multi dir comma separated spaces",
+			args: []string{"dir1, dir2", "dir3"},
+			want: []string{"dir1", "dir2", "dir3"},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := splitCommaSeparatedSlice(tt.args)
+			require.Equal(t, tt.want, got)
+		})
+	}
+}