io_buffer: add safe callback wrappers for String views. (#17473)

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

depend

@@ -7354,6 +7354,7 @@ io_buffer.$(OBJEXT): $(top_srcdir)/internal/fixnum.h
 io_buffer.$(OBJEXT): $(top_srcdir)/internal/gc.h
 io_buffer.$(OBJEXT): $(top_srcdir)/internal/imemo.h
 io_buffer.$(OBJEXT): $(top_srcdir)/internal/io.h
+io_buffer.$(OBJEXT): $(top_srcdir)/internal/io_buffer.h
 io_buffer.$(OBJEXT): $(top_srcdir)/internal/numeric.h
 io_buffer.$(OBJEXT): $(top_srcdir)/internal/sanitizers.h
 io_buffer.$(OBJEXT): $(top_srcdir)/internal/serial.h

ext/-test-/io_buffer/extconf.rb

@@ -0,0 +1,3 @@
+# frozen_string_literal: false
+require_relative "../auto_ext.rb"
+auto_ext(inc: true)

ext/-test-/io_buffer/io_buffer.c

@@ -0,0 +1,105 @@
+#include "ruby.h"
+#include "internal/io_buffer.h"
+
+static VALUE
+io_buffer_get_string(VALUE buffer, VALUE argument)
+{
+    (void)argument;
+    return rb_funcall(buffer, rb_intern("get_string"), 0);
+}
+
+static VALUE
+io_buffer_readonly_p(VALUE buffer, VALUE argument)
+{
+    (void)argument;
+    return rb_funcall(buffer, rb_intern("readonly?"), 0);
+}
+
+static VALUE
+io_buffer_object_id(VALUE buffer, VALUE argument)
+{
+    (void)argument;
+    return rb_obj_id(buffer);
+}
+
+static VALUE
+io_buffer_set_string(VALUE buffer, VALUE string)
+{
+    return rb_funcall(buffer, rb_intern("set_string"), 1, string);
+}
+
+static VALUE
+io_buffer_modify_string(VALUE buffer, VALUE string)
+{
+    (void)buffer;
+    rb_str_cat(string, "!", 1);
+    return Qnil;
+}
+
+static VALUE
+io_buffer_raise(VALUE buffer, VALUE argument)
+{
+    (void)buffer;
+    (void)argument;
+    rb_raise(rb_eRuntimeError, "interrupted");
+}
+
+static VALUE
+io_buffer_for_reading_get_string(VALUE self, VALUE object)
+{
+    return rb_io_buffer_for_reading(object, io_buffer_get_string, Qnil);
+}
+
+static VALUE
+io_buffer_for_reading_readonly_p(VALUE self, VALUE object)
+{
+    return rb_io_buffer_for_reading(object, io_buffer_readonly_p, Qnil);
+}
+
+static VALUE
+io_buffer_for_reading_object_id(VALUE self, VALUE object)
+{
+    return rb_io_buffer_for_reading(object, io_buffer_object_id, Qnil);
+}
+
+static VALUE
+io_buffer_for_reading_raise(VALUE self, VALUE string)
+{
+    StringValue(string);
+    return rb_io_buffer_for_reading(string, io_buffer_raise, Qnil);
+}
+
+static VALUE
+io_buffer_for_writing_set_string(VALUE self, VALUE object, VALUE string)
+{
+    StringValue(string);
+    return rb_io_buffer_for_writing(object, io_buffer_set_string, string);
+}
+
+static VALUE
+io_buffer_for_writing_readonly_p(VALUE self, VALUE object)
+{
+    return rb_io_buffer_for_writing(object, io_buffer_readonly_p, Qnil);
+}
+
+static VALUE
+io_buffer_for_writing_modify_string(VALUE self, VALUE string)
+{
+    StringValue(string);
+    return rb_io_buffer_for_writing(string, io_buffer_modify_string, string);
+}
+
+void
+Init_io_buffer(void)
+{
+    VALUE mBug = rb_define_module("Bug");
+    VALUE mIOBuffer = rb_define_module_under(mBug, "IOBuffer");
+
+    rb_define_singleton_method(mIOBuffer, "for_reading_get_string", io_buffer_for_reading_get_string, 1);
+    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_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);
+}

internal/io_buffer.h

@@ -0,0 +1,30 @@
+#pragma once
+
+#include <ruby/ruby.h>
+#include <ruby/io/buffer.h>
+
+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.
+ * - 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.
+ * - Other: TypeError raised.
+ */
+VALUE rb_io_buffer_for_reading(VALUE string_or_buffer, VALUE (*callback)(VALUE buffer, VALUE argument), VALUE argument);
+
+/**
+ * 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.
+ * - String: locked, wrapped in a writable IO::Buffer, callback called inside
+ *   rb_ensure, buffer 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);
+
+RUBY_SYMBOL_EXPORT_END

io_buffer.c

@@ -20,6 +20,7 @@
 #include "internal/numeric.h"
 #include "internal/string.h"
 #include "internal/io.h"
+#include "internal/io_buffer.h"
 
 VALUE rb_cIOBuffer;
 VALUE rb_eIOBufferLockedError;
@@ -556,6 +557,105 @@ io_buffer_for_yield_instance_ensure(VALUE _arguments)
     return Qnil;
 }
 
+struct io_buffer_for_callback_arguments {
+    VALUE klass;
+    VALUE string;
+    VALUE instance;
+    enum rb_io_buffer_flags flags;
+    int locked;
+    VALUE (*callback)(VALUE, VALUE);
+    VALUE argument;
+};
+
+static VALUE
+io_buffer_for_callback_call(VALUE _arguments)
+{
+    struct io_buffer_for_callback_arguments *arguments = (struct io_buffer_for_callback_arguments *)_arguments;
+
+    arguments->instance = io_buffer_for_make_instance(arguments->klass, arguments->string, arguments->flags);
+
+    if (!RB_OBJ_FROZEN(arguments->string)) {
+        rb_str_locktmp(arguments->string);
+        arguments->locked = 1;
+    }
+
+    return arguments->callback(arguments->instance, arguments->argument);
+}
+
+static VALUE
+io_buffer_for_callback_ensure(VALUE _arguments)
+{
+    struct io_buffer_for_callback_arguments *arguments = (struct io_buffer_for_callback_arguments *)_arguments;
+
+    if (arguments->instance != Qnil) {
+        rb_io_buffer_free(arguments->instance);
+    }
+
+    if (arguments->locked) {
+        rb_str_unlocktmp(arguments->string);
+    }
+
+    return Qnil;
+}
+
+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);
+    }
+    else if (RB_TYPE_P(string_or_buffer, T_STRING)) {
+        StringValue(string_or_buffer);
+        struct io_buffer_for_callback_arguments arguments = {
+            .klass = rb_cIOBuffer,
+            .string = string_or_buffer,
+            .instance = Qnil,
+            .flags = RB_IO_BUFFER_READONLY,
+            .locked = 0,
+            .callback = callback,
+            .argument = argument,
+        };
+        return rb_ensure(io_buffer_for_callback_call, (VALUE)&arguments,
+                         io_buffer_for_callback_ensure, (VALUE)&arguments);
+    }
+    else {
+        rb_raise(rb_eTypeError, "expected String or IO::Buffer, not %"PRIsVALUE,
+                 rb_obj_class(string_or_buffer));
+    }
+}
+
+/* Forward declaration: rb_io_buffer_readonly_p is defined later in this file. */
+int rb_io_buffer_readonly_p(VALUE self);
+
+VALUE
+rb_io_buffer_for_writing(VALUE string_or_buffer, VALUE (*callback)(VALUE, VALUE), VALUE argument)
+{
+    if (rb_obj_is_kind_of(string_or_buffer, rb_cIOBuffer)) {
+        if (rb_io_buffer_readonly_p(string_or_buffer)) {
+            rb_raise(rb_eArgError, "buffer is read-only");
+        }
+        return callback(string_or_buffer, argument);
+    }
+    else if (RB_TYPE_P(string_or_buffer, T_STRING)) {
+        StringValue(string_or_buffer);
+        struct io_buffer_for_callback_arguments arguments = {
+            .klass = rb_cIOBuffer,
+            .string = string_or_buffer,
+            .instance = Qnil,
+            .flags = 0,
+            .locked = 0,
+            .callback = callback,
+            .argument = argument,
+        };
+        return rb_ensure(io_buffer_for_callback_call, (VALUE)&arguments,
+                         io_buffer_for_callback_ensure, (VALUE)&arguments);
+    }
+    else {
+        rb_raise(rb_eTypeError, "expected String or IO::Buffer, not %"PRIsVALUE,
+                 rb_obj_class(string_or_buffer));
+    }
+}
+
 /*
  *  call-seq:
  *    IO::Buffer.for(string) -> readonly io_buffer

test/ruby/test_io_buffer.rb

@@ -2,6 +2,7 @@
 
 require 'tempfile'
 require 'rbconfig/sizeof'
+require '-test-/io_buffer'
 
 class TestIOBuffer < Test::Unit::TestCase
   experimental = Warning[:experimental]
@@ -31,6 +32,58 @@ class TestIOBuffer < Test::Unit::TestCase
     assert_equal 128, IO::Buffer::READONLY
   end
 
+  def test_internal_for_reading_with_string
+    string = "hello"
+
+    assert_equal "hello", Bug::IOBuffer.for_reading_get_string(string)
+    assert_equal true, Bug::IOBuffer.for_reading_readonly?(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)
+  end
+
+  def test_internal_for_writing_with_string
+    string = "hello"
+
+    Bug::IOBuffer.for_writing_set_string(string, "world")
+
+    assert_equal "world", string
+    assert_equal false, Bug::IOBuffer.for_writing_readonly?(string)
+  end
+
+  def test_internal_for_writing_rejects_readonly_buffer
+    buffer = IO::Buffer.for("hello")
+
+    assert_raise(ArgumentError) do
+      Bug::IOBuffer.for_writing_set_string(buffer, "world")
+    end
+  end
+
+  def test_internal_for_writing_unlocks_after_callback_exception
+    string = "hello"
+
+    assert_raise(RuntimeError) do
+      Bug::IOBuffer.for_writing_modify_string(string)
+    end
+
+    string << "!"
+    assert_equal "hello!", string
+  end
+
+  def test_internal_for_reading_unlocks_after_callback_exception
+    string = "hello"
+
+    assert_raise(RuntimeError) do
+      Bug::IOBuffer.for_reading_raise(string)
+    end
+
+    string << "!"
+    assert_equal "hello!", string
+  end
+
   def test_endian
     assert_equal 4, IO::Buffer::LITTLE_ENDIAN
     assert_equal 8, IO::Buffer::BIG_ENDIAN