Differentiate #close and recovery

CHANGELOG.md

@@ -1,3 +1,9 @@
+* When rescuing a failed `write_file`, differentiate between `#close`
+  and `#release_resources_on_failure!`. Closing a Writable can still try
+  to do things to the Streamer output, it can try to write to the destination
+  IO which is no longer accepting writes and so on. What we do want is to
+  safely destroy the zlib deflaters.
+
 ## 6.3.2
 
 * Make sure `rollback!` correctly works with `write_file` and the original exception gets re-raised from `write_file` if

lib/zip_kit.rb

@@ -24,6 +24,7 @@ module ZipKit
   autoload :WriteShovel, File.dirname(__FILE__) + "/zip_kit/write_shovel.rb"
   autoload :RackChunkedBody, File.dirname(__FILE__) + "/zip_kit/rack_chunked_body.rb"
   autoload :RackTempfileBody, File.dirname(__FILE__) + "/zip_kit/rack_tempfile_body.rb"
+  autoload :ZlibCleanup, File.dirname(__FILE__) + "/zip_kit/zlib_cleanup.rb"
 
   require_relative "zip_kit/railtie" if defined?(::Rails)
 end

lib/zip_kit/streamer.rb

@@ -524,7 +524,7 @@ def yield_or_return_writable(writable, &block_to_pass_writable_to)
         yield(writable)
         writable.close
       rescue
-        writable.close
+        writable.release_resources_on_failure!
         rollback!
         raise
       end

lib/zip_kit/streamer/deflated_writer.rb

@@ -5,6 +5,7 @@
 # interchangeable with the StoredWriter in terms of interface.
 class ZipKit::Streamer::DeflatedWriter
   include ZipKit::WriteShovel
+  include ZipKit::ZlibCleanup
 
   # The amount of bytes we will buffer before computing the intermediate
   # CRC32 checksums. Benchmarks show that the optimum is 64KB (see
@@ -42,4 +43,8 @@ def finish
   ensure
     @deflater.close
   end
+
+  def release_resources_on_failure!
+    safely_dispose_of_incomplete_deflater(@deflater)
+  end
 end

lib/zip_kit/streamer/heuristic.rb

@@ -13,6 +13,8 @@
 # on the Streamer passed into it once it knows which compression
 # method should be applied
 class ZipKit::Streamer::Heuristic < ZipKit::Streamer::Writable
+  include ZipKit::ZlibCleanup
+
   BYTES_WRITTEN_THRESHOLD = 128 * 1024
   MINIMUM_VIABLE_COMPRESSION = 0.75
 
@@ -48,6 +50,11 @@ def close
     @winner.close
   end
 
+  def release_resources_on_failure!
+    safely_dispose_of_incomplete_deflater(@deflater)
+    @winner&.release_resources_on_failure!
+  end
+
   private def decide
     # Finish and then close the deflater - it has likely buffered some data
     @bytes_deflated += @deflater.finish.bytesize until @deflater.finished?

lib/zip_kit/streamer/stored_writer.rb

@@ -36,4 +36,8 @@ def finish
     @crc.flush
     {crc32: @crc_compute.to_i, compressed_size: @io.tell, uncompressed_size: @io.tell}
   end
+
+  def release_resources_on_failure!
+    # Nothing to do
+  end
 end

lib/zip_kit/streamer/writable.rb

@@ -33,4 +33,10 @@ def close
     @streamer.update_last_entry_and_write_data_descriptor(**@writer.finish)
     @closed = true
   end
+
+  def release_resources_on_failure!
+    return if @closed
+    @closed = true
+    @writer.release_resources_on_failure!
+  end
 end

lib/zip_kit/zlib_cleanup.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module ZipKit::ZlibCleanup
+  # This method is used to flush and close the native zlib handles
+  # should an archiving routine encounter an error. This is necessary,
+  # since otherwise unclosed deflaters may hang around in memory
+  # indefinitely, creating leaks.
+  #
+  # @param [Zlib::Deflater?]deflater the deflater to safely finish and close
+  # @return void
+  def safely_dispose_of_incomplete_deflater(deflater)
+    return unless deflater
+
+    # It can be a bit tricky to close and dealloc the deflater correctly.
+    # We want to do the right things for it to be GCd, including the
+    # native zlib handle. Also, leaving zlib handles dangling around
+    # creates warnings with "...with N bytes remaining to read", which are an
+    # eyesore. But they are there for a reason - so that we don't forget to do
+    # exactly this.
+    if !deflater.closed? && !deflater.finished?
+      deflater.finish until deflater.finished?
+    end
+    deflater.close unless deflater.closed?
+  end
+end

spec/zip_kit/zlib_cleanup_spec.rb

@@ -0,0 +1,31 @@
+require_relative "../spec_helper"
+
+describe "ZipKit::ZlibCleanup" do
+  it "does not raise when cleaning up the deflater which is in any step" do
+    cleaner = Object.new
+    class << cleaner
+      include ZipKit::ZlibCleanup
+    end
+
+    steps = [
+      ->(flate) { flate.deflate(Random.bytes(14)) },
+      ->(flate) { flate.deflate(Random.bytes(14)) },
+      ->(flate) { flate.deflate(Random.bytes(14)) },
+      ->(flate) { flate.deflate(Random.bytes(14)) },
+      ->(flate) { flate.finish until flate.finished? },
+      ->(flate) { flate.close }
+    ]
+
+    safe_close = cleaner.method(:safely_dispose_of_incomplete_deflater).to_proc
+    steps.length.times do |at_offset|
+      steps_with_failure = steps.dup
+      steps_with_failure.insert(at_offset, safe_close)
+
+      deflater = ::Zlib::Deflate.new(Zlib::DEFAULT_COMPRESSION, -::Zlib::MAX_WBITS)
+      steps_with_failure.each do |step_proc|
+        step_proc.call(deflater)
+        break if step_proc == safe_close
+      end
+    end
+  end
+end