From c9423b016cfeab852bc5a829e55e0a11f80b3ab7 Mon Sep 17 00:00:00 2001
From: Nobuyoshi Nakada <nobu@ruby-lang.org>
Date: Mon, 1 Jul 2019 15:02:27 +0900
Subject: [PATCH] marshal.c: check instance variable count

* marshal.c (w_obj_each): ensure that no instance variable was
  added while dumping other instance variables.  [Bug #15968]
---
 marshal.c                 | 20 +++++++++++++++++---
 test/ruby/test_marshal.rb | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/marshal.c b/marshal.c
index 852305110e..a1ee01e2ca 100644
--- a/marshal.c
+++ b/marshal.c
@@ -561,14 +561,25 @@ w_uclass(VALUE obj, VALUE super, struct dump_arg *arg)
 
 #define to_be_skipped_id(id) (id == rb_id_encoding() || id == rb_intern("E") || !rb_id2str(id))
 
+struct w_ivar_arg {
+    struct dump_call_arg *dump;
+    st_data_t num_ivar;
+};
+
 static int
 w_obj_each(st_data_t key, st_data_t val, st_data_t a)
 {
     ID id = (ID)key;
     VALUE value = (VALUE)val;
-    struct dump_call_arg *arg = (struct dump_call_arg *)a;
+    struct w_ivar_arg *ivarg = (struct w_ivar_arg *)a;
+    struct dump_call_arg *arg = ivarg->dump;
 
     if (to_be_skipped_id(id)) return ST_CONTINUE;
+    if (!ivarg->num_ivar) {
+        rb_raise(rb_eRuntimeError, "instance variable added to %"PRIsVALUE" instance",
+                 CLASS_OF(arg->obj));
+    }
+    --ivarg->num_ivar;
     w_symbol(ID2SYM(id), arg->arg);
     w_object(value, arg->arg, arg->limit);
     return ST_CONTINUE;
@@ -659,7 +670,8 @@ w_ivar(st_index_t num, VALUE ivobj, VALUE encname, struct dump_call_arg *arg)
     w_long(num, arg->arg);
     w_encoding(encname, arg);
     if (ivobj != Qundef) {
-	rb_ivar_foreach(ivobj, w_obj_each, (st_data_t)arg);
+        struct w_ivar_arg ivarg = {arg, num};
+        rb_ivar_foreach(ivobj, w_obj_each, (st_data_t)&ivarg);
     }
 }
 
@@ -671,7 +683,8 @@ w_objivar(VALUE obj, struct dump_call_arg *arg)
     rb_ivar_foreach(obj, obj_count_ivars, (st_data_t)&num);
     w_long(num, arg->arg);
     if (num != 0) {
-        rb_ivar_foreach(obj, w_obj_each, (st_data_t)arg);
+        struct w_ivar_arg ivarg = {arg, num};
+        rb_ivar_foreach(obj, w_obj_each, (st_data_t)&ivarg);
     }
 }
 
@@ -691,6 +704,7 @@ w_object(VALUE obj, struct dump_arg *arg, int limit)
     if (limit > 0) limit--;
     c_arg.limit = limit;
     c_arg.arg = arg;
+    c_arg.obj = obj;
 
     if (st_lookup(arg->data, obj, &num)) {
 	w_byte(TYPE_LINK, arg);
diff --git a/test/ruby/test_marshal.rb b/test/ruby/test_marshal.rb
index e269428dda..0565a1c04f 100644
--- a/test/ruby/test_marshal.rb
+++ b/test/ruby/test_marshal.rb
@@ -779,4 +779,36 @@ class TestMarshal < Test::Unit::TestCase
     obj = Bug14314.new(foo: 42)
     assert_equal obj, Marshal.load(Marshal.dump(obj))
   end
+
+  class Bug15968
+    attr_accessor :bar, :baz
+
+    def initialize
+      self.bar = Bar.new(self)
+    end
+
+    class Bar
+      attr_accessor :foo
+
+      def initialize(foo)
+        self.foo = foo
+      end
+
+      def marshal_dump
+        self.foo.baz = :problem
+        {foo: self.foo}
+      end
+
+      def marshal_load(data)
+        self.foo = data[:foo]
+      end
+    end
+  end
+
+  def test_marshal_dump_adding_instance_variable
+    obj = Bug15968.new
+    assert_raise_with_message(RuntimeError, /instance variable added/) do
+      Marshal.dump(obj)
+    end
+  end
 end