From 1152d283fed8cc35e717c7e7146d19ecf8a99e28 Mon Sep 17 00:00:00 2001
From: Roeniss Moon <roeniss2@gmail.com>
Date: Sat, 22 Mar 2025 00:55:11 +0900
Subject: [PATCH] feat(ssh): Load hosts in other ssh config files

If the ~/.ssh/config contains `Include` directives, it'd be better
to parse the target config files if possible. This commit implements that.
The included config can also have Include directives, but it's not parsed
because of possibility of infinite loop, for now.
---
 plugins/ssh/README.md      |  5 ++--
 plugins/ssh/ssh.plugin.zsh | 57 ++++++++++++++++++++++++++++----------
 2 files changed, 46 insertions(+), 16 deletions(-)

diff --git a/plugins/ssh/README.md b/plugins/ssh/README.md
index 3dd32ec76..bbe3e6c65 100644
--- a/plugins/ssh/README.md
+++ b/plugins/ssh/README.md
@@ -1,7 +1,8 @@
 # ssh plugin
 
-This plugin provides host completion based off of your `~/.ssh/config` file, and adds
-some utility functions to work with SSH keys.
+This plugin provides host completion based off of your `~/.ssh/config` file,
+and adds some utility functions to work with SSH keys. If the `~/.ssh/config`
+contains `Include` directives, the plugin also parse hosts in related files.
 
 To use it, add `ssh` to the plugins array in your zshrc file:
 
diff --git a/plugins/ssh/ssh.plugin.zsh b/plugins/ssh/ssh.plugin.zsh
index b5b050536..3b1b28421 100644
--- a/plugins/ssh/ssh.plugin.zsh
+++ b/plugins/ssh/ssh.plugin.zsh
@@ -2,20 +2,49 @@
 # Take all host sections in .ssh/config and offer them for
 # completion as hosts (e.g. for ssh, rsync, scp and the like)
 # Filter out wildcard host sections.
-_ssh_configfile="$HOME/.ssh/config"
-if [[ -f "$_ssh_configfile" ]]; then
-  _ssh_hosts=($(
-    egrep '^Host.*' "$_ssh_configfile" |\
-    awk '{for (i=2; i<=NF; i++) print $i}' |\
-    sort |\
-    uniq |\
-    grep -v '^*' |\
-    sed -e 's/\.*\*$//'
-  ))
-  zstyle ':completion:*:hosts' hosts $_ssh_hosts
-  unset _ssh_hosts
-fi
-unset _ssh_configfile
+# If the .ssh/config has Include directives, load them too.
+function _load_ssh_hosts {
+  local conf="$1"
+  if [[ -f "$conf" ]]; then
+    local _ssh_hosts=($(
+      egrep '^Host\ .*' "$conf" |\
+      awk '{for (i=2; i<=NF; i++) print $i}' |\
+      sort |\
+      uniq |\
+      grep -v '^*' |\
+      sed -e 's/\.*\*$//'
+    ))
+    echo "${_ssh_hosts[@]}"
+  fi
+}
+
+# XXX: We could make it recursive but won't implement for now
+function _find_include_files {
+  local conf="$1"
+  if [[ -f "$conf" ]]; then
+    egrep '^Include\ .*' "$conf" |\
+    awk '{for (i=2; i<=NF; i++) print $i}'
+  fi
+}
+
+all_hosts=($(_load_ssh_hosts "$HOME/.ssh/config"))
+
+_find_include_files "$HOME/.ssh/config" | while read include_file; do
+  # omz doesn't know "~" directory
+  if [[ "$include_file" == ~* ]]; then
+    include_file="${HOME}${include_file:1}"
+  fi
+  if [[ -f "$include_file" ]]; then
+    hosts=($(_load_ssh_hosts "$include_file"))
+    all_hosts+=("${hosts[@]}")
+  fi
+done
+
+zstyle ':completion:*:hosts' hosts "${all_hosts[@]}"
+
+unset -f _load_ssh_hosts
+unset -f _find_include_files
+unset all_hosts
 
 ############################################################
 # Remove host key from known hosts based on a host section