marshal.c Marshal.load accepts a freeze: true option.

Fixes [Feature #18148]

When set, all the loaded objects are returned as frozen.

If a proc is provided, it is called with the objects already frozen.

7 files changed, 184 insertions, 30 deletions

diff --git a/.document b/.document
index 2c68af227d..6e08f42698 100644
--- a/.document
+++ b/.document
@@ -17,6 +17,7 @@ dir.rb
 gc.rb
 io.rb
 kernel.rb
+marshal.rb
 numeric.rb
 nilclass.rb
 pack.rb
diff --git a/common.mk b/common.mk
index fd15d429b9..2a582175bd 100644
--- a/common.mk
+++ b/common.mk
@@ -1047,6 +1047,7 @@ BUILTIN_RB_SRCS = \
 		$(srcdir)/gc.rb \
 		$(srcdir)/numeric.rb \
 		$(srcdir)/io.rb \
+		$(srcdir)/marshal.rb \
 		$(srcdir)/pack.rb \
 		$(srcdir)/trace_point.rb \
 		$(srcdir)/warning.rb \
@@ -7732,6 +7733,7 @@ marshal.$(OBJEXT): {$(VPATH)}backward/2/limits.h
 marshal.$(OBJEXT): {$(VPATH)}backward/2/long_long.h
 marshal.$(OBJEXT): {$(VPATH)}backward/2/stdalign.h
 marshal.$(OBJEXT): {$(VPATH)}backward/2/stdarg.h
+marshal.$(OBJEXT): {$(VPATH)}builtin.h
 marshal.$(OBJEXT): {$(VPATH)}config.h
 marshal.$(OBJEXT): {$(VPATH)}defines.h
 marshal.$(OBJEXT): {$(VPATH)}encindex.h
@@ -7889,6 +7891,8 @@ marshal.$(OBJEXT): {$(VPATH)}internal/warning_push.h
 marshal.$(OBJEXT): {$(VPATH)}internal/xmalloc.h
 marshal.$(OBJEXT): {$(VPATH)}io.h
 marshal.$(OBJEXT): {$(VPATH)}marshal.c
+marshal.$(OBJEXT): {$(VPATH)}marshal.rb
+marshal.$(OBJEXT): {$(VPATH)}marshal.rbinc
 marshal.$(OBJEXT): {$(VPATH)}missing.h
 marshal.$(OBJEXT): {$(VPATH)}onigmo.h
 marshal.$(OBJEXT): {$(VPATH)}oniguruma.h
@@ -8422,6 +8426,7 @@ miniinit.$(OBJEXT): {$(VPATH)}internal/xmalloc.h
 miniinit.$(OBJEXT): {$(VPATH)}io.rb
 miniinit.$(OBJEXT): {$(VPATH)}iseq.h
 miniinit.$(OBJEXT): {$(VPATH)}kernel.rb
+miniinit.$(OBJEXT): {$(VPATH)}marshal.rb
 miniinit.$(OBJEXT): {$(VPATH)}method.h
 miniinit.$(OBJEXT): {$(VPATH)}mini_builtin.c
 miniinit.$(OBJEXT): {$(VPATH)}miniinit.c
diff --git a/inits.c b/inits.c
index 3e04c26111..f69ee73a89 100644
--- a/inits.c
+++ b/inits.c
@@ -98,6 +98,7 @@ rb_call_builtin_inits(void)
     BUILTIN(kernel);
     BUILTIN(timev);
     BUILTIN(nilclass);
+    BUILTIN(marshal);
     Init_builtin_prelude();
 }
 #undef CALL
diff --git a/marshal.c b/marshal.c
index d8fcf56685..0f746d9805 100644
--- a/marshal.c
+++ b/marshal.c
@@ -37,6 +37,7 @@
 #include "ruby/ruby.h"
 #include "ruby/st.h"
 #include "ruby/util.h"
+#include "builtin.h"
 
 #define BITSPERSHORT (2*CHAR_BIT)
 #define SHORTMASK ((1<<BITSPERSHORT)-1)
@@ -123,7 +124,7 @@ typedef struct {
 static st_table *compat_allocator_tbl;
 static VALUE compat_allocator_tbl_wrapper;
 static VALUE rb_marshal_dump_limited(VALUE obj, VALUE port, int limit);
-static VALUE rb_marshal_load_with_proc(VALUE port, VALUE proc);
+static VALUE rb_marshal_load_with_proc(VALUE port, VALUE proc, bool freeze);
 
 static int
 mark_marshal_compat_i(st_data_t key, st_data_t value, st_data_t _)
@@ -1164,6 +1165,7 @@ struct load_arg {
     st_table *partial_objects;
     VALUE proc;
     st_table *compat_tbl;
+    bool freeze;
 };
 
 static VALUE
@@ -1580,6 +1582,17 @@ r_leave(VALUE v, struct load_arg *arg, bool partial)
 	st_data_t data;
 	st_data_t key = (st_data_t)v;
 	st_delete(arg->partial_objects, &key, &data);
+	if (arg->freeze) {
+	    if (RB_TYPE_P(v, T_MODULE) || RB_TYPE_P(v, T_CLASS)) {
+		// noop
+	    }
+	    else if (RB_TYPE_P(v, T_STRING)) {
+		v = rb_str_to_interned_str(v);
+	    }
+	    else {
+		OBJ_FREEZE(v);
+	    }
+        }
 	v = r_post_proc(v, arg);
     }
     return v;
@@ -2191,33 +2204,8 @@ clear_load_arg(struct load_arg *arg)
     }
 }
 
-/*
- * call-seq:
- *     load( source [, proc] ) -> obj
- *     restore( source [, proc] ) -> obj
- *
- * Returns the result of converting the serialized data in source into a
- * Ruby object (possibly with associated subordinate objects). source
- * may be either an instance of IO or an object that responds to
- * to_str. If proc is specified, each object will be passed to the proc, as the object
- * is being deserialized.
- *
- * Never pass untrusted data (including user supplied input) to this method.
- * Please see the overview for further details.
- */
-static VALUE
-marshal_load(int argc, VALUE *argv, VALUE _)
-{
-    VALUE port, proc;
-
-    rb_check_arity(argc, 1, 2);
-    port = argv[0];
-    proc = argc > 1 ? argv[1] : Qnil;
-    return rb_marshal_load_with_proc(port, proc);
-}
-
 VALUE
-rb_marshal_load_with_proc(VALUE port, VALUE proc)
+rb_marshal_load_with_proc(VALUE port, VALUE proc, bool freeze)
 {
     int major, minor;
     VALUE v;
@@ -2243,6 +2231,7 @@ rb_marshal_load_with_proc(VALUE port, VALUE proc)
     arg->compat_tbl = 0;
     arg->proc = 0;
     arg->readable = 0;
+    arg->freeze = freeze;
 
     if (NIL_P(v))
 	arg->buf = xmalloc(BUFSIZ);
@@ -2271,6 +2260,13 @@ rb_marshal_load_with_proc(VALUE port, VALUE proc)
     return v;
 }
 
+static VALUE marshal_load(rb_execution_context_t *ec, VALUE mod, VALUE source, VALUE proc, VALUE freeze)
+{
+    return rb_marshal_load_with_proc(source, proc, RTEST(freeze));
+}
+
+#include "marshal.rbinc"
+
 /*
  * The marshaling library converts collections of Ruby objects into a
  * byte stream, allowing them to be stored outside the currently
@@ -2403,8 +2399,6 @@ Init_marshal(void)
     set_id(s_ruby2_keywords_flag);
 
     rb_define_module_function(rb_mMarshal, "dump", marshal_dump, -1);
-    rb_define_module_function(rb_mMarshal, "load", marshal_load, -1);
-    rb_define_module_function(rb_mMarshal, "restore", marshal_load, -1);
 
     /* major version */
     rb_define_const(rb_mMarshal, "MAJOR_VERSION", INT2FIX(MARSHAL_MAJOR));
@@ -2434,5 +2428,5 @@ rb_marshal_dump(VALUE obj, VALUE port)
 VALUE
 rb_marshal_load(VALUE port)
 {
-    return rb_marshal_load_with_proc(port, Qnil);
+    return rb_marshal_load_with_proc(port, Qnil, false);
 }
diff --git a/marshal.rb b/marshal.rb
new file mode 100644
index 0000000000..b8b5ce9e82
--- /dev/null
+++ b/marshal.rb
@@ -0,0 +1,21 @@
+module Marshal
+  # call-seq:
+  #     load( source [, proc] ) -> obj
+  #     restore( source [, proc] ) -> obj
+  #
+  # Returns the result of converting the serialized data in source into a
+  # Ruby object (possibly with associated subordinate objects). source
+  # may be either an instance of IO or an object that responds to
+  # to_str. If proc is specified, each object will be passed to the proc, as the object
+  # is being deserialized.
+  #
+  # Never pass untrusted data (including user supplied input) to this method.
+  # Please see the overview for further details.
+  def self.load(source, proc = nil, freeze: false)
+    Primitive.marshal_load(source, proc, freeze)
+  end
+
+  class << self
+    alias restore load
+  end
+end
diff --git a/spec/ruby/core/marshal/shared/load.rb b/spec/ruby/core/marshal/shared/load.rb
index 37c29d7cc6..e04a30c02f 100644
--- a/spec/ruby/core/marshal/shared/load.rb
+++ b/spec/ruby/core/marshal/shared/load.rb
@@ -19,6 +19,100 @@ describe :marshal_load, shared: true do
     -> { Marshal.send(@method, kaboom) }.should raise_error(ArgumentError)
   end
 
+  ruby_version_is "3.1" do
+    describe "when called with freeze: true" do
+      it "returns frozen strings" do
+        string = Marshal.send(@method, Marshal.dump("foo"), freeze: true)
+        string.should == "foo"
+        string.should.frozen?
+
+        utf8_string = "foo".encode(Encoding::UTF_8)
+        string = Marshal.send(@method, Marshal.dump(utf8_string), freeze: true)
+        string.should == utf8_string
+        string.should.frozen?
+      end
+
+      it "returns frozen arrays" do
+        array = Marshal.send(@method, Marshal.dump([1, 2, 3]), freeze: true)
+        array.should == [1, 2, 3]
+        array.should.frozen?
+      end
+
+      it "returns frozen hashes" do
+        hash = Marshal.send(@method, Marshal.dump({foo: 42}), freeze: true)
+        hash.should == {foo: 42}
+        hash.should.frozen?
+      end
+
+      it "returns frozen regexps" do
+        regexp = Marshal.send(@method, Marshal.dump(/foo/), freeze: true)
+        regexp.should == /foo/
+        regexp.should.frozen?
+      end
+
+      it "returns frozen objects" do
+        source_object = Object.new
+        source_object.instance_variable_set(:@foo, "bar")
+
+        object = Marshal.send(@method, Marshal.dump(source_object), freeze: true)
+        object.should.frozen?
+        object.instance_variable_get(:@foo).should.frozen?
+      end
+
+      it "does not freeze modules" do
+        Marshal.send(@method, Marshal.dump(Kernel), freeze: true)
+        Kernel.should_not.frozen?
+      end
+
+      it "does not freeze classes" do
+        Marshal.send(@method, Marshal.dump(Object), freeze: true)
+        Object.should_not.frozen?
+      end
+
+      describe "when called with a proc" do
+        it "call the proc with frozen objects" 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,
+            freeze: true,
+          )
+
+          arr.should == [
+            false, 5, "hi", 10, "hi", "hi", 0.0, false, "none", st,
+            :b, :c, 2, a, false, "ant", ["hi", 10, "hi", "hi", st, [:a, :b, :c]],
+          ]
+
+          arr.each do |obj|
+            obj.should.frozen?
+          end
+
+          Struct.send(:remove_const, :Brittle)
+        end
+
+        it "does not freeze the object returned by the proc" do
+          string = Marshal.send(@method, Marshal.dump("foo"), proc { |o| o.upcase }, freeze: true)
+          string.should == "FOO"
+          string.should_not.frozen?
+        end
+      end
+    end
+  end
+
   describe "when called with a proc" do
     ruby_bug "#18141", ""..."3.1" do
       it "call the proc with fully initialized strings" do
diff --git a/test/ruby/test_marshal.rb b/test/ruby/test_marshal.rb
index 4f25344bf6..19f41de27e 100644
--- a/test/ruby/test_marshal.rb
+++ b/test/ruby/test_marshal.rb
@@ -889,4 +889,42 @@ class TestMarshal < Test::Unit::TestCase
   def test_hash_default_compared_by_identity
     _test_hash_compared_by_identity(Hash.new(true))
   end
+
+  class TestMarshalFreeze < Test::Unit::TestCase
+    include MarshalTestLib
+
+    def encode(o)
+      Marshal.dump(o)
+    end
+
+    def decode(s)
+      Marshal.load(s, freeze: true)
+    end
+
+    def test_return_objects_are_frozen
+      source = ["foo", {}, /foo/, 1..2]
+      objects = decode(encode(source))
+      assert_equal source, objects
+      assert_predicate objects, :frozen?
+      objects.each do |obj|
+        assert_predicate obj, :frozen?
+      end
+    end
+
+    def test_proc_returned_object_are_not_frozen
+      source = ["foo", {}, /foo/, 1..2]
+      objects = Marshal.load(encode(source), ->(o) { o.dup }, freeze: true)
+      assert_equal source, objects
+      refute_predicate objects, :frozen?
+      objects.each do |obj|
+        refute_predicate obj, :frozen?
+      end
+    end
+
+    def test_modules_and_classes_are_not_frozen
+      objects = Marshal.load(encode([Object, Kernel]), freeze: true)
+      refute_predicate Object, :frozen?
+      refute_predicate Kernel, :frozen?
+    end
+  end
 end