Replace RubyZip with ZipKit

ZipKit can read and write files directly from S3 which makes both imports and exports way more efficient in terms of disk usage.

Gemfile

@@ -34,7 +34,7 @@ gem "platform_agent"
 gem "aws-sdk-s3", require: false
 gem "web-push"
 gem "net-http-persistent"
-gem "rubyzip", require: "zip"
+gem "zip_kit"
 gem "mittens"
 gem "useragent", bc: "useragent"
 

Gemfile.lock

@@ -479,6 +479,7 @@ GEM
     xpath (3.2.0)
       nokogiri (~> 1.8)
     zeitwerk (2.7.4)
+    zip_kit (6.3.4)
 
 PLATFORMS
   aarch64-linux
@@ -523,7 +524,6 @@ DEPENDENCIES
   rouge
   rqrcode
   rubocop-rails-omakase
-  rubyzip
   selenium-webdriver
   solid_cable (>= 3.0)
   solid_cache (~> 1.0)
@@ -538,6 +538,7 @@ DEPENDENCIES
   web-console
   web-push
   webmock
+  zip_kit
 
 BUNDLED WITH
    2.7.2

app/models/account/import.rb

@@ -14,9 +14,8 @@ class Account::Import < ApplicationRecord
 
   def process(start: nil, callback: nil)
     processing!
-    ensure_downloaded
 
-    ZipFile.open(download_path) do |zip|
+    ZipFile.read_from(file.blob) do |zip|
       Account::DataTransfer::Manifest.new(account).each_record_set(start: start) do |record_set, last_id|
         record_set.import(from: zip, start: last_id, callback: callback)
       end
@@ -31,9 +30,8 @@ class Account::Import < ApplicationRecord
 
   def validate(start: nil, callback: nil)
     processing!
-    ensure_downloaded
 
-    ZipFile.open(download_path) do |zip|
+    ZipFile.read_from(file.blob) do |zip|
       Account::DataTransfer::Manifest.new(account).each_record_set(start: start) do |record_set, last_id|
         record_set.validate(from: zip, start: last_id, callback: callback)
       end
@@ -41,18 +39,6 @@ class Account::Import < ApplicationRecord
   end
 
   private
-    def ensure_downloaded
-      unless download_path.exist?
-        download_path.open("wb") do |f|
-          file.download { |chunk| f.write(chunk) }
-        end
-      end
-    end
-
-    def download_path
-      Pathname.new("/tmp/account-import-#{id}.zip")
-    end
-
     def mark_completed
       completed!
       ImportMailer.completed(identity, account).deliver_later

app/models/export.rb

@@ -20,19 +20,16 @@ class Export < ApplicationRecord
   def build
     processing!
 
-    tempfile = nil
     with_account_context do
-      tempfile = ZipFile.create { |zip| populate_zip(zip) }
-      file.attach io: File.open(tempfile.path), filename: "fizzy-export-#{id}.zip", content_type: "application/zip"
+      ZipFile.create_for(file, filename: "fizzy-export-#{id}.zip") do |zip|
+        populate_zip(zip)
+      end
       mark_completed
       ExportMailer.completed(self).deliver_later
     end
   rescue => e
     update!(status: :failed)
     raise e
-  ensure
-    tempfile&.close
-    tempfile&.unlink
   end
 
   def mark_completed

app/models/zip_file.rb

@@ -1,57 +1,80 @@
 class ZipFile
   class << self
-    def create
+    def create_for(attachment, filename:)
       raise ArgumentError, "No block given" unless block_given?
 
-      Tempfile.new([ "export", ".zip" ]).tap do |tempfile|
-        Zip::File.open(tempfile.path, create: true) do |zip|
-          yield new(zip)
+      reflection = attachment.record.class.reflect_on_attachment(attachment.name)
+      service_name = reflection.options[:service_name] || ActiveStorage::Blob.service.name
+      service = ActiveStorage::Blob.services.fetch(service_name)
+
+      if s3_service?(service)
+        create_for_s3(attachment, filename: filename, service: service) { |zip| yield zip }
+      else
+        create_for_disk(attachment, filename: filename) { |zip| yield zip }
+      end
+    end
+
+    def read_from(blob)
+      raise ArgumentError, "No block given" unless block_given?
+
+      if s3_service?(blob.service)
+        read_from_s3(blob) { |zip| yield zip }
+      else
+        read_from_disk(blob) { |zip| yield zip }
+      end
+    end
+
+    private
+      def s3_service?(service)
+        # The S3 service doesn't get loaded in development unless it's used
+        defined?(ActiveStorage::Service::S3Service) && service.is_a?(ActiveStorage::Service::S3Service)
+      end
+
+      def create_for_s3(attachment, filename:, service:)
+        blob = ActiveStorage::Blob.create_before_direct_upload!(
+          filename: filename,
+          content_type: "application/zip",
+          byte_size: 0,
+          checksum: "pending"
+        )
+
+        writer = Writer.new
+        service.upload(blob.key, writer.io) do |io|
+          writer.stream_to(io)
+          yield writer
+        end
+
+        blob.update!(byte_size: writer.byte_size, checksum: writer.checksum)
+        attachment.attach(blob)
+      end
+
+      def create_for_disk(attachment, filename:)
+        tempfile = Tempfile.new([ "export", ".zip" ])
+        tempfile.binmode
+
+        writer = Writer.new(tempfile)
+        yield writer
+        writer.close
+
+        tempfile.rewind
+        attachment.attach(io: tempfile, filename: filename, content_type: "application/zip")
+      ensure
+        tempfile&.close
+        tempfile&.unlink
+      end
+
+      def read_from_s3(blob)
+        url = blob.url(expires_in: 6.hour)
+        remote_io = ZipKit::RemoteIO.new(url)
+        reader = Reader.new(remote_io)
+        yield reader
+      end
+
+      def read_from_disk(blob)
+        blob.open do |file|
+          reader = Reader.new(file)
+          yield reader
         end
       end
-    end
-
-    def open(path)
-      raise ArgumentError, "No block given" unless block_given?
-
-      Zip::File.open(path.to_s) do |zip|
-        yield new(zip)
-      end
-    end
   end
-
-  def initialize(zip)
-    @zip = zip
-  end
-
-  def add_file(path, content = nil, compress: true, &block)
-    if block_given?
-      compression = compress ? nil : Zip::Entry::STORED
-      zip.get_output_stream(path, compression_method: compression, &block)
-    else
-      zip.get_output_stream(path) { |f| f.write(content) }
-    end
-  end
-
-  def glob(pattern)
-    zip.glob(pattern).map(&:name).sort
-  end
-
-  def read(file_path, &block)
-    entry = zip.find_entry(file_path)
-    raise ArgumentError, "File not found in zip: #{file_path}" unless entry
-    raise ArgumentError, "Cannot read directory entry: #{file_path}" if entry.directory?
-
-    if block_given?
-      yield entry.get_input_stream
-    else
-      entry.get_input_stream.read
-    end
-  end
-
-  def exists?(file_path)
-    zip.find_entry(file_path).present?
-  end
-
-  private
-    attr_reader :zip
 end

app/models/zip_file/reader.rb

@@ -0,0 +1,29 @@
+class ZipFile::Reader
+  def initialize(io)
+    @io = io
+    @reader = ZipKit::FileReader.read_zip_structure(io: io)
+  end
+
+  def read(file_path)
+    entry = @reader.find { |e| e.filename == file_path }
+    raise ArgumentError, "File not found in zip: #{file_path}" unless entry
+    raise ArgumentError, "Cannot read directory entry: #{file_path}" if entry.filename.end_with?("/")
+
+    extractor = entry.extractor_from(@io)
+    content = extractor.extract
+
+    if block_given?
+      yield StringIO.new(content)
+    else
+      content
+    end
+  end
+
+  def glob(pattern)
+    @reader.map(&:filename).select { |name| File.fnmatch(pattern, name) }.sort
+  end
+
+  def exists?(file_path)
+    @reader.any? { |e| e.filename == file_path }
+  end
+end

app/models/zip_file/writer.rb

@@ -0,0 +1,77 @@
+class ZipFile::Writer
+  attr_reader :byte_size
+
+  def initialize(io = nil)
+    @entries = []
+    @byte_size = 0
+    @output_io = io
+    @streamer = nil
+    @digest = Digest::MD5.new
+  end
+
+  def io
+    @counting_io ||= CountingIO.new(self)
+  end
+
+  def stream_to(io)
+    @output_io = io
+  end
+
+  def write(data)
+    @output_io.write(data)
+    @byte_size += data.bytesize
+    @digest.update(data)
+    data.bytesize
+  end
+
+  def add_file(path, content = nil, compress: true)
+    @entries << path
+    compression = compress ? :deflate : :stored
+
+    if block_given?
+      streamer.write_deflated_file(path) do |sink|
+        yield sink
+      end
+    else
+      streamer.write_deflated_file(path) do |sink|
+        sink.write(content)
+      end
+    end
+  end
+
+  def glob(pattern)
+    @entries.select { |e| File.fnmatch(pattern, e) }.sort
+  end
+
+  def exists?(path)
+    @entries.include?(path)
+  end
+
+  def close
+    streamer.close
+  end
+
+  def checksum
+    Base64.strict_encode64(@digest.digest)
+  end
+
+  private
+    def streamer
+      @streamer ||= ZipKit::Streamer.new(@output_io)
+    end
+
+  class CountingIO
+    def initialize(writer)
+      @writer = writer
+    end
+
+    def write(data)
+      @writer.write(data)
+    end
+
+    def <<(data)
+      write(data)
+      self
+    end
+  end
+end

test/models/account/export_test.rb

@@ -11,7 +11,7 @@ class Account::ExportTest < ActiveSupport::TestCase
 
   test "build sets status to failed on error" do
     export = Account::Export.create!(account: Current.account, user: users(:david))
-    ZipFile.stubs(:create).raises(StandardError.new("Test error"))
+    ZipFile.stubs(:create_for).raises(StandardError.new("Test error"))
 
     assert_raises(StandardError) do
       export.build
@@ -54,9 +54,9 @@ class Account::ExportTest < ActiveSupport::TestCase
 
     assert export.completed?
     export.file.open do |file|
-      Zip::File.open(file.path) do |zip|
-        assert zip.find_entry("storage/#{blob.key}"), "Expected blob file in zip"
-      end
+      reader = ZipKit::FileReader.read_zip_structure(io: file)
+      entry = reader.find { |e| e.filename == "storage/#{blob.key}" }
+      assert entry, "Expected blob file in zip"
     end
   end
 end

test/models/user/data_export_test.rb

@@ -42,21 +42,21 @@ class User::DataExportTest < ActiveSupport::TestCase
       export.file.download { |chunk| temp.write(chunk) }
       temp.rewind
 
-      Zip::File.open(temp.path) do |zip|
-        json_files = zip.glob("*.json")
-        assert json_files.any?, "Zip should contain at least one JSON file"
+      reader = ZipKit::FileReader.read_zip_structure(io: temp)
+      json_files = reader.select { |e| e.filename.end_with?(".json") }
+      assert json_files.any?, "Zip should contain at least one JSON file"
 
-        json_content = JSON.parse(zip.read(json_files.first.name))
-        assert json_content.key?("number")
-        assert json_content.key?("title")
-        assert json_content.key?("board")
-        assert json_content.key?("creator")
-        assert json_content["creator"].key?("id")
-        assert json_content["creator"].key?("name")
-        assert json_content["creator"].key?("email")
-        assert json_content.key?("description")
-        assert json_content.key?("comments")
-      end
+      extractor = json_files.first.extractor_from(temp)
+      json_content = JSON.parse(extractor.extract)
+      assert json_content.key?("number")
+      assert json_content.key?("title")
+      assert json_content.key?("board")
+      assert json_content.key?("creator")
+      assert json_content["creator"].key?("id")
+      assert json_content["creator"].key?("name")
+      assert json_content["creator"].key?("email")
+      assert json_content.key?("description")
+      assert json_content.key?("comments")
     end
   end