Respect transfer stream ownership and IO contracts

lib/net/sftp/operations/download.rb

@@ -204,7 +204,7 @@ def []=(name, value)
 
       # A simple struct for encapsulating information about a single remote
       # file or directory that needs to be downloaded.
-      Entry = Struct.new(:remote, :local, :directory, :size, :handle, :offset, :sink)
+      Entry = Struct.new(:remote, :local, :directory, :size, :handle, :offset, :sink, :owned)
 
       #--
       # "ruby -w" hates private attributes, so we have to do these longhand
@@ -308,7 +308,13 @@ def on_open(response)
         raise StatusException.new(response, "open #{entry.remote}") unless response.ok?
 
         entry.handle = response[:handle]
-        entry.sink = entry.local.respond_to?(:write) ? entry.local : ::File.open(entry.local, "wb")
+        if entry.local.respond_to?(:write)
+          entry.sink = entry.local
+          entry.owned = false
+        else
+          entry.sink = ::File.open(entry.local, "wb")
+          entry.owned = true
+        end
         entry.offset = 0
 
         download_next_chunk(entry)
@@ -327,9 +333,9 @@ def download_next_chunk(entry)
       def on_read(response)
         entry = response.request[:entry]
 
-        if response.eof?
+        if response.eof? || (response.ok? && response[:data].empty?)
           update_progress(:close, entry)
-          entry.sink.close
+          entry.sink.close if entry.owned
           request = sftp.close(entry.handle, &method(:on_close))
           request[:entry] = entry
         elsif !response.ok?

lib/net/sftp/operations/upload.rb

@@ -227,7 +227,7 @@ def progress; @progress; end
 
       # A simple struct for recording metadata about the file currently being
       # uploaded.
-      LiveFile = Struct.new(:local, :remote, :io, :size, :handle)
+      LiveFile = Struct.new(:local, :remote, :io, :size, :handle, :owned, :offset)
 
       # The default # of bytes to read from disk at a time.
       DEFAULT_READ_SIZE   = 32_000
@@ -283,18 +283,20 @@ def open_file(local, remote)
         if local.respond_to?(:read)
           file = local
           name = options[:name] || "<memory>"
+          owned = false
         else
           file = ::File.open(local, "rb")
           name = local
+          owned = true
         end
 
         if file.respond_to?(:stat)
           size = file.stat.size
-        else
+        elsif file.respond_to?(:size)
           size = file.size
         end
 
-        metafile = LiveFile.new(name, remote, file, size)
+        metafile = LiveFile.new(name, remote, file, size, nil, owned, 0)
         update_progress(:open, metafile)
 
         request = sftp.open(remote, "w", &method(:on_open))
@@ -357,16 +359,17 @@ def write_next_chunk(file)
           process_next_entry
         else
           @active += 1
-          offset = file.io.pos
+          offset = file.offset
           data = file.io.read(options[:read_size] || DEFAULT_READ_SIZE)
-          if data.nil?
+          if data.nil? || data.empty?
             update_progress(:close, file)
             request = sftp.close(file.handle, &method(:on_close))
             request[:file] = file
-            file.io.close
+            file.io.close if file.owned
             file.io = nil
             @uploads.delete(file)
           else
+            file.offset += data.bytesize
             update_progress(:put, file, offset, data)
             request = sftp.write(file.handle, offset, data, &method(:on_write))
             request[:file] = file

test/test_download.rb

@@ -135,6 +135,22 @@ def test_download_file_should_transfer_remote_to_local_buffer
 
     assert_scripted_command { sftp.download(remote, local) }
     assert_equal text, local.string
+    assert !local.closed?
+  end
+
+  def test_download_should_treat_empty_data_as_end_of_file
+    expect_sftp_session :server_version => 3 do |channel|
+      channel.sends_packet(FXP_OPEN, :long, 0, :string, "/path/to/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, "")
+      channel.sends_packet(FXP_CLOSE, :long, 2, :string, "handle")
+      channel.gets_packet(FXP_STATUS, :long, 2, :long, 0)
+    end
+    local = StringIO.new
+
+    assert_scripted_command { sftp.download("/path/to/remote", local) }
+    assert_equal "", local.string
   end
 
   def test_download_directory_to_buffer_should_fail
@@ -286,4 +302,4 @@ def prepare_directory_tree_download(local, remote)
 
       [file1, file2]
     end
-end
+end

test/test_upload.rb

@@ -151,9 +151,32 @@ def test_upload_io_should_send_io_as_file
       channel.gets_packet(FXP_STATUS, :long, 2, :long, 0)
     end
 
-    assert_scripted_command do
-      sftp.upload(StringIO.new("this is some text"), "/path/to/remote")
+    source = StringIO.new("this is some text")
+    assert_scripted_command { sftp.upload(source, "/path/to/remote") }
+
+    assert !source.closed?
+  end
+
+  def test_upload_should_accept_a_read_only_stream
+    expect_file_transfer_from_stream("/path/to/remote", "this is some text")
+    chunks = ["this is some text", nil]
+    source = Object.new
+    source.define_singleton_method(:read) { |_size| chunks.shift }
+
+    assert_scripted_command { sftp.upload(source, "/path/to/remote") }
+  end
+
+  def test_upload_should_treat_an_empty_chunk_as_end_of_file
+    expect_sftp_session :server_version => 3 do |channel|
+      channel.sends_packet(FXP_OPEN, :long, 0, :string, "/path/to/remote", :long, 0x1A, :long, 0)
+      channel.gets_packet(FXP_HANDLE, :long, 0, :string, "handle")
+      channel.sends_packet(FXP_CLOSE, :long, 1, :string, "handle")
+      channel.gets_packet(FXP_STATUS, :long, 1, :long, 0)
     end
+    source = Object.new
+    source.define_singleton_method(:read) { |_size| "" }
+
+    assert_scripted_command { sftp.upload(source, "/path/to/remote") }
   end
 
   private
@@ -230,4 +253,15 @@ def expect_file_transfer(local, remote, data)
 
       expect_file(local, data)
     end
+
+    def expect_file_transfer_from_stream(remote, data)
+      expect_sftp_session :server_version => 3 do |channel|
+        channel.sends_packet(FXP_OPEN, :long, 0, :string, remote, :long, 0x1A, :long, 0)
+        channel.gets_packet(FXP_HANDLE, :long, 0, :string, "handle")
+        channel.sends_packet(FXP_WRITE, :long, 1, :string, "handle", :int64, 0, :string, data)
+        channel.sends_packet(FXP_CLOSE, :long, 2, :string, "handle")
+        channel.gets_packet(FXP_STATUS, :long, 1, :long, 0)
+        channel.gets_packet(FXP_STATUS, :long, 2, :long, 0)
+      end
+    end
 end