From 584f8ca37707cb3bfad237acd8fc13236bb5b98d Mon Sep 17 00:00:00 2001
From: Vadym Popov <vadym.popov@goteleport.com>
Date: Wed, 19 Feb 2025 21:14:58 -0800
Subject: [PATCH] Prevent keystore cleanup to remove bin directory (#52331)

---
 lib/client/keystore.go      | 12 +++++++-----
 lib/client/keystore_test.go | 15 ++++++++++++++-
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/lib/client/keystore.go b/lib/client/keystore.go
index ae8dca55b0168..7a929bd98cd7b 100644
--- a/lib/client/keystore.go
+++ b/lib/client/keystore.go
@@ -54,6 +54,10 @@ const (
 	// tshAzureDirName is the name of the directory containing the
 	// az cli app-specific profiles.
 	tshAzureDirName = "azure"
+
+	// tshBin is the name of the directory containing the
+	// updated binaries of client tools.
+	tshBin = "bin"
 )
 
 // KeyStore is a storage interface for client session keys and certificates.
@@ -282,13 +286,11 @@ func (fs *FSKeyStore) DeleteKeys() error {
 	if err != nil {
 		return trace.ConvertSystemError(err)
 	}
+	ignoreDirs := map[string]struct{}{tshConfigFileName: {}, tshAzureDirName: {}, tshBin: {}}
 	for _, file := range files {
-		// Don't delete 'config' and 'azure' directories.
+		// Don't delete 'config', 'azure' and 'bin' directories.
 		// TODO: this is hackish and really shouldn't be needed, but fs.KeyDir is `~/.tsh` while it probably should be `~/.tsh/keys` instead.
-		if file.IsDir() && file.Name() == tshConfigFileName {
-			continue
-		}
-		if file.IsDir() && file.Name() == tshAzureDirName {
+		if _, ok := ignoreDirs[file.Name()]; ok && file.IsDir() {
 			continue
 		}
 		if file.IsDir() {
diff --git a/lib/client/keystore_test.go b/lib/client/keystore_test.go
index 77ef869044cc3..fe708e7d7bce2 100644
--- a/lib/client/keystore_test.go
+++ b/lib/client/keystore_test.go
@@ -276,7 +276,7 @@ func TestAddKey_withoutSSHCert(t *testing.T) {
 	require.Len(t, keyCopy.DBTLSCerts, 1)
 }
 
-func TestConfigDirNotDeleted(t *testing.T) {
+func TestProtectedDirsNotDeleted(t *testing.T) {
 	t.Parallel()
 	auth := newTestAuthority(t)
 	keyStore := newTestFSKeyStore(t)
@@ -285,8 +285,21 @@ func TestConfigDirNotDeleted(t *testing.T) {
 	keyStore.AddKey(auth.makeSignedKey(t, idx, false))
 	configPath := filepath.Join(keyStore.KeyDir, "config")
 	require.NoError(t, os.Mkdir(configPath, 0700))
+
+	azurePath := filepath.Join(keyStore.KeyDir, "azure")
+	require.NoError(t, os.Mkdir(azurePath, 0700))
+
+	binPath := filepath.Join(keyStore.KeyDir, "bin")
+	require.NoError(t, os.Mkdir(binPath, 0700))
+
+	testPath := filepath.Join(keyStore.KeyDir, "test")
+	require.NoError(t, os.Mkdir(testPath, 0700))
+
 	require.NoError(t, keyStore.DeleteKeys())
 	require.DirExists(t, configPath)
+	require.DirExists(t, azurePath)
+	require.DirExists(t, binPath)
+	require.NoDirExists(t, testPath)
 
 	require.NoDirExists(t, filepath.Join(keyStore.KeyDir, "keys"))
 }