Marshal.load: do not call the proc until strings have their encoding

Ref: https://bugs.ruby-lang.org/issues/18141

3 files changed, 64 insertions, 22 deletions

diff --git a/marshal.c b/marshal.c
index ea43c62b94..a817cf9742 100644
--- a/marshal.c
+++ b/marshal.c
@@ -1690,6 +1690,9 @@ r_object0(struct load_arg *arg, int *ivp, VALUE extmod)
 
 	    v = r_object0(arg, &ivar, extmod);
 	    if (ivar) r_ivar(v, NULL, arg);
+	    if (RB_TYPE_P(v, T_STRING)) {
+	        v = r_leave(v, arg);
+	    }
 	}
 	break;
 
@@ -1819,7 +1822,9 @@ r_object0(struct load_arg *arg, int *ivp, VALUE extmod)
 
       case TYPE_STRING:
 	v = r_entry(r_string(arg), arg);
-        v = r_leave(v, arg);
+	if (!ivp) {
+	    v = r_leave(v, arg);
+	}
 	break;
 
       case TYPE_REGEXP:
diff --git a/spec/ruby/core/marshal/shared/load.rb b/spec/ruby/core/marshal/shared/load.rb
index a79fb51cd7..18f11d70d6 100644
--- a/spec/ruby/core/marshal/shared/load.rb
+++ b/spec/ruby/core/marshal/shared/load.rb
@@ -20,6 +20,24 @@ describe :marshal_load, shared: true do
   end
 
   describe "when called with a proc" do
+    ruby_bug "#18141", ""..."3.1" do
+      it "call the proc with fully initialized strings" do
+        utf8_string = "foo".encode(Encoding::UTF_8)
+        Marshal.send(@method, Marshal.dump(utf8_string), proc { |arg|
+          if arg.is_a?(String)
+            arg.should == utf8_string
+            arg.encoding.should == Encoding::UTF_8
+          end
+          arg
+        })
+      end
+
+      it "no longer mutate the object after it was passed to the proc" do
+        string = Marshal.load(Marshal.dump("foo"), :freeze.to_proc)
+        string.should.frozen?
+      end
+    end
+
     it "returns the value of the proc" do
       Marshal.send(@method, Marshal.dump([1,2]), proc { [3,4] }).should ==  [3,4]
     end
@@ -35,27 +53,29 @@ describe :marshal_load, shared: true do
       ret.size.should == 3
     end
 
-    it "loads an Array with proc" do
-      arr = []
-      s = 'hi'
-      s.instance_variable_set(:@foo, 5)
-      st = Struct.new("Brittle", :a).new
-      st.instance_variable_set(:@clue, 'none')
-      st.a = 0.0
-      h = Hash.new('def')
-      h['nine'] = 9
-      a = [:a, :b, :c]
-      a.instance_variable_set(:@two, 2)
-      obj = [s, 10, s, s, st, a]
-      obj.instance_variable_set(:@zoo, 'ant')
-      proc = Proc.new { |o| arr << o; o}
-
-      Marshal.send(@method, "\x04\bI[\vI\"\ahi\a:\x06EF:\t@fooi\ni\x0F@\x06@\x06IS:\x14Struct::Brittle\x06:\x06af\x060\x06:\n@clueI\"\tnone\x06;\x00FI[\b;\b:\x06b:\x06c\x06:\t@twoi\a\x06:\t@zooI\"\bant\x06;\x00F", proc)
-
-      arr.should == ["hi", false, 5, 10, "hi", "hi", 0.0, st, "none", false,
-                     :b, :c, a, 2, ["hi", 10, "hi", "hi", st, [:a, :b, :c]], "ant", false]
-
-      Struct.send(:remove_const, :Brittle)
+    ruby_bug "#18141", ""..."3.1" do
+      it "loads an Array with proc" do
+        arr = []
+        s = 'hi'
+        s.instance_variable_set(:@foo, 5)
+        st = Struct.new("Brittle", :a).new
+        st.instance_variable_set(:@clue, 'none')
+        st.a = 0.0
+        h = Hash.new('def')
+        h['nine'] = 9
+        a = [:a, :b, :c]
+        a.instance_variable_set(:@two, 2)
+        obj = [s, 10, s, s, st, a]
+        obj.instance_variable_set(:@zoo, 'ant')
+        proc = Proc.new { |o| arr << o; o}
+
+        Marshal.send(@method, "\x04\bI[\vI\"\ahi\a:\x06EF:\t@fooi\ni\x0F@\x06@\x06IS:\x14Struct::Brittle\x06:\x06af\x060\x06:\n@clueI\"\tnone\x06;\x00FI[\b;\b:\x06b:\x06c\x06:\t@twoi\a\x06:\t@zooI\"\bant\x06;\x00F", proc)
+
+        arr.should == [false, 5, "hi", 10, "hi", "hi", 0.0, st, false, "none",
+                       :b, :c, a, 2, ["hi", 10, "hi", "hi", st, [:a, :b, :c]], false, "ant"]
+
+        Struct.send(:remove_const, :Brittle)
+      end
     end
   end
 
diff --git a/test/ruby/test_marshal.rb b/test/ruby/test_marshal.rb
index 5e6891a6a0..a0d9d96a1f 100644
--- a/test/ruby/test_marshal.rb
+++ b/test/ruby/test_marshal.rb
@@ -676,6 +676,23 @@ class TestMarshal < Test::Unit::TestCase
     assert_equal(['X', 'X'], Marshal.load(Marshal.dump(obj), ->(v) { v == str ? v.upcase : v }))
   end
 
+  def test_marshal_proc_string_encoding
+    string = "foo"
+    payload = Marshal.dump(string)
+    Marshal.load(payload, ->(v) {
+      if v.is_a?(String)
+        assert_equal(string, v)
+        assert_equal(string.encoding, v.encoding)
+      end
+      v
+    })
+  end
+
+  def test_marshal_proc_freeze
+    object = { foo: [42, "bar"] }
+    assert_equal object, Marshal.load(Marshal.dump(object), :freeze.to_proc)
+  end
+
   def test_marshal_load_extended_class_crash
     assert_separately([], "#{<<-"begin;"}\n#{<<-"end;"}")
     begin;