Cap what an account import reads out of a zip entry

app/models/account/import.rb

@@ -85,6 +85,9 @@ def cleanup
   end
 
   private
+    # Only the Disk service stages the archive on local disk. On S3 the reader
+    # range-requests it and blob contents upload straight back, so there is no
+    # local free space to preflight.
     def ensure_sufficient_storage_space
       return unless path = ZipFile.path_on_disk(file.blob)
 

app/models/zip_file.rb

@@ -1,6 +1,13 @@
 class ZipFile
   class InvalidFileError < StandardError; end
 
+  # Both inherit from InvalidFileError so an archive that busts a limit is
+  # handled everywhere a bad export already is: the import stops with an
+  # "invalid export" reason, and the job discards it rather than resuming
+  # against the same entry.
+  class EntryTooLargeError < InvalidFileError; end
+  class ArchiveTooLargeError < InvalidFileError; end
+
   class << self
     def create_for(attachment, filename:)
       raise ArgumentError, "No block given" unless block_given?

app/models/zip_file/reader.rb

@@ -1,20 +1,45 @@
 class ZipFile::Reader
+  # Entries read into memory hold one database record each, and the largest
+  # column any of them can fill is a rich text body; every other text column the
+  # export writes tops out at 64KB. Reading them with no ceiling let a
+  # two-megabyte upload inflate to gigabytes inside a jobs worker.
+  MAX_BUFFERED_ENTRY_SIZE = 2.megabytes
+
+  # The extractor is fed compressed bytes, so how much it hands back in one call
+  # is the archive author's choice: a maximally compressible slice expands about
+  # a thousandfold. Slicing here keeps every real record a single read of the
+  # archive, which matters on S3 where each read is its own range request, while
+  # capping what one call can return at tens of megabytes.
+  EXTRACT_SLICE_SIZE = 64.kilobytes
+
+  # Deflating each entry on its own is what makes an archive's total expansion a
+  # usable signal: a real export barely shrinks, since records are small and per
+  # entry overhead eats the savings, while a crafted one expands a thousandfold.
+  # Budgeting everything read out of an archive against what was actually
+  # uploaded bounds the streaming path too, where an entry's own declared size is
+  # the archive author's word. The floor keeps a small archive holding one large
+  # rich text body importable.
+  MAX_TOTAL_EXPANSION = 100
+  MIN_EXPANSION_BUDGET = 64.megabytes
+
   def initialize(io)
     @io = io
     @reader = ZipKit::FileReader.read_zip_structure(io: io)
+    @expanded = 0
+    @budget = [ io.size * MAX_TOTAL_EXPANSION, MIN_EXPANSION_BUDGET ].max
   rescue ZipKit::FileReader::ReadError, ZipKit::FileReader::MissingEOCD, ZipKit::FileReader::UnsupportedFeature => e
     raise ZipFile::InvalidFileError, e.message
   end
 
-  def read(file_path)
+  def read(file_path, max_bytes: MAX_BUFFERED_ENTRY_SIZE)
     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?("/")
 
     if block_given?
-      yield ZipFile::Reader::IO.new(entry, @io)
+      yield ZipFile::Reader::IO.new(entry, @io, self)
     else
-      entry.extractor_from(@io).extract
+      extract_within(entry, max_bytes)
     end
   end
 
@@ -25,4 +50,42 @@ def glob(pattern)
   def exists?(file_path)
     @reader.any? { |e| e.filename == file_path }
   end
+
+  # Called for every byte handed out, buffered or streamed.
+  def count_expanded(bytes)
+    @expanded += bytes
+
+    if @expanded > @budget
+      raise ZipFile::ArchiveTooLargeError,
+        "archive has produced #{@expanded} bytes, over the #{@budget} byte limit for its size"
+    end
+  end
+
+  private
+    def extract_within(entry, max_bytes)
+      ensure_within entry, entry.uncompressed_size, max_bytes
+
+      extractor = entry.extractor_from(@io)
+      content = "".b
+
+      until extractor.eof?
+        chunk = extractor.extract(EXTRACT_SLICE_SIZE)
+        break if chunk.nil?
+
+        count_expanded chunk.bytesize
+        content << chunk
+        ensure_within entry, content.bytesize, max_bytes
+      end
+
+      content
+    end
+
+    # The size an entry declares is the archive author's word, so the bytes that
+    # come out of the extractor are counted as well.
+    def ensure_within(entry, bytes, max_bytes)
+      if bytes > max_bytes
+        raise ZipFile::EntryTooLargeError,
+          "#{entry.filename} expands to at least #{bytes} bytes, over the #{max_bytes} byte limit"
+      end
+    end
 end

app/models/zip_file/reader/io.rb

@@ -1,15 +1,18 @@
 class ZipFile::Reader::IO
-  def initialize(entry, io)
+  STORED = 0
+
+  def initialize(entry, io, reader)
     @entry = entry
     @io = io
-    @extractor = @entry.extractor_from(@io)
+    @reader = reader
+    rewind
   end
 
   def read(length = nil, buffer = nil)
-    return nil if @extractor.eof?
+    fill_buffer_for(length)
+    return nil if eof?
 
-    data = @extractor.extract(length)
-    return nil if data.nil?
+    data = take(length)
 
     if buffer
       buffer.replace(data)
@@ -20,15 +23,73 @@ def read(length = nil, buffer = nil)
   end
 
   def eof?
-    @extractor.eof?
+    buffered.zero? && drained?
   end
 
   def rewind
     @extractor = @entry.extractor_from(@io)
+    @buffer = "".b
+    @consumed = 0
+    @drained = false
     0
   end
 
   def size
     @entry.uncompressed_size
   end
+
+  private
+    def fill_buffer_for(length)
+      until drained? || (length && buffered >= length)
+        drop_consumed
+
+        chunk = @extractor.extract(slice_size(length))
+
+        if chunk.nil?
+          @drained = true
+        else
+          @reader.count_expanded chunk.bytesize
+          @buffer << chunk
+        end
+      end
+    end
+
+    # An entry whose compressed data runs out before the deflate stream ends
+    # leaves the extractor short of eof while it has nothing left to give, so
+    # take its first nil as the end too rather than asking again forever.
+    def drained?
+      @drained || @extractor.eof?
+    end
+
+    # A deflated entry hands back whatever it inflates to, which is the archive
+    # author's choice rather than the caller's: asking for five megabytes of a
+    # crafted entry gets gigabytes back in one string. Read those in slices and
+    # keep what the caller didn't ask for. A stored entry returns exactly the
+    # bytes asked of it, so its reads pass straight through.
+    def slice_size(length)
+      if @entry.storage_mode == STORED
+        length
+      else
+        [ ZipFile::Reader::EXTRACT_SLICE_SIZE, length ].compact.min
+      end
+    end
+
+    def take(length)
+      wanted = length ? [ length, buffered ].min : buffered
+
+      @buffer.byteslice(@consumed, wanted).tap { @consumed += wanted }
+    end
+
+    def buffered
+      @buffer.bytesize - @consumed
+    end
+
+    # Handing bytes out by moving a cursor rather than trimming the front keeps a
+    # long run of small reads from recopying the rest of the buffer every time.
+    def drop_consumed
+      if @consumed > 0
+        @buffer = @buffer.byteslice(@consumed..)
+        @consumed = 0
+      end
+    end
 end

test/jobs/account/data_import_job_test.rb

@@ -94,6 +94,17 @@ class Account::DataImportJobTest < ActiveJob::TestCase
     tampered_tempfile&.unlink
   end
 
+  test "discards the job when an export entry inflates past the size limit" do
+    import = Account::Import.create!(identity: identities(:david), account: Account.create!(name: "Import Test"))
+    Account::Import.any_instance.stubs(:check).raises(ZipFile::EntryTooLargeError)
+
+    assert_nothing_raised do
+      assert_no_enqueued_jobs only: Account::DataImportJob do
+        Account::DataImportJob.perform_now(import)
+      end
+    end
+  end
+
   private
     # Simulates a hand-edited export: an extra board_publications record with a
     # fresh id but a key copied from another record in the same export.

test/models/account/import_test.rb

@@ -187,6 +187,37 @@ class Account::ImportTest < ActiveSupport::TestCase
     tempfile&.unlink
   end
 
+  test "check sets failure_reason to invalid_export when a record entry inflates past the size limit" do
+    source_account = accounts("37s")
+    exporter = users(:david)
+    identity = exporter.identity
+
+    export = Account::Export.create!(account: source_account, user: exporter)
+    export.build
+
+    export_tempfile = Tempfile.new([ "export", ".zip" ])
+    export.file.open { |f| FileUtils.cp(f.path, export_tempfile.path) }
+
+    bomb_tempfile = with_oversized_tag(export_tempfile)
+
+    source_account.destroy!
+
+    target_account = Account.create_with_owner(account: { name: "Import Test" }, owner: { identity: identity, name: exporter.name })
+    import = Account::Import.create!(identity: identity, account: target_account)
+    Current.set(account: target_account) do
+      import.file.attach(io: File.open(bomb_tempfile.path), filename: "export.zip", content_type: "application/zip")
+    end
+
+    assert_raises(ZipFile::EntryTooLargeError) { import.check }
+
+    assert import.reload.failed_due_to_invalid_export?
+  ensure
+    export_tempfile&.close
+    export_tempfile&.unlink
+    bomb_tempfile&.close
+    bomb_tempfile&.unlink
+  end
+
   test "check sets failure_reason to conflict when records already exist" do
     source_account = accounts("37s")
     exporter = users(:david)
@@ -311,6 +342,29 @@ class Account::ImportTest < ActiveSupport::TestCase
   end
 
   private
+    def with_oversized_tag(export_tempfile)
+      bomb = Tempfile.new([ "oversized_export", ".zip" ])
+      bomb.binmode
+
+      File.open(export_tempfile.path, "rb") do |file|
+        reader = ZipFile::Reader.new(file)
+        writer = ZipFile::Writer.new(bomb)
+
+        target = reader.glob("data/tags/*.json").first
+        tag = JSON.parse(reader.read(target))
+
+        reader.glob("**/*").reject { |name| name.end_with?("/") || name == target }.each do |entry|
+          writer.add_file(entry, reader.read(entry))
+        end
+
+        writer.add_file target, tag.merge("title" => "a" * (ZipFile::Reader::MAX_BUFFERED_ENTRY_SIZE + 1)).to_json
+        writer.close
+      end
+
+      bomb.rewind
+      bomb
+    end
+
     def import_with_attached_zip
       account = Account.create!(name: "Disk Check")
       import = Account::Import.create!(account: account, identity: identities(:david))