SHort circuit api check if auto keys exist

pkg/cmd/codespace/ssh.go

@@ -245,6 +245,11 @@ func useAutomaticSSHKeys(
 		}
 	}
 
+	if automaticKeySSHKeysExist(sshContext) {
+		// If the automatic keys already exist, just use them
+		return true, nil
+	}
+
 	// If there is a public key uploaded which matches a configured
 	// private key, there is no need for automatic key generation
 	hasPublicKey, err := hasUploadedPublicKeyForConfig(ctx, sshContext, apiClient, customConfigPath, opts.profile)
@@ -252,6 +257,26 @@ func useAutomaticSSHKeys(
 	return !hasPublicKey, err
 }
 
+func automaticKeySSHKeysExist(sshContext ssh.Context) bool {
+	publicKeys, err := sshContext.LocalPublicKeys()
+	if err != nil {
+		return false
+	}
+
+	for _, publicKey := range publicKeys {
+		if filepath.Base(publicKey) != automaticPrivateKeyName+".pub" {
+			continue
+		}
+
+		privateKey := strings.TrimSuffix(publicKey, ".pub")
+
+		_, err := os.Stat(privateKey)
+		return err == nil
+	}
+
+	return false
+}
+
 func setupAutomaticSSHKeys(sshContext ssh.Context) (*ssh.KeyPair, error) {
 	keyPair := checkAndUpdateOldKeyPair(sshContext)
 	if keyPair != nil {

pkg/cmd/codespace/ssh_test.go

@@ -123,7 +123,48 @@ func TestAutomaticSSHKeyPairs(t *testing.T) {
 			}
 		}
 	}
+}
 
+func TestDontUseAutomaticSSHKeysWithCustomPrivateKey(t *testing.T) {
+	args := []string{"-i", "private-key"}
+	result, err := useAutomaticSSHKeys(context.Background(), ssh.Context{}, nil, args, sshOptions{})
+
+	if err != nil {
+		t.Errorf("Unexpected error from useAutomaticSSHKeys: %v", err)
+	}
+
+	if result != false {
+		t.Errorf("Want useAutomaticSSHKeys to be false, got true")
+	}
+}
+
+func TestUseAutomaticSSHKeysWithAutoKeysExist(t *testing.T) {
+	dir := t.TempDir()
+
+	f, err := os.Create(path.Join(dir, automaticPrivateKeyName))
+	if err != nil {
+		t.Errorf("Failed to create test private key: %v", err)
+	}
+	f.Close()
+
+	f, err = os.Create(path.Join(dir, automaticPrivateKeyName+".pub"))
+	if err != nil {
+		t.Errorf("Failed to create test public key: %v", err)
+	}
+	f.Close()
+
+	sshContext := ssh.Context{
+		ConfigDir: dir,
+	}
+	result, err := useAutomaticSSHKeys(context.Background(), sshContext, nil, nil, sshOptions{})
+
+	if err != nil {
+		t.Errorf("Unexpected error from useAutomaticSSHKeys: %v", err)
+	}
+
+	if result != true {
+		t.Errorf("Want useAutomaticSSHKeys to be true, got false")
+	}
 }
 
 func TestHasUploadedPublicKeyForConfig(t *testing.T) {