Merge pull request #174 from jkeiser/remote_callback

Add callback to directly find out remote forwarding response

2 files changed, 128 insertions, 39 deletions

diff --git a/lib/net/ssh/service/forward.rb b/lib/net/ssh/service/forward.rb
index 0d5d55d..b6349c0 100644
--- a/lib/net/ssh/service/forward.rb
+++ b/lib/net/ssh/service/forward.rb
@@ -88,13 +88,13 @@ module Net; module SSH; module Service
         end
 
         prepare_client(client, channel, :local)
-  
+
         channel.on_open_failed do |ch, code, description|
           channel.error { "could not establish direct channel: #{description} (#{code})" }
           channel[:socket].close
         end
       end
-      
+
       local_port
     end
 
@@ -131,22 +131,52 @@ module Net; module SSH; module Service
     # the port number. The assigned port will show up in the # #active_remotes
     # list.
     #
+    # remote_host is interpreted by the server per RFC 4254, which has these
+    # special values:
+    #
+    # - "" means that connections are to be accepted on all protocol
+    #   families supported by the SSH implementation.
+    # - "0.0.0.0" means to listen on all IPv4 addresses.
+    # - "::" means to listen on all IPv6 addresses.
+    # - "localhost" means to listen on all protocol families supported by
+    #   the SSH implementation on loopback addresses only ([RFC3330] and
+    #   [RFC3513]).
+    # - "127.0.0.1" and "::1" indicate listening on the loopback
+    #   interfaces for IPv4 and IPv6, respectively.
+    #
+    # You may pass a block that will be called when the the port forward
+    # request receives a response.  This block will be passed the remote_port
+    # that was actually bound to, or nil if the binding failed.  If the block
+    # returns :no_exception, the "failed binding" exception will not be thrown.
+    #
     # If you want to block until the port is active, you could do something
     # like this:
     #
-    #   old_active_remotes = ssh.forward.active_remotes
-    #   ssh.forward.remote(80, "www.google.com", 0, "0.0.0.0")
-    #   ssh.loop { !(ssh.forward.active_remotes.length > old_active_remotes.length) }
-    #   assigned_port = (ssh.forward.active_remotes - old_active_remotes).first[0]
+    #   got_remote_port = nil
+    #   remote(port, host, remote_port, remote_host) do |actual_remote_port|
+    #     got_remote_port = actual_remote_port || :error
+    #     :no_exception # will yield the exception on my own thread
+    #   end
+    #   session.loop { !got_remote_port }
+    #   if got_remote_port == :error
+    #     raise Net::SSH::Exception, "remote forwarding request failed"
+    #   end
+    #
     def remote(port, host, remote_port, remote_host="127.0.0.1")
       session.send_global_request("tcpip-forward", :string, remote_host, :long, remote_port) do |success, response|
         if success
           remote_port = response.read_long if remote_port == 0
           debug { "remote forward from remote #{remote_host}:#{remote_port} to #{host}:#{port} established" }
           @remote_forwarded_ports[[remote_port, remote_host]] = Remote.new(host, port)
+          yield remote_port, remote_host if block_given?
         else
-          error { "remote forwarding request failed" }
-          raise Net::SSH::Exception, "remote forwarding request failed"
+          instruction = if block_given?
+            yield :error
+          end
+          unless instruction == :no_exception
+            error { "remote forwarding request failed" }
+            raise Net::SSH::Exception, "remote forwarding request failed"
+          end
         end
       end
     end
@@ -183,6 +213,16 @@ module Net; module SSH; module Service
       @remote_forwarded_ports.keys
     end
 
+    # Returns all active remote forwarded ports and where they forward to. The
+    # returned value is a hash from [<forwarding port on the local host>, <local forwarding address>]
+    # to [<port on the remote host>, <remote bind address>].
+    def active_remote_destinations
+      @remote_forwarded_ports.inject({}) do |result, (remote, local)|
+        result[[local.port, local.host]] = remote
+        result
+      end
+    end
+
     # Enables SSH agent forwarding on the given channel. The forwarded agent
     # will remain active even after the channel closes--the channel is only
     # used as the transport for enabling the forwarded connection. You should
@@ -216,7 +256,7 @@ module Net; module SSH; module Service
     end
 
     private
-      
+
       # Perform setup operations that are common to all forwarded channels.
       # +client+ is a socket, +channel+ is the channel that was just created,
       # and +type+ is an arbitrary string describing the type of the channel.
@@ -228,11 +268,11 @@ module Net; module SSH; module Service
         session.listen_to(client)
         channel[:socket] = client
 
-        channel.on_data do |ch, data|  
+        channel.on_data do |ch, data|
           debug { "data:#{data.length} on #{type} forwarded channel" }
           ch[:socket].enqueue(data)
         end
-        
+
         # Handles server close on the sending side by Miklós Fazekas
         channel.on_eof do |ch|
           debug { "eof #{type} on #{type} forwarded channel" }
@@ -251,7 +291,7 @@ module Net; module SSH; module Service
             debug { "enotconn in on_eof => shallowing exception:#{e}" }
           end
         end
-        
+
         channel.on_close do |ch|
           debug { "closing #{type} forwarded channel" }
           ch[:socket].close if !client.closed?
diff --git a/test/manual/test_forward.rb b/test/manual/test_forward.rb
index 5794077..7e864e2 100644
--- a/test/manual/test_forward.rb
+++ b/test/manual/test_forward.rb
@@ -3,15 +3,15 @@
 # Tests for the following patch:
 #
 #   http://github.com/net-ssh/net-ssh/tree/portfwfix
-# 
+#
 # It fixes 3 issues, regarding closing forwarded ports:
-# 
+#
 # 1.) if client closes a forwarded connection, but the server is reading, net-ssh terminates with IOError socket closed.
 # 2.) if client force closes (RST) a forwarded connection, but server is reading, net-ssh terminates with
 # 3.) if server closes the sending side, the on_eof is not handled.
-# 
+#
 # More info:
-# 
+#
 # http://net-ssh.lighthouseapp.com/projects/36253/tickets/7
 
 require 'common'
@@ -21,26 +21,26 @@ require 'timeout'
 require 'tempfile'
 
 class TestForward < Test::Unit::TestCase
-  
+
   def localhost
     'localhost'
   end
-  
+
   def ssh_start_params
     [localhost ,ENV['USER'], {:keys => "~/.ssh/id_rsa", :verbose => :debug}]
   end
-  
+
   def start_server_sending_lot_of_data(exceptions)
     server = TCPServer.open(0)
     Thread.start do
       loop do
         Thread.start(server.accept) do |client|
           begin
-            10000.times do |i| 
+            10000.times do |i|
               client.puts "item#{i}"
             end
             client.close
-          rescue 
+          rescue
             exceptions << $!
             raise
           end
@@ -49,14 +49,14 @@ class TestForward < Test::Unit::TestCase
     end
     return server
   end
-  
+
   def start_server_closing_soon(exceptions=nil)
     server = TCPServer.open(0)
     Thread.start do
       loop do
         Thread.start(server.accept) do |client|
           begin
-            client.recv(1024) 
+            client.recv(1024)
             client.setsockopt(Socket::SOL_SOCKET, Socket::SO_LINGER, [1, 0].pack("ii"))
             client.close
           rescue
@@ -68,20 +68,20 @@ class TestForward < Test::Unit::TestCase
     end
     return server
   end
-  
+
   def test_local_ephemeral_port_should_work_correctly
     session = Net::SSH.start(*ssh_start_params)
-    
+
     assert_nothing_raised do
       assigned_port = session.forward.local(0, localhost, 22)
       assert_not_nil assigned_port
       assert_operator assigned_port, :>, 0
     end
   end
-  
+
   def test_remote_ephemeral_port_should_work_correctly
     session = Net::SSH.start(*ssh_start_params)
-    
+
     assert_nothing_raised do
       session.forward.remote(22, localhost, 0, localhost)
       session.loop { !(session.forward.active_remotes.length > 0) }
@@ -90,9 +90,58 @@ class TestForward < Test::Unit::TestCase
       assert_operator assigned_port, :>, 0
     end
   end
-  
+
+  def test_remote_callback_should_fire
+    session = Net::SSH.start(*ssh_start_params)
+
+    assert_nothing_raised do
+      got_port = nil
+      session.forward.remote(22, localhost, 0, localhost) do |port|
+        got_port = port
+      end
+      session.loop { !(session.forward.active_remotes.length > 0) }
+      assert_operator session.forward.active_remote_destinations.length, :==, 1
+      assert_operator session.forward.active_remote_destinations.keys.first, :==, [ 22, localhost ]
+      assert_operator session.forward.active_remote_destinations.values.first, :==, [ got_port, localhost ]
+      assert_operator session.forward.active_remotes.first, :==, [ got_port, localhost ]
+      assigned_port = session.forward.active_remotes.first[0]
+      assert_operator got_port, :==, assigned_port
+      assert_not_nil assigned_port
+      assert_operator assigned_port, :>, 0
+    end
+  end
+
+  def test_remote_callback_should_fire_on_error_and_still_throw_exception
+    session = Net::SSH.start(*ssh_start_params)
+
+    assert_nothing_raised do
+      session.forward.remote(22, localhost, 22, localhost) do |port|
+        assert_operator port, :==, :error
+      end
+    end
+    assert_raises(Net::SSH::Exception) do
+      session.loop { true }
+    end
+  end
+
+  def test_remote_callback_should_fire_on_error_but_not_throw_exception_if_asked_not_to
+    session = Net::SSH.start(*ssh_start_params)
+
+    assert_nothing_raised do
+      got_port = nil
+      session.forward.remote(22, localhost, 22, localhost) do |port|
+        assert_operator port, :==, :error
+        got_port = port
+        :no_exception
+      end
+      session.loop { !got_port }
+      assert_operator port, :==, :error
+      assert_operator session.forward.active_remotes.length, :==, 0
+    end
+  end
+
   def test_loop_should_not_abort_when_local_side_of_forward_is_closed
-    session = Net::SSH.start(*ssh_start_params) 
+    session = Net::SSH.start(*ssh_start_params)
     server_exc = Queue.new
     server = start_server_sending_lot_of_data(server_exc)
     remote_port = server.addr[1]
@@ -112,10 +161,10 @@ class TestForward < Test::Unit::TestCase
     session.loop(0.1) { client_done.empty? }
     assert_equal "Broken pipe", "#{server_exc.pop}" unless server_exc.empty?
   end
-  
+
   def test_loop_should_not_abort_when_local_side_of_forward_is_reset
     session = Net::SSH.start(*ssh_start_params)
-    server_exc = Queue.new    
+    server_exc = Queue.new
     server = start_server_sending_lot_of_data(server_exc)
     remote_port = server.addr[1]
     local_port = 0 # request ephemeral port
@@ -143,9 +192,9 @@ class TestForward < Test::Unit::TestCase
     yield UNIXServer.open(path)
     File.delete(path)
   end if defined?(UNIXServer)
-  
+
   def test_forward_local_unix_socket_to_remote_port
-    session = Net::SSH.start(*ssh_start_params) 
+    session = Net::SSH.start(*ssh_start_params)
     server_exc = Queue.new
     server = start_server_sending_lot_of_data(server_exc)
     remote_port = server.addr[1]
@@ -174,7 +223,7 @@ class TestForward < Test::Unit::TestCase
   end if defined?(UNIXSocket)
 
   def test_loop_should_not_abort_when_server_side_of_forward_is_closed
-    session = Net::SSH.start(*ssh_start_params)    
+    session = Net::SSH.start(*ssh_start_params)
     server = start_server_closing_soon
     remote_port = server.addr[1]
     local_port = 0 # request ephemeral port
@@ -183,18 +232,18 @@ class TestForward < Test::Unit::TestCase
     Thread.start do
       begin
         client = TCPSocket.new(localhost, local_port)
-        1.times do |i| 
+        1.times do |i|
           client.puts "item#{i}"
         end
         client.close
         sleep(0.1)
-      ensure                 
+      ensure
         client_done << true
       end
     end
     session.loop(0.1) { client_done.empty? }
   end
-  
+
   def start_server
     server = TCPServer.open(0)
     Thread.start do
@@ -206,9 +255,9 @@ class TestForward < Test::Unit::TestCase
     end
     return server
   end
-  
+
   def test_server_eof_should_be_handled
-    session = Net::SSH.start(*ssh_start_params)    
+    session = Net::SSH.start(*ssh_start_params)
     server = start_server do |client|
       client.write "This is a small message!"
       client.close