Fix Zlib::Inflate#inflate(buffer:) when the buffer changes between calls

ext/zlib/zlib.c

@@ -815,6 +815,7 @@ zstream_detach_buffer(struct zstream *z)
     }
 
     z->buf = Qnil;
+    z->flags &= ~ZSTREAM_REUSE_BUFFER;
     z->stream.next_out = 0;
     z->stream.avail_out = 0;
 
@@ -1154,6 +1155,12 @@ zstream_run_try(VALUE value_arg)
     }
     z->flags |= ZSTREAM_IN_PROGRESS;
 
+    if (ZSTREAM_REUSE_BUFFER_P(z)) {
+	/* always re-derive next_out; the user may have modified z->buf */
+	z->stream.next_out = NULL;
+	z->stream.avail_out = 0;
+    }
+
     if (NIL_P(z->input) && len == 0) {
 	z->stream.next_in = (Bytef*)"";
 	z->stream.avail_in = 0;
@@ -2195,12 +2202,24 @@ rb_inflate_inflate_body(VALUE _arguments)
                 rb_str_modify_expand(buffer, len);
             }
             rb_str_set_len(buffer, 0);
+            if (ZSTREAM_REUSE_BUFFER_P(z) && !NIL_P(z->buf) && RSTRING_LEN(z->buf) > 0) {
+                /* carry over the output pending in the previous buffer */
+                rb_str_buf_cat(buffer, RSTRING_PTR(z->buf), RSTRING_LEN(z->buf));
+            }
             z->flags |= ZSTREAM_REUSE_BUFFER;
             z->buf = buffer;
         }
     } else if (ZSTREAM_REUSE_BUFFER_P(z)) {
+        VALUE pending = z->buf;
         z->flags &= ~ZSTREAM_REUSE_BUFFER;
         z->buf = Qnil;
+        z->stream.next_out = NULL;
+        z->stream.avail_out = 0;
+        if (RSTRING_LEN(pending) > 0) {
+            /* carry over the output pending in the previous buffer */
+            zstream_append_buffer2(z, pending);
+        }
+        RB_GC_GUARD(pending);
     }
     rb_scan_args(argc, argv, "10", &src);
 

test/zlib/test_zlib.rb

@@ -427,6 +427,151 @@ def test_inflate_buffer_chunked
       }
     end
 
+    def test_inflate_buffer_modified_between_calls
+      # In the block form the buffer is kept across calls while the stream
+      # is in progress; the application may modify it in between.
+      src = "0" * 100_000
+      deflated = Zlib::Deflate.deflate(src)
+      half = deflated.bytesize / 2
+
+      # buffer cleared between calls: the output pending in the buffer is
+      # discarded by the application, the rest must still be inflated
+      z = Zlib::Inflate.new
+      buf = String.new(capacity: 1 << 14)
+      inflated = "".dup
+      z.inflate(deflated.byteslice(0, half), buffer: buf) { |chunk| inflated << chunk }
+      discarded = buf.bytesize
+      buf.clear
+      z.inflate(deflated.byteslice(half..), buffer: buf) { |chunk| inflated << chunk }
+      z.inflate(nil, buffer: buf) { |chunk| inflated << chunk }
+      assert_equal src.bytesize - discarded, inflated.bytesize
+      assert_match(/\A0*\z/, inflated)
+
+      # buffer reallocated between calls, content preserved
+      z = Zlib::Inflate.new
+      buf = String.new(capacity: 1 << 14)
+      inflated = "".dup
+      z.inflate(deflated.byteslice(0, half), buffer: buf) { |chunk| inflated << chunk }
+      pending = buf.dup
+      buf << ("x" * 50_000)
+      buf.replace(pending)
+      z.inflate(deflated.byteslice(half..), buffer: buf) { |chunk| inflated << chunk }
+      z.inflate(nil, buffer: buf) { |chunk| inflated << chunk }
+      assert_equal src, inflated
+
+      # buffer cleared, then continued with Inflate#<< and ZStream#finish
+      z = Zlib::Inflate.new
+      buf = String.new(capacity: 1 << 14)
+      inflated = "".dup
+      z.inflate(deflated.byteslice(0, half), buffer: buf) { |chunk| inflated << chunk }
+      discarded = buf.bytesize
+      buf.clear
+      z << deflated.byteslice(half..)
+      inflated << z.finish
+      assert_equal src.bytesize - discarded, inflated.bytesize
+      assert_match(/\A0*\z/, inflated)
+    end
+
+    def test_inflate_buffer_replaced_between_calls
+      # Output pending in the previous buffer is carried over when another
+      # buffer (or no buffer) is given for the next call.
+      src = "0" * 100_000
+      deflated = Zlib::Deflate.deflate(src)
+
+      # block form, a different buffer object for each call
+      z = Zlib::Inflate.new
+      inflated = "".dup
+      (0...deflated.bytesize).step(200) do |i|
+        z.inflate(deflated.byteslice(i, 200), buffer: String.new(capacity: 1 << 14)) { |chunk| inflated << chunk }
+      end
+      z.inflate(nil, buffer: String.new) { |chunk| inflated << chunk }
+      assert_equal src, inflated
+
+      # block form, buffer: dropped mid-stream
+      z = Zlib::Inflate.new
+      buf = String.new(capacity: 1 << 14)
+      inflated = "".dup
+      half = deflated.bytesize / 2
+      z.inflate(deflated.byteslice(0, half), buffer: buf) { |chunk| inflated << chunk }
+      z.inflate(deflated.byteslice(half..)) { |chunk| inflated << chunk }
+      inflated << z.finish
+      assert_equal src, inflated
+
+      # NeedDict handled with a different buffer object, with GC in between
+      dict = "dictionary-" * 4
+      d = Zlib::Deflate.new
+      d.set_dictionary(dict)
+      deflated = d.deflate("A" * 12_000, Zlib::FINISH)
+      z = Zlib::Inflate.new
+      head, rest = deflated.byteslice(0, 6), deflated.byteslice(6..)
+      assert_raise(Zlib::NeedDict) { z.inflate(head, buffer: String.new(capacity: 1 << 14)) }
+      z.set_dictionary(dict)
+      EnvUtil.under_gc_stress do
+        assert_equal "A" * 12_000, z.inflate(rest, buffer: String.new(capacity: 1 << 14))
+      end
+    end
+
+    def test_inflate_buffer_then_addstr_and_finish
+      src = "A" * 50_000
+      deflated = Zlib::Deflate.deflate(src)
+      half = deflated.bytesize / 2
+
+      z = Zlib::Inflate.new
+      buf = String.new(capacity: 1 << 14)
+      inflated = z.inflate(deflated.byteslice(0, half), buffer: buf).dup
+      z << deflated.byteslice(half..)
+      rest = z.finish
+      assert_instance_of String, rest
+      assert_equal src, inflated + rest
+    end
+
+    def test_inflate_buffer_dropped_while_handling_need_dict
+      dict = "dictionary-" * 4
+      d = Zlib::Deflate.new
+      d.set_dictionary(dict)
+      deflated = d.deflate("A" * 12_000, Zlib::FINISH)
+
+      z = Zlib::Inflate.new
+      buf = String.new(capacity: 1 << 14)
+      assert_raise(Zlib::NeedDict) { z.inflate(deflated, buffer: buf) }
+      buf.clear
+      z.set_dictionary(dict)
+      assert_equal "A" * 12_000, z.inflate("")
+    end
+
+    def test_inflate_buffer_modified_while_handling_need_dict
+      dict = "dictionary-" * 4
+      d = Zlib::Deflate.new
+      d.set_dictionary(dict)
+      deflated = d.deflate("A" * 12_000, Zlib::FINISH)
+
+      # buffer cleared (its heap block is freed) before retrying
+      z = Zlib::Inflate.new
+      buf = String.new(capacity: 1 << 14)
+      assert_raise(Zlib::NeedDict) { z.inflate(deflated, buffer: buf) }
+      buf.clear
+      z.set_dictionary(dict)
+      assert_equal "A" * 12_000, z.inflate("", buffer: buf)
+
+      # buffer grown beyond its capacity (its heap block is reallocated)
+      z = Zlib::Inflate.new
+      buf = String.new(capacity: 1 << 14)
+      assert_raise(Zlib::NeedDict) { z.inflate(deflated, buffer: buf) }
+      buf << ("x" * 20_000)
+      buf.clear
+      z.set_dictionary(dict)
+      assert_equal "A" * 12_000, z.inflate("", buffer: buf)
+
+      # buffer replaced by another string
+      z = Zlib::Inflate.new
+      buf = String.new(capacity: 1 << 14)
+      assert_raise(Zlib::NeedDict) { z.inflate(deflated, buffer: buf) }
+      buf.replace("x" * 20_000)
+      buf.clear
+      z.set_dictionary(dict)
+      assert_equal "A" * 12_000, z.inflate("", buffer: buf)
+    end
+
     def test_inflate_chunked_break
       # zeros = Zlib::Deflate.deflate("0" * 100_000)
       zeros = "x\234\355\3011\001\000\000\000\302\240J\353\237\316\032\036@" \