Lock IO::Buffer coercion callbacks. (#18490)

Merged-By: ioquatix <samuel@codeotaku.com>

ext/-test-/io_buffer/io_buffer.c

@@ -44,6 +44,13 @@ io_buffer_raise(VALUE buffer, VALUE argument)
     rb_raise(rb_eRuntimeError, "interrupted");
 }
 
+static VALUE
+io_buffer_locked_p(VALUE buffer, VALUE argument)
+{
+    (void)argument;
+    return rb_funcall(buffer, rb_intern("locked?"), 0);
+}
+
 static VALUE
 io_buffer_for_reading_get_string(VALUE self, VALUE object)
 {
@@ -63,10 +70,15 @@ io_buffer_for_reading_object_id(VALUE self, VALUE object)
 }
 
 static VALUE
-io_buffer_for_reading_raise(VALUE self, VALUE string)
+io_buffer_for_reading_raise(VALUE self, VALUE object)
 {
-    StringValue(string);
-    return rb_io_buffer_for_reading(string, io_buffer_raise, Qnil);
+    return rb_io_buffer_for_reading(object, io_buffer_raise, Qnil);
+}
+
+static VALUE
+io_buffer_for_reading_locked_p(VALUE self, VALUE object)
+{
+    return rb_io_buffer_for_reading(object, io_buffer_locked_p, Qnil);
 }
 
 static VALUE
@@ -89,6 +101,18 @@ io_buffer_for_writing_modify_string(VALUE self, VALUE string)
     return rb_io_buffer_for_writing(string, io_buffer_modify_string, string);
 }
 
+static VALUE
+io_buffer_for_writing_raise(VALUE self, VALUE object)
+{
+    return rb_io_buffer_for_writing(object, io_buffer_raise, Qnil);
+}
+
+static VALUE
+io_buffer_for_writing_locked_p(VALUE self, VALUE object)
+{
+    return rb_io_buffer_for_writing(object, io_buffer_locked_p, Qnil);
+}
+
 static VALUE
 io_buffer_locked_for_reading_callback(const void *base, size_t size, VALUE buffer)
 {
@@ -191,9 +215,12 @@ Init_io_buffer(void)
     rb_define_singleton_method(mIOBuffer, "for_reading_readonly?", io_buffer_for_reading_readonly_p, 1);
     rb_define_singleton_method(mIOBuffer, "for_reading_object_id", io_buffer_for_reading_object_id, 1);
     rb_define_singleton_method(mIOBuffer, "for_reading_raise", io_buffer_for_reading_raise, 1);
+    rb_define_singleton_method(mIOBuffer, "for_reading_locked?", io_buffer_for_reading_locked_p, 1);
     rb_define_singleton_method(mIOBuffer, "for_writing_set_string", io_buffer_for_writing_set_string, 2);
     rb_define_singleton_method(mIOBuffer, "for_writing_readonly?", io_buffer_for_writing_readonly_p, 1);
     rb_define_singleton_method(mIOBuffer, "for_writing_modify_string", io_buffer_for_writing_modify_string, 1);
+    rb_define_singleton_method(mIOBuffer, "for_writing_raise", io_buffer_for_writing_raise, 1);
+    rb_define_singleton_method(mIOBuffer, "for_writing_locked?", io_buffer_for_writing_locked_p, 1);
     rb_define_singleton_method(mIOBuffer, "locked_for_reading", io_buffer_locked_for_reading, 1);
     rb_define_singleton_method(mIOBuffer, "locked_for_reading_raise", io_buffer_locked_for_reading_raise, 1);
     rb_define_singleton_method(mIOBuffer, "locked_for_writing", io_buffer_locked_for_writing, 1);

internal/io_buffer.h

@@ -8,10 +8,13 @@ RUBY_SYMBOL_EXPORT_BEGIN
 /**
  * Wrap string_or_buffer as a read-only IO::Buffer view and invoke callback(buffer, argument).
  *
- * - IO::Buffer: callback is called directly with no wrapping.
+ * The resulting buffer's backing allocation is locked for the duration of the
+ * callback and automatically unlocked when the callback returns or raises.
+ *
+ * - IO::Buffer: locked and passed directly to the callback.
  * - String: locked to prevent GC compaction from moving the backing memory,
- *   wrapped in a read-only IO::Buffer, callback called inside rb_ensure, buffer
- *   freed and string unlocked on exit.
+ *   wrapped in a locked read-only IO::Buffer, callback called inside
+ *   rb_ensure, buffer freed and string unlocked on exit.
  * - Other: TypeError raised.
  */
 VALUE rb_io_buffer_for_reading(VALUE string_or_buffer, VALUE (*callback)(VALUE buffer, VALUE argument), VALUE argument);
@@ -20,9 +23,12 @@ VALUE rb_io_buffer_for_reading(VALUE string_or_buffer, VALUE (*callback)(VALUE b
  * Wrap string_or_buffer as a writable IO::Buffer view and invoke callback(buffer, argument).
  *
  * - Read-only IO::Buffer: ArgumentError raised.
- * - IO::Buffer: callback is called directly with no wrapping.
+ * The resulting buffer's backing allocation is locked for the duration of the
+ * callback and automatically unlocked when the callback returns or raises.
+ *
+ * - IO::Buffer: locked and passed directly to the callback.
  * - String: locked, wrapped in a writable IO::Buffer, callback called inside
- *   rb_ensure, buffer freed and string unlocked on exit.
+ *   rb_ensure, buffer unlocked and freed, and string unlocked on exit.
  * - Other: TypeError raised.
  */
 VALUE rb_io_buffer_for_writing(VALUE string_or_buffer, VALUE (*callback)(VALUE buffer, VALUE argument), VALUE argument);

io_buffer.c

@@ -584,6 +584,35 @@ struct io_buffer_for_callback_arguments {
     VALUE argument;
 };
 
+static VALUE rb_io_buffer_locked_ensure(VALUE self);
+
+struct io_buffer_for_locked_callback_arguments {
+    VALUE buffer;
+    VALUE (*callback)(VALUE, VALUE);
+    VALUE argument;
+};
+
+static VALUE
+io_buffer_for_locked_callback_call(VALUE _arguments)
+{
+    struct io_buffer_for_locked_callback_arguments *arguments = (void *)_arguments;
+
+    return arguments->callback(arguments->buffer, arguments->argument);
+}
+
+static VALUE
+io_buffer_for_locked_callback(VALUE buffer, VALUE (*callback)(VALUE, VALUE), VALUE argument)
+{
+    struct io_buffer_for_locked_callback_arguments arguments = {
+        .buffer = buffer,
+        .callback = callback,
+        .argument = argument,
+    };
+
+    rb_io_buffer_lock(buffer);
+    return rb_ensure(io_buffer_for_locked_callback_call, (VALUE)&arguments, rb_io_buffer_locked_ensure, buffer);
+}
+
 static VALUE
 io_buffer_for_callback_call(VALUE _arguments)
 {
@@ -596,7 +625,7 @@ io_buffer_for_callback_call(VALUE _arguments)
         arguments->locked = 1;
     }
 
-    return arguments->callback(arguments->instance, arguments->argument);
+    return io_buffer_for_locked_callback(arguments->instance, arguments->callback, arguments->argument);
 }
 
 static VALUE
@@ -619,7 +648,7 @@ VALUE
 rb_io_buffer_for_reading(VALUE string_or_buffer, VALUE (*callback)(VALUE, VALUE), VALUE argument)
 {
     if (rb_obj_is_kind_of(string_or_buffer, rb_cIOBuffer)) {
-        return callback(string_or_buffer, argument);
+        return io_buffer_for_locked_callback(string_or_buffer, callback, argument);
     }
     else if (RB_TYPE_P(string_or_buffer, T_STRING)) {
         StringValue(string_or_buffer);
@@ -652,7 +681,7 @@ rb_io_buffer_for_writing(VALUE string_or_buffer, VALUE (*callback)(VALUE, VALUE)
         if (io_buffer_readonly_p(buffer)) {
             rb_raise(rb_eArgError, "buffer is read-only");
         }
-        return callback(string_or_buffer, argument);
+        return io_buffer_for_locked_callback(string_or_buffer, callback, argument);
     }
     else if (RB_TYPE_P(string_or_buffer, T_STRING)) {
         StringValue(string_or_buffer);

test/ruby/test_io_buffer.rb

@@ -47,12 +47,15 @@ class TestIOBuffer < Test::Unit::TestCase
 
     assert_equal "hello", Bug::IOBuffer.for_reading_get_string(string)
     assert_equal true, Bug::IOBuffer.for_reading_readonly?(string)
+    assert_equal true, Bug::IOBuffer.for_reading_locked?(string)
   end
 
   def test_internal_for_reading_with_io_buffer
     buffer = IO::Buffer.for("hello")
 
     assert_equal buffer.object_id, Bug::IOBuffer.for_reading_object_id(buffer)
+    assert_equal true, Bug::IOBuffer.for_reading_locked?(buffer)
+    refute_predicate buffer, :locked?
   end
 
   def test_internal_for_writing_with_string
@@ -62,6 +65,14 @@ class TestIOBuffer < Test::Unit::TestCase
 
     assert_equal "world", string
     assert_equal false, Bug::IOBuffer.for_writing_readonly?(string)
+    assert_equal true, Bug::IOBuffer.for_writing_locked?(string)
+  end
+
+  def test_internal_for_writing_with_io_buffer
+    buffer = IO::Buffer.new(5)
+
+    assert_equal true, Bug::IOBuffer.for_writing_locked?(buffer)
+    refute_predicate buffer, :locked?
   end
 
   def test_internal_for_writing_rejects_readonly_buffer
@@ -94,6 +105,46 @@ class TestIOBuffer < Test::Unit::TestCase
     assert_equal "hello!", string
   end
 
+  def test_internal_for_reading_unlocks_io_buffer_after_callback_exception
+    buffer = IO::Buffer.new(5)
+
+    assert_raise(RuntimeError) do
+      Bug::IOBuffer.for_reading_raise(buffer)
+    end
+
+    refute_predicate buffer, :locked?
+  end
+
+  def test_internal_for_writing_unlocks_io_buffer_after_callback_exception
+    buffer = IO::Buffer.new(5)
+
+    assert_raise(RuntimeError) do
+      Bug::IOBuffer.for_writing_raise(buffer)
+    end
+
+    refute_predicate buffer, :locked?
+  end
+
+  def test_internal_for_reading_preserves_existing_lock
+    buffer = IO::Buffer.new(5)
+    Bug::IOBuffer.lock(buffer)
+
+    assert_equal true, Bug::IOBuffer.for_reading_locked?(buffer)
+    assert_predicate buffer, :locked?
+  ensure
+    Bug::IOBuffer.unlock(buffer) if buffer&.locked?
+  end
+
+  def test_internal_for_writing_preserves_existing_lock
+    buffer = IO::Buffer.new(5)
+    Bug::IOBuffer.lock(buffer)
+
+    assert_equal true, Bug::IOBuffer.for_writing_locked?(buffer)
+    assert_predicate buffer, :locked?
+  ensure
+    Bug::IOBuffer.unlock(buffer) if buffer&.locked?
+  end
+
   def test_internal_locked_for_reading
     buffer = IO::Buffer.for("hello")