Merge pull request #421 from mfazekas/read-data-return-before-raiseignore-disconnect-exception-when-all-closed

non blocking next_packet shall return the available data before raisi…

2 files changed, 23 insertions, 2 deletions

diff --git a/lib/net/ssh/transport/packet_stream.rb b/lib/net/ssh/transport/packet_stream.rb
index 8e9ffe5..f38b5d7 100644
--- a/lib/net/ssh/transport/packet_stream.rb
+++ b/lib/net/ssh/transport/packet_stream.rb
@@ -86,7 +86,12 @@ module Net; module SSH; module Transport
       when :nonblock then
         if available_for_read?
           if fill <= 0
-            raise Net::SSH::Disconnect, "connection closed by remote host"
+            result = poll_next_packet
+            if result.nil?
+              raise Net::SSH::Disconnect, "connection closed by remote host"
+            else
+              return result
+            end
           end
         end
         poll_next_packet
diff --git a/test/transport/test_packet_stream.rb b/test/transport/test_packet_stream.rb
index f3e67f6..3e01f55 100644
--- a/test/transport/test_packet_stream.rb
+++ b/test/transport/test_packet_stream.rb
@@ -6,7 +6,7 @@ require 'net/ssh/transport/packet_stream'
 
 module Transport
 
-  class TestPacketStream < NetSSHTest
+  class TestPacketStream < NetSSHTest # rubocop:disable Metrics/ClassLength
     include Net::SSH::Transport::Constants
 
     def test_client_name_when_getnameinfo_works
@@ -142,6 +142,22 @@ module Transport
       assert_equal DEBUG, packet.type
     end
 
+    def test_nonblocking_next_packet_should_raise
+      IO.stubs(:select).returns([[stream]])
+      stream.stubs(:recv).returns("")
+      assert_raises(Net::SSH::Disconnect) { stream.next_packet(:nonblock) }
+    end
+
+    def test_nonblocking_next_packet_should_return_packet_before_raise
+      IO.stubs(:select).returns([[stream]])
+      stream.send(:input).append(packet)
+      stream.stubs(:recv).returns("")
+      packet = stream.next_packet(:nonblock)
+      assert_not_nil packet
+      assert_equal DEBUG, packet.type
+      assert_raises(Net::SSH::Disconnect) { stream.next_packet(:nonblock) }
+    end
+
     def test_next_packet_should_block_when_requested_until_entire_packet_is_available
       IO.stubs(:select).returns([[stream]])
       stream.stubs(:recv).returns(packet[0,10], packet[10,20], packet[20..-1])