Provide an alternative implementation of Net::SSH::KnownHost in ssh_options

As discussed in capistrano#326 (comment)
Net::SSH re-parse the known_hosts files every time it needs to lookup for a known key.
This alternative implementation parse it once and for all, and cache the result.

Author: byroot

lib/sshkit/backends/netssh.rb

@@ -1,3 +1,5 @@
+require 'strscan'
+require 'mutex_m'
 require 'net/ssh'
 require 'net/scp'
 
@@ -18,13 +20,101 @@ module SSHKit
   module Backend
 
     class Netssh < Abstract
+      class KnownHostsKeys
+        include Mutex_m
+
+        def initialize(path)
+          super()
+          @path = File.expand_path(path)
+          @hosts_keys = nil
+        end
+
+        def keys_for(hostlist)
+          keys = hosts_keys || parse_file
+          hostlist.split(',').each do |host|
+            if key_list = keys[host]
+              return key_list
+            end
+          end
+          []
+        end
+
+        private
+
+        attr_reader :path
+        attr_accessor :hosts_keys
+
+        def parse_file
+          synchronize do
+            return keys if hosts_keys
+
+            return self.hosts_keys = {} unless File.readable?(path)
+
+            new_keys = {}
+            File.open(path) do |file|
+              scanner = StringScanner.new("")
+              file.each_line do |line|
+                scanner.string = line
+                hostlist, key = parse_line(scanner)
+                next unless key
+
+                hostlist.each do |host|
+                  (new_keys[host] ||= []) << key
+                end
+              end
+            end
+            return self.hosts_keys = new_keys
+          end
+        end
+
+        def parse_line(scanner)
+          scanner.skip(/\s*/)
+          return if scanner.match?(/$|#/)
+
+          hostlist = scanner.scan(/\S+/).split(',')
+          scanner.skip(/\s*/)
+          type = scanner.scan(/\S+/)
+
+          return unless Net::SSH::KnownHosts::SUPPORTED_TYPE.include?(type)
+
+          scanner.skip(/\s*/)
+          blob = scanner.rest.unpack("m*").first
+          return hostlist, Net::SSH::Buffer.new(blob).read_key
+        end
+      end
+
+      class KnownHosts
+        include Mutex_m
+
+        def initialize
+          super()
+          @files = {}
+        end
+
+        def search_for(host, options = {})
+          ::Net::SSH::KnownHosts.hostfiles(options).map do |path|
+            known_hosts_file(path).keys_for(host)
+          end.flatten
+        end
+
+        def add(*args)
+          ::Net::SSH::KnownHosts.add(*args)
+          synchronize { @files = {} }
+        end
+
+        private
+
+        def known_hosts_file(path)
+          @files[path] || synchronize { @files[path] ||= KnownHostsKeys.new(path) }
+        end
+      end
 
       class Configuration
         attr_accessor :connection_timeout, :pty
         attr_writer :ssh_options
 
         def ssh_options
-          @ssh_options || {}
+          @ssh_options ||= {known_hosts: KnownHosts.new}
         end
       end
 

test/known_hosts/github

@@ -0,0 +1 @@
+github.com ssh-rsa AAAAB3NzaC1yc2EAAAABIwAAAQEAq2A7hRGmdnm9tUDbO9IDSwBK6TbQa+PXYPCPy6rbTrTtw7PHkccKrpp0yVhp5HdEIcKr6pLlVDBfOLX9QUsyCOV0wzfjIJNlGEYsdlLJizHhbn2mUjvSAHQqZETYP81eFzLQNnPHt4EVVUh7VfDESU84KezmD5QlWpXLmvU31/yMf+Se8xhHTvKSCZIFImWwoG6mbUoWf9nzpIoaSjB+weqqUUmpaaasXVal72J+UX2B+2RPW3RcT0eOzQgqlJL3RKrTJvdsjE3JEAvGq3lGHSZXy28G3skua2SmVi/w4yCE6gbODqnTWlg7+wC604ydGXA8VJiS5ap43JXiUFFAaQ==

test/known_hosts/github_hash

@@ -0,0 +1 @@
+|1|eKp+6E0rZ3lONgsIziurXEnaIik=|rcQB/rlJMUquUyFta64KugPjX4o= ssh-rsa AAAAB3NzaC1yc2EAAAABIwAAAQEAq2A7hRGmdnm9tUDbO9IDSwBK6TbQa+PXYPCPy6rbTrTtw7PHkccKrpp0yVhp5HdEIcKr6pLlVDBfOLX9QUsyCOV0wzfjIJNlGEYsdlLJizHhbn2mUjvSAHQqZETYP81eFzLQNnPHt4EVVUh7VfDESU84KezmD5QlWpXLmvU31/yMf+Se8xhHTvKSCZIFImWwoG6mbUoWf9nzpIoaSjB+weqqUUmpaaasXVal72J+UX2B+2RPW3RcT0eOzQgqlJL3RKrTJvdsjE3JEAvGq3lGHSZXy28G3skua2SmVi/w4yCE6gbODqnTWlg7+wC604ydGXA8VJiS5ap43JXiUFFAaQ==

test/unit/backends/test_netssh.rb

@@ -1,4 +1,5 @@
 require 'helper'
+require 'tempfile'
 
 module SSHKit
   module Backend
@@ -53,6 +54,23 @@ def test_transfer_summarizer
         end
       end
 
+      def test_known_hosts_for_when_all_hosts_are_recognized
+        perform_known_hosts_test("github")
+      end
+
+      def test_known_hosts_for_when_an_host_hash_is_recognized
+        perform_known_hosts_test("github_hash")
+      end
+
+      private
+
+      def perform_known_hosts_test(hostfile)
+        source = File.join(File.dirname(__FILE__), '../../known_hosts', hostfile)
+        kh = Netssh::KnownHosts.new
+        keys = kh.search_for('github.com', user_known_hosts_file: source, global_known_hosts_file: Tempfile.new('sshkit-test').path)
+        assert_equal(1, keys.count)
+        assert_equal("ssh-rsa", keys[0].ssh_type)
+      end
     end
   end
 end