Support "IdentitiesOnly" directivemusybite-identities-only

4 files changed, 77 insertions, 30 deletions

diff --git a/lib/net/ssh.rb b/lib/net/ssh.rb
index c752c71..acc23a7 100644
--- a/lib/net/ssh.rb
+++ b/lib/net/ssh.rb
@@ -66,7 +66,7 @@ module Net
       :logger, :paranoid, :password, :port, :proxy, :rekey_blocks_limit,
       :rekey_limit, :rekey_packet_limit, :timeout, :verbose,
       :global_known_hosts_file, :user_known_hosts_file, :host_key_alias,
-      :host_name, :user, :properties, :passphrase
+      :host_name, :user, :properties, :passphrase, :keys_only
     ]
 
     # The standard means of starting a new SSH connection. When used with a
@@ -125,6 +125,10 @@ module Net
     #   and hostbased authentication
     # * :key_data => an array of strings, with each element of the array being
     #   a raw private key in PEM format.
+    # * :keys_only => set to +true+ to use only private keys from +keys+ and
+    #   +key_data+ parameters, even if ssh-agent offers more identities. This
+    #   option is intended for situations where ssh-agent offers many different
+    #   identites.
     # * :logger => the logger instance to use when logging
     # * :paranoid => either true, false, or :very, specifying how strict
     #   host-key verification should be
diff --git a/lib/net/ssh/authentication/key_manager.rb b/lib/net/ssh/authentication/key_manager.rb
index 1a2f386..c96e18c 100644
--- a/lib/net/ssh/authentication/key_manager.rb
+++ b/lib/net/ssh/authentication/key_manager.rb
@@ -90,40 +90,30 @@ module Net
         # The origin of the identities may be from files on disk or from an
         # ssh-agent. Note that identities from an ssh-agent are always listed
         # first in the array, with other identities coming after.
+        #
+        # If key manager was created with :keys_only option, any identity
+        # from ssh-agent will be ignored unless it present in key_files or
+        # key_data.
         def each_identity
+          user_identities = load_identities_from_files + load_identities_from_data
+
           if agent
             agent.identities.each do |key|
-              known_identities[key] = { :from => :agent }
-              yield key
-            end
-          end
-          
-          key_files.each do |file|
-            public_key_file = file + ".pub"
-            if File.readable?(public_key_file)
-              begin
-                key = KeyFactory.load_public_key(public_key_file)
-                known_identities[key] = { :from => :file, :file => file }
-                yield key
-              rescue Exception => e
-                error { "could not load public key file `#{public_key_file}': #{e.class} (#{e.message})" }
-              end
-            elsif File.readable?(file)
-              begin
-                private_key = KeyFactory.load_private_key(file, options[:passphrase])
-                key = private_key.send(:public_key)
-                known_identities[key] = { :from => :file, :file => file, :key => private_key }
+              corresponding_user_identity = user_identities.detect { |identity|
+                identity[:public_key].to_pem == key.to_pem
+              }
+              user_identities.delete(corresponding_user_identity) if corresponding_user_identity
+
+              if !options[:keys_only] || corresponding_user_identity
+                known_identities[key] = { :from => :agent }
                 yield key
-              rescue Exception => e
-                error { "could not load private key file `#{file}': #{e.class} (#{e.message})" }
               end
             end
           end
 
-          key_data.each do |data|
-            private_key = KeyFactory.load_data_private_key(data)
-            key = private_key.send(:public_key)
-            known_identities[key] = { :from => :key_data, :data => data, :key => private_key }
+          user_identities.each do |identity|
+            key = identity.delete(:public_key)
+            known_identities[key] = identity
             yield key
           end
 
@@ -186,8 +176,44 @@ module Net
           @use_agent = false
           nil
         end
-      end
 
+        private
+
+        # Extracts identities from user key_files, preserving their order and sources.
+        def load_identities_from_files
+          key_files.map do |file|
+            public_key_file = file + ".pub"
+            if File.readable?(public_key_file)
+              begin
+                key = KeyFactory.load_public_key(public_key_file)
+                { :public_key => key, :from => :file, :file => file }
+              rescue Exception => e
+                error { "could not load public key file `#{public_key_file}': #{e.class} (#{e.message})" }
+                nil
+              end
+            elsif File.readable?(file)
+              begin
+                private_key = KeyFactory.load_private_key(file, options[:passphrase])
+                key = private_key.send(:public_key)
+                { :public_key => key, :from => :file, :file => file, :key => private_key }
+              rescue Exception => e
+                error { "could not load private key file `#{file}': #{e.class} (#{e.message})" }
+                nil
+              end
+            end
+          end.compact
+        end
+
+        # Extraccts identities from user key_data, preserving their order and sources.
+        def load_identities_from_data
+          key_data.map do |data|
+            private_key = KeyFactory.load_data_private_key(data)
+            key = private_key.send(:public_key)
+            { :public_key => key, :from => :key_data, :data => data, :key => private_key }
+          end
+        end
+
+      end
     end
   end
 end
diff --git a/lib/net/ssh/config.rb b/lib/net/ssh/config.rb
index 6eb9c2e..213e1b6 100644
--- a/lib/net/ssh/config.rb
+++ b/lib/net/ssh/config.rb
@@ -19,6 +19,7 @@ module Net; module SSH
   # * HostKeyAlias => :host_key_alias
   # * HostName => :host_name
   # * IdentityFile => maps to the :keys option
+  # * IdentitiesOnly => :keys_only
   # * Macs => maps to the :hmac option
   # * PasswordAuthentication => maps to the :auth_methods option
   # * Port => :port
@@ -128,6 +129,8 @@ module Net; module SSH
             hash[:timeout] = value
           when 'forwardagent' then
             hash[:forward_agent] = value
+          when 'identitiesonly' then
+            hash[:keys_only] = value
           when 'globalknownhostsfile'
             hash[:global_known_hosts_file] = value
           when 'hostbasedauthentication' then
diff --git a/test/authentication/test_key_manager.rb b/test/authentication/test_key_manager.rb
index dc8beda..ad171ac 100644
--- a/test/authentication/test_key_manager.rb
+++ b/test/authentication/test_key_manager.rb
@@ -59,6 +59,20 @@ module Authentication
       assert_equal({:from => :agent}, manager.known_identities[dsa])
     end
 
+    def test_only_identities_with_key_files_should_load_from_agent_of_keys_only_set
+      manager(:keys_only => true).stubs(:agent).returns(agent)
+
+      stub_file_key "/first", rsa
+
+      identities = []
+      manager.each_identity { |identity| identities << identity }
+
+      assert_equal 1, identities.length
+      assert_equal rsa.to_blob, identities.first.to_blob
+
+      assert_equal({:from => :agent}, manager.known_identities[rsa])
+    end
+
     def test_sign_with_agent_originated_key_should_request_signature_from_agent
       manager.stubs(:agent).returns(agent)
       manager.each_identity { |identity| } # preload the known_identities
@@ -96,8 +110,8 @@ module Authentication
         @agent ||= stub("agent", :identities => [rsa, dsa])
       end
 
-      def manager
-        @manager ||= Net::SSH::Authentication::KeyManager.new(nil)
+      def manager(options = {})
+        @manager ||= Net::SSH::Authentication::KeyManager.new(nil, options)
       end
 
   end