lib/knife-solo/ssh_connection.rb in knife-solo-0.4.1 vs lib/knife-solo/ssh_connection.rb in knife-solo-0.4.2

@@ -30,11 +30,19 @@
       @password_hook = sudo_password_hook
     end
 
     attr_reader :host, :user, :connection_options
 
-    def session
-      @session ||= Net::SSH.start(host, user, connection_options)
+    def session(&block)
+      @session ||= begin
+        if connection_options[:gateway]
+          co = connection_options
+          gw_user,gw =  co.delete(:gateway).split '@'
+          Net::SSH::Gateway.new(gw, gw_user).ssh(host, user, co, &block)
+        else
+          Net::SSH.start(host, user, connection_options, &block)
+        end
+      end
     end
 
     def password
       @password ||= @password_hook.call
     end