Fix open two connections with changed cache key (#392)

* Fix open two connections with changed cache key

`args.to_s` is used for cache key of connection pooling.
`ssh_options` in `args` is changed while creating connections in
`connection_factory.call(*args)`.
Updating cache key with changed `args` prevents cache miss.

* Simplify test and workaround for mutated conn args

* Fix default options and updating key

:password_prompt and :logger options are added in net-ssh
`Net::SSH.assign_defaults`.
Set default options early for ConnectionPool cache key.
The key in Cache object should be updated for comparing next time.

CHANGELOG.md

@@ -8,6 +8,7 @@ appear at the top.
   * Your contribution here!
   * [#390](https://github.com/capistrano/sshkit/pull/390): Properly wrap Ruby StandardError w/ add'l context - [@mattbrictson](https://github.com/mattbrictson)
   * [#372](https://github.com/capistrano/sshkit/pull/372): Use cp_r in local backend with recursive option - [@okuramasafumi](https://github.com/okuramasafumi)
+  * [#392](https://github.com/capistrano/sshkit/pull/392): Fix open two connections with changed cache key - [@shirosaki](https://github.com/shirosaki)
 
 ## [1.12.0][] (2017-02-10)
 

lib/sshkit/backends/connection_pool.rb

@@ -61,6 +61,9 @@ def with(connection_factory, *args)
     yield(conn)
   ensure
     cache.push(conn) unless conn.nil?
+    # Sometimes the args mutate as a result of opening a connection. In this
+    # case we need to update the cache key to match the new args.
+    update_key_if_args_changed(cache, args)
   end
 
   # Immediately remove all cached connections, without closing them. This only
@@ -84,14 +87,18 @@ def close_connections
 
   private
 
+  def cache_key_for_connection_args(args)
+    args.to_s
+  end
+
   def cache_enabled?
     idle_timeout && idle_timeout > 0
   end
 
   # Look up a Cache that matches the given connection arguments.
   def find_cache(args)
     if cache_enabled?
-      key = args.to_s
+      key = cache_key_for_connection_args(args)
       caches[key] || thread_safe_find_or_create_cache(key)
     else
       NilCache.new(method(:silently_close_connection))
@@ -103,11 +110,22 @@ def find_cache(args)
   def thread_safe_find_or_create_cache(key)
     caches.synchronize do
       caches[key] ||= begin
-        Cache.new(idle_timeout, method(:silently_close_connection_later))
+        Cache.new(key, idle_timeout, method(:silently_close_connection_later))
       end
     end
   end
 
+  # Update cache key with changed args to prevent cache miss
+  def update_key_if_args_changed(cache, args)
+    new_key = cache_key_for_connection_args(args)
+    return if cache.same_key?(new_key)
+
+    caches.synchronize do
+      caches[new_key] = caches.delete(cache.key)
+      cache.key = new_key
+    end
+  end
+
   # Loops indefinitely to close connections and to find abandoned connections
   # that need to be closed.
   def run_eviction_loop

lib/sshkit/backends/connection_pool/cache.rb

@@ -1,7 +1,10 @@
 # A Cache holds connections for a given key. Each connection is stored along
 # with an expiration time so that its idle duration can be measured.
 class SSHKit::Backend::ConnectionPool::Cache
-  def initialize(idle_timeout, closer)
+  attr_accessor :key
+
+  def initialize(key, idle_timeout, closer)
+    @key = key
     @connections = []
     @connections.extend(MonitorMixin)
     @idle_timeout = idle_timeout
@@ -53,6 +56,10 @@ def clear
     end
   end
 
+  def same_key?(other_key)
+    key == other_key
+  end
+
   protected
 
   attr_reader :connections, :idle_timeout, :closer

lib/sshkit/backends/connection_pool/nil_cache.rb

@@ -8,4 +8,8 @@ def pop
   def push(conn)
     closer.call(conn)
   end
+
+  def same_key?(_key)
+    true
+  end
 end

lib/sshkit/backends/netssh.rb

@@ -36,12 +36,20 @@ def ssh_options
         if Net::SSH::VALID_OPTIONS.include?(:known_hosts)
           def default_options
             @default_options ||= {known_hosts: SSHKit::Backend::Netssh::KnownHosts.new}
+            assign_defaults
           end
         else
           def default_options
             @default_options ||= {}
+            assign_defaults
           end
         end
+
+        # Set default options early for ConnectionPool cache key
+        def assign_defaults
+          Net::SSH.assign_defaults(@default_options)
+          @default_options
+        end
       end
 
       def upload!(local, remote, options = {})

test/functional/backends/test_netssh.rb

@@ -59,7 +59,7 @@ def test_ssh_option_merge
           capture(:uname)
           host_ssh_options = host.ssh_options
         end.run
-        assert_equal [:forward_agent, :paranoid, :known_hosts].sort, host_ssh_options.keys.sort
+        assert_equal [:forward_agent, :paranoid, :known_hosts, :logger, :password_prompt].sort, host_ssh_options.keys.sort
         assert_equal false, host_ssh_options[:forward_agent]
         assert_equal true, host_ssh_options[:paranoid]
         assert_instance_of SSHKit::Backend::Netssh::KnownHosts, host_ssh_options[:known_hosts]

test/unit/backends/test_connection_pool.rb

@@ -134,6 +134,15 @@ def test_close_connections
           pool.close_connections
         end
       end
+
+      def test_connections_with_changed_args_is_reused
+        options = { known_hosts: "foo" }
+        connect_change_options = ->(*args) { args.last[:known_hosts] = "bar"; Object.new }
+        conn1 = pool.with(connect_change_options, "arg", options) { |c| c }
+        conn2 = pool.with(connect_change_options, "arg", options) { |c| c }
+
+        assert_equal conn1, conn2
+      end
     end
   end
 end