From d410830f7f15199bbe31db9a7bbcfe780c0572b2 Mon Sep 17 00:00:00 2001
From: Caleb Brose <5447118+cmbrose@users.noreply.github.com>
Date: Tue, 9 Aug 2022 11:11:11 -0500
Subject: [PATCH] Handle comment in local ssh key

---
 pkg/cmd/codespace/ssh.go      | 19 +++++++++++--------
 pkg/cmd/codespace/ssh_test.go | 31 +++++++++++++++++++------------
 2 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/pkg/cmd/codespace/ssh.go b/pkg/cmd/codespace/ssh.go
index 1c2360d83..f5b9dec9f 100644
--- a/pkg/cmd/codespace/ssh.go
+++ b/pkg/cmd/codespace/ssh.go
@@ -360,7 +360,7 @@ func hasUploadedPublicKeyForConfig(
 		return false, fmt.Errorf("getting local ssh keys: %w", err)
 	}
 
-	var configuredPublicKeys []string
+	configuredPublicKeys := make(map[string]bool)
 	for _, privateKeyPath := range configuredPrivateKeyPaths {
 		publicKeyPath := privateKeyPath + ".pub"
 
@@ -371,7 +371,15 @@ func hasUploadedPublicKeyForConfig(
 			continue
 		}
 
-		configuredPublicKeys = append(configuredPublicKeys, string(publicKeyContent))
+		parts := strings.SplitN(string(publicKeyContent), " ", 3)
+		if len(parts) < 2 {
+			// Unexpected format, skip it
+			continue
+		}
+
+		publicKeyWithoutComment := strings.Join(parts[:2], " ")
+
+		configuredPublicKeys[publicKeyWithoutComment] = true
 	}
 
 	if len(configuredPublicKeys) == 0 {
@@ -379,11 +387,6 @@ func hasUploadedPublicKeyForConfig(
 		return false, nil
 	}
 
-	publicKeyMap := make(map[string]bool)
-	for _, publicKey := range configuredPublicKeys {
-		publicKeyMap[publicKey] = true
-	}
-
 	user, err := apiClient.GetUser(ctx)
 	if err != nil {
 		return false, fmt.Errorf("fetching user account: %w", err)
@@ -399,7 +402,7 @@ func hasUploadedPublicKeyForConfig(
 	}
 
 	for _, uploadedPublicKey := range uploadedPublicKeys {
-		if publicKeyMap[uploadedPublicKey] {
+		if configuredPublicKeys[uploadedPublicKey] {
 			return true, nil
 		}
 	}
diff --git a/pkg/cmd/codespace/ssh_test.go b/pkg/cmd/codespace/ssh_test.go
index b9e9f9596..7c0a02b47 100644
--- a/pkg/cmd/codespace/ssh_test.go
+++ b/pkg/cmd/codespace/ssh_test.go
@@ -185,40 +185,47 @@ func TestHasUploadedPublicKeyForConfig(t *testing.T) {
 		},
 		{
 			// Has API keys, but no local keys
-			apiAuthorizedPublicKeys: []string{"test-key"},
+			apiAuthorizedPublicKeys: []string{"ssh-rsa test-key"},
 			wantResult:              false,
 		},
 		{
 			// No API keys, but has local keys
-			localKeyPairs: []testLocalKeyPair{{"keyfile", "test-key"}},
+			localKeyPairs: []testLocalKeyPair{{"keyfile", "ssh-rsa test-key"}},
 			wantResult:    false,
 		},
 		{
 			// API keys and local keys, but not matching
-			apiAuthorizedPublicKeys: []string{"test-api-key"},
-			localKeyPairs:           []testLocalKeyPair{{"keyfile", "test-local-key"}},
+			apiAuthorizedPublicKeys: []string{"ssh-rsa test-api-key"},
+			localKeyPairs:           []testLocalKeyPair{{"keyfile", "ssh-rsa test-local-key"}},
 			wantResult:              false,
 		},
 
 		// Successful tests
 		{
-			apiAuthorizedPublicKeys: []string{"test-key"},
-			localKeyPairs:           []testLocalKeyPair{{"keyfile", "test-key"}},
+			apiAuthorizedPublicKeys: []string{"ssh-rsa test-key"},
+			localKeyPairs:           []testLocalKeyPair{{"keyfile", "ssh-rsa test-key"}},
 			wantResult:              true,
 		},
 		{
-			apiAuthorizedPublicKeys: []string{"test-key-1", "test-key-2"},
-			localKeyPairs:           []testLocalKeyPair{{"keyfile1", "test-key-1"}},
+			apiAuthorizedPublicKeys: []string{"ssh-rsa test-key-1", "ssh-rsa test-key-2"},
+			localKeyPairs:           []testLocalKeyPair{{"keyfile1", "ssh-rsa test-key-1"}},
 			wantResult:              true,
 		},
 		{
-			apiAuthorizedPublicKeys: []string{"test-key-1"},
-			localKeyPairs:           []testLocalKeyPair{{"keyfile1", "test-key-1"}, {"keyfile2", "test-key-2"}},
+			apiAuthorizedPublicKeys: []string{"ssh-rsa test-key-1"},
+			localKeyPairs:           []testLocalKeyPair{{"keyfile1", "ssh-rsa test-key-1"}, {"keyfile2", "ssh-rsa test-key-2"}},
 			wantResult:              true,
 		},
 		{
-			apiAuthorizedPublicKeys: []string{"test-key-1", "test-key-2"},
-			localKeyPairs:           []testLocalKeyPair{{"keyfile3", "test-key-3"}, {"keyfile2", "test-key-2"}},
+			apiAuthorizedPublicKeys: []string{"ssh-rsa test-key-1", "ssh-rsa test-key-2"},
+			localKeyPairs:           []testLocalKeyPair{{"keyfile3", "ssh-rsa test-key-3"}, {"keyfile2", "ssh-rsa test-key-2"}},
+			wantResult:              true,
+		},
+
+		// Extra case - local key contain comments
+		{
+			apiAuthorizedPublicKeys: []string{"ssh-rsa test-key-1", "ssh-rsa test-key"},
+			localKeyPairs:           []testLocalKeyPair{{"keyfile3", "ssh-rsa test-key a comment on the key"}},
 			wantResult:              true,
 		},
 	}