From 372abbdb0cbc6c9c5d29370040e88ef53cc39d56 Mon Sep 17 00:00:00 2001
From: Yusuke Endoh <mame@ruby-lang.org>
Date: Sun, 25 Oct 2026 00:32:15 +0900
Subject: [PATCH] Keep ZSTREAM_IN_PROGRESS when raising InProgressError

The ensure that clears the flag also ran for a call rejected by
InProgressError.  After rescuing it, the stream was no longer marked
as in progress, so the recursive call prevented by [Bug #18358] was
allowed again.

  s = Zlib.deflate(Random.bytes(1024**2))
  zi = Zlib::Inflate.new
  GC.stress = true
  zi.inflate(s) do
    zi << s rescue nil
    zi << s
  end
  #=> heap-use-after-free (with ASAN)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 ext/zlib/zlib.c        |  9 +++++----
 test/zlib/test_zlib.rb | 11 +++++++++++
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/ext/zlib/zlib.c b/ext/zlib/zlib.c
index bcdd32a..7857cb4 100644
--- a/ext/zlib/zlib.c
+++ b/ext/zlib/zlib.c
@@ -1149,10 +1149,6 @@ zstream_run_try(VALUE value_arg)
     int err;
     VALUE old_input = Qnil;
 
-    /* Cannot start zstream while it is in progress. */
-    if (z->flags & ZSTREAM_IN_PROGRESS) {
-	rb_raise(cInProgressError, "zlib stream is in progress");
-    }
     z->flags |= ZSTREAM_IN_PROGRESS;
 
     if (ZSTREAM_REUSE_BUFFER_P(z)) {
@@ -1254,6 +1250,11 @@ zstream_run(struct zstream *z, Bytef *src, rb_len_t len, int flush)
         .stream_output = !ZSTREAM_IS_GZFILE(z) && rb_block_given_p(),
     };
 
+    /* Cannot start zstream while it is in progress. */
+    if (z->flags & ZSTREAM_IN_PROGRESS) {
+	rb_raise(cInProgressError, "zlib stream is in progress");
+    }
+
     rb_ensure(zstream_run_try, (VALUE)&args, zstream_run_ensure, (VALUE)&args);
     if (args.jump_state)
 	rb_jump_tag(args.jump_state);
diff --git a/test/zlib/test_zlib.rb b/test/zlib/test_zlib.rb
index e1ecbf8..756cc57 100644
--- a/test/zlib/test_zlib.rb
+++ b/test/zlib/test_zlib.rb
@@ -704,6 +704,17 @@ def test_recursive_deflate
       zd&.close
     end
 
+    def test_in_progress_error_does_not_end_in_progress
+      zd = Zlib::Deflate.new
+      s = Random.new(0).bytes(100_000)
+      zd.deflate(s, Zlib::FINISH) do
+        assert_raise(Zlib::InProgressError) { zd.flush }
+        assert_raise(Zlib::InProgressError) { zd << s }
+        break
+      end
+      zd.finish
+    end
+
     def test_recursive_inflate
       original_gc_stress = GC.stress
       GC.stress = true