xlsxrb 0.1.12 → 0.1.13

data/lib/xlsxrb/ooxml/zip_reader.rb

@@ -9,129 +9,356 @@ require "tempfile"
 module Xlsxrb
   module Ooxml
     # Reads ZIP archives using only stdlib (zlib).
-    # Scans local file headers sequentially — works with non-seekable IO.
+    # Supports streaming entry decompression with bounded memory consumption.
     class ZipReader
       LOCAL_HEADER_SIG = "PK\x03\x04".b
       MAX_UNCOMPRESSED_SIZE = 500 * 1024 * 1024 # 500MB per file limit
       attr_accessor :max_uncompressed_size
 
+      # Metadata record for a single archive entry.
+      class Entry
+        attr_reader :name, :method, :compressed_size, :uncompressed_size, :crc32, :local_header_offset, :data_offset
+
+        #: (name: String, method: Integer, compressed_size: Integer, uncompressed_size: Integer, crc32: Integer, local_header_offset: Integer, data_offset: Integer) -> void
+        def initialize(name:, method:, compressed_size:, uncompressed_size:, crc32:, local_header_offset:, data_offset:)
+          @name = name
+          @method = method
+          @compressed_size = compressed_size
+          @uncompressed_size = uncompressed_size
+          @crc32 = crc32
+          @local_header_offset = local_header_offset
+          @data_offset = data_offset
+        end
+      end
+
+      # IO-like stream wrapper that decompresses entry chunks on demand without full buffering.
+      class EntryStreamIO
+        DEFAULT_CHUNK_SIZE = 65_536
+
+        attr_reader :data_offset, :compressed_size, :compression_method
+
+        #: (untyped io, Integer data_offset, Integer compressed_size, Integer compression_method, ?max_uncompressed_size: Integer?) -> void
+        def initialize(io, data_offset, compressed_size, compression_method, max_uncompressed_size: nil)
+          @io = io
+          @data_offset = data_offset
+          @compressed_size = compressed_size
+          @compression_method = compression_method
+          @max_uncompressed_size = max_uncompressed_size || ZipReader::MAX_UNCOMPRESSED_SIZE
+          @buffer = +""
+          @bytes_read_compressed = 0
+          @total_uncompressed = 0
+          @eof = false
+          @closed = false
+          @inflater = @compression_method == 8 ? Zlib::Inflate.new(-Zlib::MAX_WBITS) : nil
+        end
+
+        #: (?Integer? length) -> String?
+        def read(length = nil)
+          raise IOError, "closed stream" if @closed
+
+          if length.nil?
+            result = @buffer.dup
+            @buffer.clear
+            while (chunk = fetch_chunk(DEFAULT_CHUNK_SIZE))
+              result << chunk
+            end
+            return result
+          end
+
+          return "" if length.zero?
+
+          while @buffer.bytesize < length && !@eof
+            chunk = fetch_chunk([length - @buffer.bytesize, DEFAULT_CHUNK_SIZE].max)
+            break unless chunk
+
+            @buffer << chunk
+          end
+
+          return nil if @buffer.empty? && @eof
+
+          @buffer.slice!(0, length)
+        end
+
+        #: (?Integer chunk_size) ?{ (String) -> void } -> (void | Enumerator[String, void])
+        def each_chunk(chunk_size = DEFAULT_CHUNK_SIZE, &block)
+          return enum_for(:each_chunk, chunk_size) unless block
+          raise IOError, "closed stream" if @closed
+
+          while (chunk = read(chunk_size))
+            break if chunk.nil? || chunk.empty?
+
+            yield chunk
+          end
+        end
+
+        #: () -> void
+        def close
+          return if @closed
+
+          @closed = true
+          @inflater&.close
+          @inflater = nil
+          @buffer.clear
+          nil
+        end
+
+        #: () -> bool
+        def closed?
+          @closed
+        end
+
+        #: () -> bool
+        def eof?
+          @eof && @buffer.empty?
+        end
+
+        private
+
+        def fetch_chunk(chunk_size)
+          return nil if @eof || @closed
+
+          if @compression_method.zero?
+            remaining = @compressed_size - @bytes_read_compressed
+            if remaining <= 0
+              @eof = true
+              return nil
+            end
+
+            read_len = [remaining, chunk_size].min
+            @io.seek(@data_offset + @bytes_read_compressed, IO::SEEK_SET)
+            raw = @io.read(read_len)
+            if raw.nil? || raw.empty?
+              @eof = true
+              return nil
+            end
+
+            @bytes_read_compressed += raw.bytesize
+            @total_uncompressed += raw.bytesize
+            raise ArgumentError, "ZIP bomb detected: Uncompressed size exceeds #{@max_uncompressed_size} bytes" if @total_uncompressed > @max_uncompressed_size
+
+            raw
+          elsif @compression_method == 8
+            inflated = +""
+            while inflated.empty? && !@eof
+              if @bytes_read_compressed < @compressed_size
+                remaining_compressed = @compressed_size - @bytes_read_compressed
+                read_len = [remaining_compressed, chunk_size].min
+                @io.seek(@data_offset + @bytes_read_compressed, IO::SEEK_SET)
+                raw = @io.read(read_len)
+                if raw.nil? || raw.empty?
+                  @eof = true
+                else
+                  @bytes_read_compressed += raw.bytesize
+                  inflater = @inflater
+                  inflated << inflater.inflate(raw) if inflater
+                end
+              else
+                inflater = @inflater
+                if inflater && !inflater.finished?
+                  begin
+                    inflated << inflater.finish
+                  rescue Zlib::BufError, Zlib::DataError
+                    # Stream ended
+                  end
+                end
+                @eof = true
+              end
+            end
+
+            @eof = true if @inflater&.finished?
+
+            @total_uncompressed += inflated.bytesize
+            raise ArgumentError, "ZIP bomb detected: Uncompressed size exceeds #{@max_uncompressed_size} bytes" if @total_uncompressed > @max_uncompressed_size
+
+            inflated.empty? && @eof ? nil : inflated
+          else
+            @eof = true
+            nil
+          end
+        end
+      end
+
       # Opens a ZIP from a file path or IO and yields the reader.
       #: (untyped source, ?max_uncompressed_size: Integer) ?{ (ZipReader) -> untyped } -> (ZipReader | untyped)
       def self.open(source, max_uncompressed_size: MAX_UNCOMPRESSED_SIZE)
-        io = source.is_a?(String) ? File.open(source, "rb") : source
-        reader = new(io, max_uncompressed_size: max_uncompressed_size)
+        should_close = source.is_a?(String) && File.file?(source)
+        io = if source.is_a?(String)
+               if source.start_with?(LOCAL_HEADER_SIG) || source.include?("\x00")
+                 StringIO.new(source.b)
+               else
+                 File.open(source, "rb")
+               end
+             else
+               source
+             end
+        reader = new(io, max_uncompressed_size: max_uncompressed_size, should_close_io: should_close)
         if block_given?
           begin
             yield reader
           ensure
-            io.close if source.is_a?(String)
+            reader.close
           end
         else
           reader
         end
       end
 
-      #: (untyped io, ?max_uncompressed_size: Integer) -> void
-      def initialize(io, max_uncompressed_size: MAX_UNCOMPRESSED_SIZE)
-        @io = io
+      #: (untyped io, ?max_uncompressed_size: Integer, ?should_close_io: bool) -> void
+      def initialize(io, max_uncompressed_size: MAX_UNCOMPRESSED_SIZE, should_close_io: false)
+        @raw_io = io
         @max_uncompressed_size = max_uncompressed_size
-        @entries = nil
+        @should_close_io = should_close_io
+        @tempfile = nil
+        @io = nil
+        @catalog = nil
+        @cached_entries = nil
       end
 
       # Returns a Hash { entry_name => raw_bytes } for all entries.
+      #: () -> Hash[String, String?]
       def read_all
         result = {}
-        each_entry { |name, data| result[name] = data }
+        catalog.each_key do |name|
+          result[name] = read_entry(name)
+        end
         result
       end
 
       # Returns raw bytes for a single entry, or nil if not found.
+      #: (String name) -> String?
       def read_entry(name)
-        entries[name]
+        return @cached_entries[name] if @cached_entries&.key?(name)
+
+        entry = catalog[name]
+        return nil unless entry
+
+        io.seek(entry.data_offset, IO::SEEK_SET)
+        raw = io.read(entry.compressed_size)
+        decompress(raw, entry.method)
       end
 
       # Yields (entry_name, data_string) for each file in the archive.
+      #: () ?{ (String, String?) -> void } -> (void | Enumerator[[String, String?], void])
       def each_entry(&block)
         return enum_for(:each_entry) unless block
 
-        entries.each_pair(&block)
+        catalog.each_key do |name|
+          block.call(name, read_entry(name))
+        end
       end
 
-      private
+      # Opens an EntryStreamIO for an entry by name.
+      #: (String entry_name) -> EntryStreamIO
+      def open_entry_io(entry_name)
+        entry = catalog[entry_name]
+        raise ArgumentError, "Entry not found in archive: #{entry_name}" unless entry
 
-      def entries
-        @entries ||= parse_entries
+        EntryStreamIO.new(io, entry.data_offset, entry.compressed_size, entry.method, max_uncompressed_size: @max_uncompressed_size)
       end
 
-      def parse_entries
-        io = if @io.is_a?(StringIO) || (@io.respond_to?(:read) && @io.respond_to?(:seek) && @io.respond_to?(:pos))
-               @io
-             elsif @io.respond_to?(:each) || @io.is_a?(Enumerator)
-               tf = Tempfile.new("xlsxrb_zip")
-               tf.binmode
-               @io.each { |chunk| tf.write(chunk) }
-               tf.rewind
-               tf
-             else
-               StringIO.new(@io.read.b)
-             end
-
-        io.binmode if io.respond_to?(:binmode)
-
-        first_sig = io.read(4)
-        raise ArgumentError, "Invalid magic number: Expected a valid ZIP/XLSX file format (PK\\x03\\x04)" unless first_sig == LOCAL_HEADER_SIG
+      # Yields inflated binary chunks directly from the compressed stream.
+      #: (String entry_name, ?chunk_size: Integer) ?{ (String) -> void } -> (void | Enumerator[String, void])
+      def each_entry_chunk(entry_name, chunk_size: 65_536, &block)
+        return enum_for(:each_entry_chunk, entry_name, chunk_size: chunk_size) unless block
 
-        # Try parsing via Central Directory first (standard and handles Data Descriptors perfectly)
-        cd_result = parse_from_central_directory(io)
-        if cd_result && !cd_result.empty?
-          if io.is_a?(Tempfile)
-            io.close
-            io.unlink
-          end
-          return cd_result
+        stream = open_entry_io(entry_name)
+        begin
+          stream.each_chunk(chunk_size, &block)
+        ensure
+          stream.close
         end
+      end
 
-        io.seek(-4, IO::SEEK_CUR) if io.respond_to?(:seek)
-        io = StringIO.new(first_sig + io.read.b) unless io.respond_to?(:seek)
-
-        result = {}
+      # Returns true if the archive contains the specified entry.
+      #: (String name) -> bool
+      def entry?(name)
+        catalog.key?(name)
+      end
 
-        loop do
-          sig = io.read(4)
-          break unless sig == LOCAL_HEADER_SIG
+      # Returns entry metadata catalog.
+      #: () -> Hash[String, Entry]
+      def entry_catalog
+        catalog
+      end
 
-          header = io.read(26)
-          break unless header && header.bytesize == 26
+      # Returns array of all entry names.
+      #: () -> Array[String]
+      def entry_names
+        catalog.keys
+      end
 
-          gp_flag         = header[2, 2].unpack1("v")
-          method          = header[4, 2].unpack1("v")
-          compressed_size = header[14, 4].unpack1("V")
-          name_len        = header[22, 2].unpack1("v")
-          extra_len       = header[24, 2].unpack1("v")
+      # Copies the specified entry to a ZipWriter directly without decompression.
+      #: (String entry_name, ZipWriter writer, ?target_path: String?) -> void
+      def copy_to_writer(entry_name, writer, target_path: nil)
+        entry = catalog[entry_name]
+        raise ArgumentError, "Entry not found in archive: #{entry_name}" unless entry
 
-          entry_name = io.read(name_len).force_encoding("UTF-8")
-          io.read(extra_len) # skip extra
+        if entry.compressed_size.zero? && entry.uncompressed_size.positive?
+          data = read_entry(entry_name)
+          writer.add_binary_entry(target_path || entry_name, data || "")
+          return
+        end
 
-          has_data_descriptor = gp_flag.anybits?(0x08)
+        writer.copy_raw_entry(
+          target_path || entry_name,
+          io,
+          data_offset: entry.data_offset,
+          compressed_size: entry.compressed_size,
+          uncompressed_size: entry.uncompressed_size,
+          crc32: entry.crc32,
+          method: entry.method
+        )
+      end
 
-          if has_data_descriptor && compressed_size.zero?
-            entry_data, = find_data_descriptor_stream(io, method)
-          else
-            raw = io.read(compressed_size)
-            entry_data = decompress(raw, method)
-          end
-          result[entry_name] = entry_data unless entry_name.end_with?("/")
+      # Closes the archive reader and cleans up any temporary resources.
+      #: () -> void
+      def close
+        @raw_io.close if @should_close_io && @raw_io.respond_to?(:close) && !@raw_io.closed?
+        if @tempfile
+          @tempfile.close
+          @tempfile.unlink
+          @tempfile = nil
         end
+        nil
+      end
+
+      private
 
-        if io.is_a?(Tempfile)
-          io.close
-          io.unlink
+      def io
+        @io ||= begin
+          prepared = if @raw_io.is_a?(StringIO) || (@raw_io.respond_to?(:read) && @raw_io.respond_to?(:seek) && @raw_io.respond_to?(:pos))
+                       @raw_io
+                     elsif @raw_io.respond_to?(:each) || @raw_io.is_a?(Enumerator)
+                       @tempfile = Tempfile.new("xlsxrb_zip")
+                       @tempfile.binmode
+                       @raw_io.each { |chunk| @tempfile.write(chunk) }
+                       @tempfile.rewind
+                       @tempfile
+                     else
+                       StringIO.new(@raw_io.read.b)
+                     end
+          prepared.binmode if prepared.respond_to?(:binmode)
+          prepared
         end
+      end
 
-        result
+      def catalog
+        @catalog ||= parse_catalog
       end
 
-      def parse_from_central_directory(io)
+      def parse_catalog
+        first_sig = io.read(4)
+        raise ArgumentError, "Invalid magic number: Expected a valid ZIP/XLSX file format (PK\\x03\\x04)" unless first_sig == LOCAL_HEADER_SIG
+
+        # Try Central Directory first
+        cd_catalog = parse_catalog_from_central_directory(io)
+        return cd_catalog if cd_catalog && !cd_catalog.empty?
+
+        io.seek(-4, IO::SEEK_CUR) if io.respond_to?(:seek)
+        parse_catalog_from_local_headers(io)
+      end
+
+      def parse_catalog_from_central_directory(io)
         return nil unless io.respond_to?(:seek) && io.respond_to?(:pos)
 
         io.seek(0, IO::SEEK_END)
@@ -165,7 +392,7 @@ module Xlsxrb
           cd_header = io.read(42)
           break unless cd_header && cd_header.bytesize == 42
 
-          _v_made, _v_need, _gp, method, _time, _date, _crc, csize, _usize, nlen, elen, clen, _dnum, _iattr, _eattr, offset = cd_header.unpack("vvvvvvVVVvvvvvVV")
+          _v_made, _v_need, _gp, method, _time, _date, crc, csize, usize, nlen, elen, clen, _dnum, _iattr, _eattr, offset = cd_header.unpack("vvvvvvVVVvvvvvVV")
           entry_name = io.read(nlen).force_encoding("UTF-8")
           io.read(elen + clen)
 
@@ -182,11 +409,17 @@ module Xlsxrb
 
           local_nlen = local_header[22, 2].unpack1("v")
           local_elen = local_header[24, 2].unpack1("v")
-          io.seek(offset + 30 + local_nlen + local_elen, IO::SEEK_SET)
-
-          raw = io.read(csize)
-          entry_data = decompress(raw, method)
-          result[entry_name] = entry_data
+          data_offset = offset + 30 + local_nlen + local_elen
+
+          result[entry_name] = Entry.new(
+            name: entry_name,
+            method: method,
+            compressed_size: csize,
+            uncompressed_size: usize,
+            crc32: crc,
+            local_header_offset: offset,
+            data_offset: data_offset
+          )
 
           io.seek(current_cd_pos, IO::SEEK_SET)
         end
@@ -198,6 +431,62 @@ module Xlsxrb
         nil
       end
 
+      def parse_catalog_from_local_headers(io)
+        result = {}
+        loop do
+          header_pos = io.pos
+          sig = io.read(4)
+          break unless sig == LOCAL_HEADER_SIG
+
+          header = io.read(26)
+          break unless header && header.bytesize == 26
+
+          gp_flag         = header[2, 2].unpack1("v")
+          method          = header[4, 2].unpack1("v")
+          crc             = header[10, 4].unpack1("V")
+          compressed_size = header[14, 4].unpack1("V")
+          uncompressed_size = header[18, 4].unpack1("V")
+          name_len        = header[22, 2].unpack1("v")
+          extra_len       = header[24, 2].unpack1("v")
+
+          entry_name = io.read(name_len).force_encoding("UTF-8")
+          io.read(extra_len)
+          data_offset = io.pos
+
+          next if entry_name.end_with?("/")
+
+          has_data_descriptor = gp_flag.anybits?(0x08)
+
+          if has_data_descriptor && compressed_size.zero?
+            entry_data, = find_data_descriptor_stream(io, method)
+            @cached_entries ||= {}
+            @cached_entries[entry_name] = entry_data
+            result[entry_name] = Entry.new(
+              name: entry_name,
+              method: method,
+              compressed_size: entry_data.bytesize,
+              uncompressed_size: entry_data.bytesize,
+              crc32: crc,
+              local_header_offset: header_pos,
+              data_offset: data_offset
+            )
+          else
+            result[entry_name] = Entry.new(
+              name: entry_name,
+              method: method,
+              compressed_size: compressed_size,
+              uncompressed_size: uncompressed_size,
+              crc32: crc,
+              local_header_offset: header_pos,
+              data_offset: data_offset
+            )
+
+            io.seek(compressed_size, IO::SEEK_CUR) if io.respond_to?(:seek)
+          end
+        end
+        result
+      end
+
       def find_data_descriptor_stream(io, method)
         if method == 8
           inflater = Zlib::Inflate.new(-Zlib::MAX_WBITS)
@@ -238,12 +527,10 @@ module Xlsxrb
       end
 
       def decompress(raw, method)
-        return raw&.dup&.force_encoding("UTF-8") || "" if method.zero? # stored
+        return raw&.dup&.force_encoding("UTF-8") || "" if method.zero?
 
-        # Deflated
         safe_inflate(raw || "", -Zlib::MAX_WBITS)
       rescue Zlib::DataError
-        # Try with raw deflate (no header)
         safe_inflate(raw || "", -Zlib::MAX_WBITS)
       end
 
@@ -264,7 +551,6 @@ module Xlsxrb
             result << inflated
             offset += chunk_size
           end
-          # Finish inflation
           inflated = inflater.finish
           raise ArgumentError, "ZIP bomb detected: Uncompressed size exceeds #{limit} bytes" if result.bytesize + inflated.bytesize > limit
 

data/lib/xlsxrb/ooxml/zip_writer.rb

@@ -109,7 +109,33 @@ module Xlsxrb
           compressed_size: @entry_compressed_size,
           uncompressed_size: @entry_uncompressed_size,
           offset: @entry_offset,
-          gp_flag: @entry_gp_flag
+          gp_flag: @entry_gp_flag,
+          method: 8
+        }
+      end
+
+      # Copies an entry from an existing archive IO without decompressing or recompressing.
+      #: (String path, untyped source_io, data_offset: Integer, compressed_size: Integer, uncompressed_size: Integer, crc32: Integer, ?method: Integer) -> void
+      def copy_raw_entry(path, source_io, data_offset:, compressed_size:, uncompressed_size:, crc32:, method: 8)
+        raise "ZipWriter is closed" if @closed
+
+        offset = @bytes_written
+        write_local_header(path, crc32, compressed_size, uncompressed_size, gp_flag: 0, method: method)
+
+        if compressed_size.positive?
+          source_io.seek(data_offset, IO::SEEK_SET)
+          IO.copy_stream(source_io, @io, compressed_size)
+          @bytes_written += compressed_size
+        end
+
+        @entries << {
+          path: path,
+          crc32: crc32,
+          compressed_size: compressed_size,
+          uncompressed_size: uncompressed_size,
+          offset: offset,
+          gp_flag: 0,
+          method: method
         }
       end
 
@@ -159,7 +185,8 @@ module Xlsxrb
           compressed_size: compressed.bytesize,
           uncompressed_size: content_bytes.bytesize,
           offset: offset,
-          gp_flag: 0
+          gp_flag: 0,
+          method: 8
         }
       end
 
@@ -172,13 +199,13 @@ module Xlsxrb
         end
       end
 
-      def write_local_header(path, crc, compressed_size, uncompressed_size, gp_flag: 0, count_bytes: true)
+      def write_local_header(path, crc, compressed_size, uncompressed_size, gp_flag: 0, count_bytes: true, method: 8)
         name_bytes = path.encode("UTF-8").b
         header = [
           0x04034B50,        # local file header signature
           20,                # version needed (2.0)
           gp_flag,           # general purpose bit flag (0x0008 if data descriptor follows)
-          8,                 # compression method (deflate)
+          method,            # compression method (0 = store, 8 = deflate)
           0,                 # last mod file time
           33,                # last mod file date
           crc,
@@ -206,7 +233,7 @@ module Xlsxrb
             20,                # version made by
             20,                # version needed
             entry[:gp_flag] || 0, # general purpose bit flag
-            8,                 # compression method
+            entry[:method] || 8, # compression method
             0,                 # last mod file time
             33,                # last mod file date
             entry[:crc32],