From 64bf633b0f6eb25c85141008c07fcec8077e262a Mon Sep 17 00:00:00 2001
From: Batuhan Taskaya <isidentical@gmail.com>
Date: Thu, 2 Sep 2021 12:48:34 +0300
Subject: [PATCH] ssh: support multiple keys (#6519)

---
 dvc/fs/ssh.py             | 19 +++++++++++--------
 tests/unit/fs/test_ssh.py | 18 ++++++++++++++++++
 2 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/dvc/fs/ssh.py b/dvc/fs/ssh.py
index 491c06ad8..23f15f888 100644
--- a/dvc/fs/ssh.py
+++ b/dvc/fs/ssh.py
@@ -2,7 +2,7 @@ import getpass
 import os.path
 import threading
 
-from funcy import cached_property, first, memoize, silent, wrap_prop, wrap_with
+from funcy import cached_property, memoize, silent, wrap_prop, wrap_with
 
 from dvc import prompt
 from dvc.scheme import Schemes
@@ -68,14 +68,17 @@ class SSHFileSystem(CallbackMixin, FSSpecWrapper):
 
         login_info["password"] = config.get("password")
 
-        if user_ssh_config.get("IdentityFile"):
-            config.setdefault(
-                "keyfile", first(user_ssh_config.get("IdentityFile"))
-            )
+        raw_keys = []
+        if config.get("keyfile"):
+            raw_keys.append(config.get("keyfile"))
+        elif user_ssh_config.get("IdentityFile"):
+            raw_keys.extend(user_ssh_config.get("IdentityFile"))
+
+        if raw_keys:
+            login_info["client_keys"] = [
+                os.path.expanduser(key) for key in raw_keys
+            ]
 
-        keyfile = config.get("keyfile")
-        if keyfile:
-            login_info["client_keys"] = [os.path.expanduser(keyfile)]
         login_info["timeout"] = config.get("timeout", _SSH_TIMEOUT)
 
         # These two settings fine tune the asyncssh to use the
diff --git a/tests/unit/fs/test_ssh.py b/tests/unit/fs/test_ssh.py
index 474a589de..da5a544e6 100644
--- a/tests/unit/fs/test_ssh.py
+++ b/tests/unit/fs/test_ssh.py
@@ -135,6 +135,24 @@ def test_ssh_keyfile(mock_file, config, expected_keyfile):
     assert fs.fs_args.get("client_keys") == expected_keyfile
 
 
+mock_ssh_multi_key_config = """
+IdentityFile file_1
+
+Host example.com
+    IdentityFile file_2
+"""
+
+
+@patch(
+    "builtins.open",
+    new_callable=mock_open,
+    read_data=mock_ssh_multi_key_config,
+)
+def test_ssh_multi_identity_files(mock_file):
+    fs = SSHFileSystem(host="example.com")
+    assert fs.fs_args.get("client_keys") == ["file_1", "file_2"]
+
+
 @pytest.mark.parametrize(
     "config,expected_gss_auth",
     [