Merge branch 'portfwfix' of http://github.com/mfazekas/net-ssh into p…

…ortfwfix

endtoendtests/common.rb

@@ -0,0 +1,5 @@
+$LOAD_PATH.unshift "#{File.dirname(__FILE__)}/../lib"
+require 'rubygems'
+gem "test-unit" # http://rubyforge.org/pipermail/test-unit-tracker/2009-July/000075.html
+require 'test/unit'
+require 'mocha'

endtoendtests/test_forward.rb

@@ -0,0 +1,176 @@
+require 'common'
+require 'net/ssh/buffer'
+require 'net/ssh'
+require 'timeout'
+
+# keyless ssh setup
+#
+# cat ~/.ssh/id_rsa.pub  > ~/.ssh/authorized_keys
+# to test:
+# ssh localhost
+#
+
+class TestForward < Test::Unit::TestCase
+
+  def localhost
+    'localhost'
+  end
+
+  def ssh_start_params
+    [localhost ,ENV['USER']]
+  end
+
+  def find_free_port 
+    server = TCPServer.open(0)
+    server.setsockopt(Socket::SOL_SOCKET,Socket::SO_REUSEADDR,true)
+    port = server.addr[1]
+    server.close
+    port
+  end
+
+  def start_server_sending_lot_of_data(exceptions=nil)
+    server = TCPServer.open(0)
+    Thread.start do
+      loop do
+        Thread.start(server.accept) do |client|
+          begin
+            10000.times do |i| 
+              client.puts "item#{i}"
+            end
+            client.close
+          rescue 
+            exceptions << $!
+            raise
+          end
+        end
+      end
+    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.setsockopt(Socket::SOL_SOCKET, Socket::SO_LINGER, [1, 0].pack("ii"))
+            client.close
+          rescue
+            exceptions <<  $!
+            raise
+          end
+        end
+      end
+    end
+    return server
+  end
+
+  def test_loop_should_not_abort_when_local_side_of_forward_is_closed
+    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]
+    local_port = find_free_port
+    session.forward.local(local_port, localhost, remote_port)
+    client_done = Queue.new
+    Thread.start do
+      begin
+        client = TCPSocket.new(localhost, local_port)
+        client.recv(1024)
+        client.close
+        sleep(0.2)
+      ensure
+        client_done << true
+      end
+    end
+    session.loop(0.1) { client_done.empty? }
+    assert_equal "Broken pipe", "#{server_exc.pop}"
+  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 = start_server_sending_lot_of_data(server_exc)
+    remote_port = server.addr[1]
+    local_port = find_free_port
+    session.forward.local(local_port, localhost, remote_port)
+    client_done = Queue.new
+    Thread.start do
+      begin
+        client = TCPSocket.new(localhost, local_port)
+        client.recv(1024)
+        client.setsockopt(Socket::SOL_SOCKET, Socket::SO_LINGER, [1, 0].pack("ii"))
+        client.close
+        sleep(0.1)
+      ensure
+        client_done << true
+      end
+    end
+    session.loop(0.1) { client_done.empty? }
+    assert_equal "Broken pipe", "#{server_exc.pop}"
+  end
+
+  def test_loop_should_not_abort_when_server_side_of_forward_is_closed
+    session = Net::SSH.start(*ssh_start_params)    
+    server = start_server_closing_soon
+    remote_port = server.addr[1]
+    local_port = find_free_port
+    session.forward.local(local_port, localhost, remote_port)
+    client_done = Queue.new
+    Thread.start do
+      begin
+        client = TCPSocket.new(localhost, local_port)
+        1.times do |i| 
+          client.puts "item#{i}"
+        end
+        client.close
+        sleep(0.1)
+      ensure                 
+        client_done << true
+      end
+    end
+    session.loop(0.1) { client_done.empty? }
+  end
+
+  def start_server
+    server = TCPServer.open(0)
+    Thread.start do
+      loop do
+        Thread.start(server.accept) do |client|
+          yield(client)
+        end
+      end
+    end
+    return server
+  end
+
+  def test_server_eof_should_be_handled
+    session = Net::SSH.start(*ssh_start_params)    
+    server = start_server do |client|
+      client.write "This is a small message!"
+      client.close
+    end
+    client_done = Queue.new
+    client_exception = Queue.new
+    client_data = Queue.new
+    remote_port = server.addr[1]
+    local_port = find_free_port
+    session.forward.local(local_port, localhost, remote_port)
+    Thread.start do
+      begin
+        client = TCPSocket.new(localhost, local_port)
+        data = client.read(4096)
+        client.close
+        client_done << data
+      rescue
+        client_done << $!
+      end
+    end
+    timeout(5) do
+      session.loop(0.1) { client_done.empty? }
+      assert_equal "This is a small message!", client_done.pop
+    end
+  end
+end

lib/net/ssh/service/forward.rb

@@ -193,21 +193,63 @@ def agent(channel)
     end
 
     private
+
+      module ForwardedBufferedIo
+        def fill(n=8192)
+          begin
+            super(n)
+          rescue Errno::ECONNRESET => e
+            debug { "connection was reset => shallowing exception:#{e}" }
+            return 0
+          rescue IOError => e                                 
+            if e.message =~ /closed/ then 
+              debug { "connection was reset => shallowing exception:#{e}" }
+              return 0
+            else
+              raise
+            end 
+          end
+        end
+
+        def send_pending
+          begin
+            super                                                          
+          rescue Errno::ECONNRESET => e
+            debug { "connection was reset => shallowing exception:#{e}" }
+            return 0
+          rescue IOError => e
+            if e.message =~ /closed/ then 
+              debug { "connection was reset => shallowing exception:#{e}" }
+              return 0
+            else
+              raise
+            end
+          end
+        end
+      end
 
       # 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.
       def prepare_client(client, channel, type)
         client.extend(Net::SSH::BufferedIo)
+        client.extend(ForwardedBufferedIo)
         client.logger = logger
 
         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
-
+
+        channel.on_eof do |ch|
+          debug { "eof #{type} on #{type} forwarded channel" }
+          ch[:socket].send_pending
+          ch[:socket].shutdown Socket::SHUT_WR
+        end
+
         channel.on_close do |ch|
           debug { "closing #{type} forwarded channel" }
           ch[:socket].close if !client.closed?