Ensure re-raised from write_file (#20)

Previously a `PathSet::Conflict` would get raised when `close` would cause a duplicate filename. `close` would then be called again, and raising the `PathSet::Conflict` _again_, overriding the actual exception raised within the block.

Also, a partial write of a local file entry header would not cause a correct rollback. This needs more state management variables to do correctly.

Closes #19 and addresses the underlying cause of #15

Author: julik

CHANGELOG.md

@@ -1,3 +1,8 @@
+## 6.3.2
+
+* Make sure `rollback!` correctly works with `write_file` and the original exception gets re-raised from `write_file` if
+  closing the current entry happens in `Writable#close`
+
 ## 6.3.1
 
 * Include `RailsStreaming` in a Rails loader callback, so that ActionController does not need to be in the namespace.

lib/zip_kit/streamer.rb

@@ -5,8 +5,12 @@
 # Is used to write ZIP archives without having to read them back or to overwrite
 # data. It outputs into any object that supports `<<` or `write`, namely:
 #
-# An `Array`, `File`, `IO`, `Socket` and even `String` all can be output destinations
-# for the `Streamer`.
+# * `Array` - will contain binary strings
+# * `File` - data will be written to it as it gets generated
+# * `IO` (`Socket`, `StringIO`) - data gets written into it
+# * `String` - in binary encoding and unfrozen - also makes a decent output target
+#
+# or anything else that responds to `#<<` or `#write`.
 #
 # You can also combine output through the `Streamer` with direct output to the destination,
 # all while preserving the correct offsets in the ZIP file structures. This allows usage
@@ -482,6 +486,10 @@ def update_last_entry_and_write_data_descriptor(crc32:, compressed_size:, uncomp
   # is likely already on the wire. However, excluding the entry from the central directory of the ZIP
   # file will allow better-behaved ZIP unarchivers to extract the entries which did store correctly,
   # provided they read the ZIP from the central directory and not straight-ahead.
+  # Rolling back does not perform any writes.
+  #
+  # `rollback!` gets called for you if an exception is raised inside the block of `write_file`,
+  # `write_deflated_file` and `write_stored_file`.
   #
   # @example
   #     zip.add_stored_entry(filename: "data.bin", size: 4.megabytes, crc32: the_crc)
@@ -493,14 +501,17 @@ def update_last_entry_and_write_data_descriptor(crc32:, compressed_size:, uncomp
   #     end
   # @return [Integer] position in the output stream / ZIP archive
   def rollback!
-    removed_entry = @files.pop
-    return @out.tell unless removed_entry
+    @files.pop if @remove_last_file_at_rollback
 
+    # Recreate the path set from remaining entries (PathSet does not support cheap deletes yet)
     @path_set.clear
     @files.each do |e|
       @path_set.add_directory_or_file_path(e.filename) unless e.filler?
     end
-    @files << Filler.new(@out.tell - removed_entry.local_header_offset)
+
+    # Create filler for the truncated or unusable local file entry that did get written into the output
+    filler_size_bytes = @out.tell - @offset_before_last_local_file_header
+    @files << Filler.new(filler_size_bytes)
 
     @out.tell
   end
@@ -554,6 +565,11 @@ def add_file_and_write_local_header(
     use_data_descriptor:,
     unix_permissions:
   )
+    # Set state needed for proper rollback later. If write_local_file_header
+    # does manage to write _some_ bytes, but fails later (we write in tiny bits sometimes)
+    # we should be able to create a filler from this offset on when we
+    @offset_before_last_local_file_header = @out.tell
+    @remove_last_file_at_rollback = false
 
     # Clean backslashes
     filename = remove_backslash(filename)
@@ -600,9 +616,11 @@ def add_file_and_write_local_header(
       mtime: e.mtime,
       filename: e.filename,
       storage_mode: e.storage_mode)
+
     e.bytes_used_for_local_header = @out.tell - e.local_header_offset
 
     @files << e
+    @remove_last_file_at_rollback = true
   end
 
   def remove_backslash(filename)

lib/zip_kit/streamer/heuristic.rb

@@ -26,6 +26,7 @@ def initialize(streamer, filename, **write_file_options)
     @bytes_deflated = 0
 
     @winner = nil
+    @started_closing = false
   end
 
   def <<(bytes)
@@ -40,13 +41,17 @@ def <<(bytes)
   end
 
   def close
+    return if @started_closing
+    @started_closing = true # started_closing because an exception may get raised inside close(), as we add an entry there
+
     decide unless @winner
     @winner.close
   end
 
   private def decide
     # Finish and then close the deflater - it has likely buffered some data
     @bytes_deflated += @deflater.finish.bytesize until @deflater.finished?
+
     # If the deflated version is smaller than the stored one
     # - use deflate, otherwise stored
     ratio = @bytes_deflated / @buf.size.to_f
@@ -55,9 +60,12 @@ def close
     else
       @streamer.write_stored_file(@filename, **@write_file_options)
     end
+
     # Copy the buffered uncompressed data into the newly initialized writable
     @buf.rewind
     IO.copy_stream(@buf, @winner)
     @buf.truncate(0)
+  ensure
+    @deflater.close
   end
 end

lib/zip_kit/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module ZipKit
-  VERSION = "6.3.1"
+  VERSION = "6.3.2"
 end

spec/zip_kit/file_reader_spec.rb

@@ -93,7 +93,7 @@ class EvilWriter < ZipKit::ZipWriter
         def write_end_of_central_directory(**kwargs)
           # Pretend there has to be more data
           kwargs[:central_directory_size] = kwargs[:central_directory_size] + 64
-          super(**kwargs)
+          super
         end
       end
 

spec/zip_kit/streamer_spec.rb

@@ -639,4 +639,40 @@ def stream_with_just_write.write(bytes)
     expect(per_filename["deflated.txt"]).to eq("this is attempt 2")
     expect(per_filename["stored.txt"]).to eq("this is attempt 2")
   end
+
+  it "correctly rolls back if an exception is raised from Writable#close when using write_file" do
+    # A Unicode string will not be happy about binary writes
+    # and will raise an exception. The exception won't be raised when
+    # starting the writes, but in `#close` of the Writable with write_file. When the size of the byte
+    # stream is small, the decision whether to call write_stored_file or write_deflated_file is going to
+    # be deferred until the Writable gets close() called on it. In that situation a correct rollback
+    # must be performed.
+    # If this is not handled correctly, the exception we will get raised won't be the original exception
+    # (Encoding::CompatibilityError that we need to see - to know what went wrong inside the writing block)
+    # but a PathSet::Conflict duplicate zip entry exception. Also, the IO offsets will be out of whack.
+    # To cause this, we need something that will raise during `Writable#close` - we will use a Unicode string
+    # for that purpose. See https://github.com/julik/zip_kit/issues/15
+    uniсode_str_buf = +"Ж"
+
+    zip = described_class.new(uniсode_str_buf)
+    4.times do
+      bytes_before_partial_write = uniсode_str_buf.bytesize
+      expect {
+        zip.write_file("deflated.txt") do |sink|
+          sink.write("x")
+        end
+      }.to raise_error(Encoding::CompatibilityError) # Should not be a PathSet::Conflict
+
+      # Ensure there was a partial write
+      expect(uniсode_str_buf.bytesize - bytes_before_partial_write).to be > 0
+    end
+
+    # We must force the string into binary so that the ZIP can be closed - we
+    # want to write out the EOCD since we want to make sure there is no byte offset
+    # mismatch (fillers have been correctly produced)
+    uniсode_str_buf.force_encoding(Encoding::BINARY)
+    expect {
+      zip.close
+    }.not_to raise_error # Offset validation should pass
+  end
 end

spec/zip_kit/zip_writer_spec.rb

@@ -3,7 +3,7 @@
 describe ZipKit::ZipWriter do
   class ByteReader < Struct.new(:io)
     def initialize(io)
-      super(io).tap { io.rewind }
+      super.tap { io.rewind }
     end
 
     def read_1b