Merge pull request #26 from accardi/fix_fragmentation_download_bug

Fix fragmentation download failure

Author: delano

lib/net/sftp/session.rb

@@ -898,7 +898,7 @@ def when_channel_polled(channel)
             @packet_length = input.read_long
           end
 
-          return unless input.length >= @packet_length
+          return unless input.length >= @packet_length + 4
           packet = Net::SFTP::Packet.new(input.read(@packet_length))
           input.consume!
           @packet_length = nil

test/common.rb

@@ -125,6 +125,18 @@ def gets_packet(type, *args)
     gets_data(sftp_packet(type, *args))
   end
 
+  def gets_packet_in_two(fragment_len, type, *args)
+    fragment_len ||= 0
+    whole_packet = sftp_packet(type, *args)
+
+    if 0 < fragment_len && fragment_len < whole_packet.length
+      gets_data(whole_packet[0, whole_packet.length - fragment_len])
+      gets_data(whole_packet[-fragment_len..-1])
+    else
+      gets_data(whole_packet)
+    end
+  end
+
   def sends_packet(type, *args)
     sends_data(sftp_packet(type, *args))
   end

test/test_download.rb

@@ -21,6 +21,20 @@ def test_download_file_should_transfer_remote_to_local
     assert_equal text, file.string
   end
 
+  def test_download_file_should_transfer_remote_to_local_in_spite_of_fragmentation
+    local = "/path/to/local"
+    remote = "/path/to/remote"
+    text = "this is some text\n"
+
+    expect_file_transfer(remote, text, :fragment_len => 1)
+
+    file = StringIO.new
+    File.stubs(:open).with(local, "wb").returns(file)
+
+    assert_scripted_command { sftp.download(remote, local) }
+    assert_equal text, file.string
+  end
+
   def test_download_large_file_should_transfer_remote_to_local
     local = "/path/to/local"
     remote = "/path/to/remote"
@@ -130,12 +144,12 @@ def test_download_directory_to_buffer_should_fail
 
   private
 
-    def expect_file_transfer(remote, text)
+    def expect_file_transfer(remote, text, opts={})
       expect_sftp_session :server_version => 3 do |channel|
         channel.sends_packet(FXP_OPEN, :long, 0, :string, remote, :long, 0x01, :long, 0)
         channel.gets_packet(FXP_HANDLE, :long, 0, :string, "handle")
         channel.sends_packet(FXP_READ, :long, 1, :string, "handle", :int64, 0, :long, 32_000)
-        channel.gets_packet(FXP_DATA, :long, 1, :string, text)
+        channel.gets_packet_in_two(opts[:fragment_len], FXP_DATA, :long, 1, :string, text)
         channel.sends_packet(FXP_READ, :long, 2, :string, "handle", :int64, text.bytesize, :long, 32_000)
         channel.gets_packet(FXP_STATUS, :long, 2, :long, 1)
         channel.sends_packet(FXP_CLOSE, :long, 3, :string, "handle")
@@ -160,7 +174,7 @@ def prepare_large_file_download(local, remote, text, requested_chunk_size = FXP_
         channel.sends_packet(FXP_CLOSE, :long, data_packet_count + 2, :string, "handle")
         channel.gets_packet(FXP_STATUS, :long, data_packet_count + 2, :long, 0)
       end
-      
+
       file = StringIO.new
       File.stubs(:open).with(local, "wb").returns(file)