Don't overwrite shape capacity when removing ivar

Other objects may be using the shape, so we can't change the capacity
otherwise the other objects may have a buffer overflow.

2 files changed, 35 insertions, 1 deletions

diff --git a/shape.c b/shape.c
index 13c8d11b82..588c24604d 100644
--- a/shape.c
+++ b/shape.c
@@ -597,7 +597,8 @@ remove_shape_recursive(VALUE obj, ID id, rb_shape_t * shape, VALUE * removed)
                     return new_child;
                 }
 
-                new_child->capacity = shape->capacity;
+                RUBY_ASSERT(new_child->capacity <= shape->capacity);
+
                 if (new_child->type == SHAPE_IVAR) {
                     move_iv(obj, id, shape->next_iv_index - 1, new_child->next_iv_index - 1);
                 }
diff --git a/test/ruby/test_shapes.rb b/test/ruby/test_shapes.rb
index 8755da9d25..9f455e4389 100644
--- a/test/ruby/test_shapes.rb
+++ b/test/ruby/test_shapes.rb
@@ -695,6 +695,39 @@ class TestShapes < Test::Unit::TestCase
     end;
   end
 
+  def test_remove_instance_variable_capacity_transition
+    assert_separately([], "#{<<~"begin;"}\n#{<<~'end;'}")
+    begin;
+      t_object_shape = RubyVM::Shape.find_by_id(GC::INTERNAL_CONSTANTS[:SIZE_POOL_COUNT])
+      assert_equal(RubyVM::Shape::SHAPE_T_OBJECT, t_object_shape.type)
+
+      initial_capacity = t_object_shape.capacity
+
+      # a does not transition in capacity
+      a = Class.new.new
+      initial_capacity.times do |i|
+        a.instance_variable_set(:"@ivar#{i + 1}", i)
+      end
+
+      # b transitions in capacity
+      b = Class.new.new
+      (initial_capacity + 1).times do |i|
+        b.instance_variable_set(:"@ivar#{i}", i)
+      end
+
+      assert_operator(RubyVM::Shape.of(a).capacity, :<, RubyVM::Shape.of(b).capacity)
+
+      # b will now have the same tree as a
+      b.remove_instance_variable(:@ivar0)
+
+      a.instance_variable_set(:@foo, 1)
+      a.instance_variable_set(:@bar, 1)
+
+      # Check that there is no heap corruption
+      GC.verify_internal_consistency
+    end;
+  end
+
   def test_freeze_after_complex
     ensure_complex