Add support of Marshal.load :freeze option

andrykonchin · eregon · commit 77c68a09baa6 · 2023-02-28T20:25:29.000+01:00
diff --git a/spec/ruby/core/marshal/fixtures/marshal_data.rb b/spec/ruby/core/marshal/fixtures/marshal_data.rb
@@ -183,12 +183,17 @@ def initialize(*args)
     end
   end
 
+  StructToDump = Struct.new(:a, :b)
+
   class BasicObjectSubWithRespondToFalse < BasicObject
     def respond_to?(method_name, include_all=false)
       false
     end
   end
 
+  module ModuleToExtendBy
+  end
+
   def self.random_data
     randomizer = Random.new(42)
     1000.times{randomizer.rand} # Make sure we exhaust his first state of 624 random words
@@ -273,6 +278,16 @@ class MultibyteぁあぃいTime < Time
     end
   ruby
 
+  class ObjectWithFreezeRaisingException < Object
+    def freeze
+      raise
+    end
+  end
+
+  class ObjectWithoutFreeze < Object
+    undef freeze
+  end
+
   DATA = {
     "nil" => [nil, "\004\b0"],
     "1..2" => [(1..2),
diff --git a/spec/ruby/core/marshal/shared/load.rb b/spec/ruby/core/marshal/shared/load.rb
@@ -54,22 +54,77 @@
         regexp.should.frozen?
       end
 
+      it "returns frozen structs" do
+        struct = Marshal.send(@method, Marshal.dump(MarshalSpec::StructToDump.new(1, 2)), freeze: true)
+        struct.should == MarshalSpec::StructToDump.new(1, 2)
+        struct.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
+
+      describe "deep freezing" do
+        it "returns hashes with frozen keys and values" do
+          key = Object.new
+          value = Object.new
+          source_object = {key => value}
+
+          hash = Marshal.send(@method, Marshal.dump(source_object), freeze: true)
+          hash.size.should == 1
+          hash.keys[0].should.frozen?
+          hash.values[0].should.frozen?
+        end
+
+        it "returns arrays with frozen elements" do
+          object = Object.new
+          source_object = [object]
+
+          array = Marshal.send(@method, Marshal.dump(source_object), freeze: true)
+          array.size.should == 1
+          array[0].should.frozen?
+        end
+
+        it "returns structs with frozen members" do
+          object1 = Object.new
+          object2 = Object.new
+          source_object = MarshalSpec::StructToDump.new(object1, object2)
+
+          struct = Marshal.send(@method, Marshal.dump(source_object), freeze: true)
+          struct.a.should.frozen?
+          struct.b.should.frozen?
+        end
+
+        it "returns objects with frozen instance variables" do
+          source_object = Object.new
+          instance_variable = Object.new
+          source_object.instance_variable_set(:@a, instance_variable)
+
+          object = Marshal.send(@method, Marshal.dump(source_object), freeze: true)
+          object.instance_variable_get(:@a).should != nil
+          object.instance_variable_get(:@a).should.frozen?
+        end
+
+        it "deduplicates frozen strings" do
+          source_object = ["foo" + "bar", "foobar"]
+          object = Marshal.send(@method, Marshal.dump(source_object), freeze: true)
+
+          object[0].should equal(object[1])
+        end
       end
 
       it "does not freeze modules" do
-        Marshal.send(@method, Marshal.dump(Kernel), freeze: true)
+        object = Marshal.send(@method, Marshal.dump(Kernel), freeze: true)
+        object.should_not.frozen?
         Kernel.should_not.frozen?
       end
 
       it "does not freeze classes" do
-        Marshal.send(@method, Marshal.dump(Object), freeze: true)
+        object = Marshal.send(@method, Marshal.dump(Object), freeze: true)
+        object.should_not.frozen?
         Object.should_not.frozen?
       end
 
@@ -85,6 +140,46 @@
         end
       end
 
+      ruby_bug "#19427", "3.1"..."3.3" do
+        it "returns frozen object having #_dump method" do
+          object = Marshal.send(@method, Marshal.dump(UserDefined.new), freeze: true)
+          object.should.frozen?
+        end
+
+        it "returns frozen object responding to #marshal_dump and #marshal_load" do
+          object = Marshal.send(@method, Marshal.dump(UserMarshal.new), freeze: true)
+          object.should.frozen?
+        end
+
+        it "returns frozen object extended by a module" do
+          object = Object.new
+          object.extend(MarshalSpec::ModuleToExtendBy)
+
+          object = Marshal.send(@method, Marshal.dump(object), freeze: true)
+          object.should.frozen?
+        end
+      end
+
+      it "does not call freeze method" do
+        object = MarshalSpec::ObjectWithFreezeRaisingException.new
+        object = Marshal.send(@method, Marshal.dump(object), freeze: true)
+        object.should.frozen?
+      end
+
+      it "returns frozen object even if object does not respond to freeze method" do
+        object = MarshalSpec::ObjectWithoutFreeze.new
+        object = Marshal.send(@method, Marshal.dump(object), freeze: true)
+        object.should.frozen?
+      end
+
+      it "returns a frozen object when is an instance of String/Array/Regexp/Hash subclass and has instance variables" do
+        source_object = UserString.new
+        source_object.instance_variable_set(:@foo, "bar")
+
+        object = Marshal.send(@method, Marshal.dump(source_object), freeze: true)
+        object.should.frozen?
+      end
+
       describe "when called with a proc" do
         it "call the proc with frozen objects" do
           arr = []
diff --git a/spec/tags/core/marshal/load_tags.txt b/spec/tags/core/marshal/load_tags.txt
@@ -1,5 +1,4 @@
 fails:Marshal.load loads a Random
-fails:Marshal.load when called with a proc returns the value of the proc
 fails:Marshal.load when called with a proc loads an Array with proc
 fails:Marshal.load for an Array loads an array containing the same objects
 fails:Marshal.load for an Object raises ArgumentError if the object from an 'o' stream is not dumpable as 'o' type user class
@@ -8,13 +7,6 @@ fails:Marshal.load for a wrapped C pointer raises ArgumentError when the local c
 fails:Marshal.load when a class does not exist in the namespace raises an ArgumentError
 fails(immutable regexp):Marshal.load for a Regexp loads an extended Regexp
 fails(immutable regexp):Marshal.load for a Regexp loads a extended_user_regexp having ivar
-fails:Marshal.load when called with freeze: true returns frozen strings
-fails:Marshal.load when called with freeze: true returns frozen arrays
-fails:Marshal.load when called with freeze: true returns frozen hashes
-fails:Marshal.load when called with freeze: true returns frozen regexps
-fails:Marshal.load when called with freeze: true returns frozen objects
-fails:Marshal.load when called with freeze: true does not freeze modules
-fails:Marshal.load when called with freeze: true does not freeze classes
 fails:Marshal.load when called with freeze: true when called with a proc call the proc with frozen objects
 fails:Marshal.load when called with freeze: true when called with a proc does not freeze the object returned by the proc
 fails:Marshal.load when called with freeze: true does freeze extended objects
diff --git a/spec/tags/core/marshal/restore_tags.txt b/spec/tags/core/marshal/restore_tags.txt
@@ -1,5 +1,4 @@
 fails:Marshal.restore loads a Random
-fails:Marshal.restore when called with a proc returns the value of the proc
 fails:Marshal.restore when called with a proc loads an Array with proc
 fails:Marshal.restore for an Array loads an array containing the same objects
 fails:Marshal.restore for an Object raises ArgumentError if the object from an 'o' stream is not dumpable as 'o' type user class
@@ -8,13 +7,6 @@ fails:Marshal.restore for a wrapped C pointer raises ArgumentError when the loca
 fails:Marshal.restore when a class does not exist in the namespace raises an ArgumentError
 fails(immutable regexp):Marshal.restore for a Regexp loads an extended Regexp
 fails(immutable regexp):Marshal.restore for a Regexp loads a extended_user_regexp having ivar
-fails:Marshal.restore when called with freeze: true returns frozen strings
-fails:Marshal.restore when called with freeze: true returns frozen arrays
-fails:Marshal.restore when called with freeze: true returns frozen hashes
-fails:Marshal.restore when called with freeze: true returns frozen regexps
-fails:Marshal.restore when called with freeze: true returns frozen objects
-fails:Marshal.restore when called with freeze: true does not freeze modules
-fails:Marshal.restore when called with freeze: true does not freeze classes
 fails:Marshal.restore when called with freeze: true when called with a proc call the proc with frozen objects
 fails:Marshal.restore when called with freeze: true when called with a proc does not freeze the object returned by the proc
 fails:Marshal.restore when called with freeze: true does freeze extended objects
diff --git a/src/main/java/org/truffleruby/core/support/TypeNodes.java b/src/main/java/org/truffleruby/core/support/TypeNodes.java
@@ -16,7 +16,6 @@
 import com.oracle.truffle.api.dsl.GenerateUncached;
 import com.oracle.truffle.api.dsl.NeverDefault;
 import com.oracle.truffle.api.profiles.ConditionProfile;
-import com.oracle.truffle.api.strings.TruffleString;
 import org.truffleruby.annotations.CoreModule;
 import org.truffleruby.annotations.Primitive;
 import org.truffleruby.builtins.PrimitiveArrayArgumentsNode;
@@ -28,7 +27,6 @@
 import org.truffleruby.core.cast.ToIntNode;
 import org.truffleruby.core.cast.ToLongNode;
 import org.truffleruby.core.cast.ToRubyIntegerNode;
-import org.truffleruby.core.encoding.Encodings;
 import org.truffleruby.core.kernel.KernelNodes;
 import org.truffleruby.core.kernel.KernelNodes.ToSNode;
 import org.truffleruby.core.klass.RubyClass;
@@ -39,6 +37,7 @@
 import org.truffleruby.language.Nil;
 import org.truffleruby.language.NotProvided;
 import org.truffleruby.language.RubyDynamicObject;
+import org.truffleruby.language.RubyGuards;
 import org.truffleruby.language.RubyNode;
 import org.truffleruby.language.RubySourceNode;
 import org.truffleruby.language.control.RaiseException;
@@ -122,6 +121,33 @@ protected Object freeze(Object self,
         }
     }
 
+    @Primitive(name = "object_freeze_with_singleton_class")
+    public abstract static class ObjectFreezeWithSingletonClassNode extends PrimitiveArrayArgumentsNode {
+        @Specialization(limit = "getRubyLibraryCacheLimit()", guards = "!isRubyDynamicObject(self)")
+        protected Object freeze(Object self,
+                @CachedLibrary("self") RubyLibrary rubyLibrary) {
+            rubyLibrary.freeze(self);
+            return self;
+        }
+
+        @Specialization(limit = "getRubyLibraryCacheLimit()", guards = "isRubyDynamicObject(self)")
+        protected Object freezeDynamicObject(Object self,
+                @CachedLibrary("self") RubyLibrary rubyLibrary,
+                @CachedLibrary(limit = "1") RubyLibrary rubyLibraryMetaClass,
+                @Cached ConditionProfile singletonClassUnfrozenProfile,
+                @Cached MetaClassNode metaClassNode) {
+            final RubyClass metaClass = metaClassNode.execute(self);
+            if (singletonClassUnfrozenProfile.profile(metaClass.isSingleton &&
+                    !(RubyGuards.isRubyClass(self) && ((RubyClass) self).isSingleton) &&
+                    !rubyLibraryMetaClass.isFrozen(metaClass))) {
+                rubyLibraryMetaClass.freeze(metaClass);
+            }
+            rubyLibrary.freeze(self);
+            return self;
+        }
+
+    }
+
     @Primitive(name = "immediate_value?")
     public abstract static class IsImmediateValueNode extends PrimitiveArrayArgumentsNode {
 
diff --git a/src/main/ruby/truffleruby/core/marshal.rb b/src/main/ruby/truffleruby/core/marshal.rb
@@ -555,7 +555,7 @@ module Marshal
 
   class State
 
-    def initialize(stream, depth, proc)
+    def initialize(stream, depth, proc, freeze)
       # shared
       @links = {}
       @symlinks = {}
@@ -576,6 +576,7 @@ def initialize(stream, depth, proc)
       @modules = nil
       @has_ivar = []
       @proc = proc
+      @freeze = freeze
       @call = true
       @user_classes = nil
     end
@@ -633,7 +634,7 @@ def assign_reserved_symlink(sz, obj)
       @symlinks[obj.__id__] = sz
     end
 
-    def construct(ivar_index = nil, call_proc = true)
+    def construct(ivar_index = nil, call_proc = true, postpone_freezing = false)
       type = consume_byte()
       obj = case type
             when 48   # ?0
@@ -699,7 +700,7 @@ def construct(ivar_index = nil, call_proc = true)
               name = get_symbol
               @modules << const_lookup(name, Module)
 
-              obj = construct nil, false
+              obj = construct nil, false, true
 
               extend_object obj
 
@@ -711,13 +712,13 @@ def construct(ivar_index = nil, call_proc = true)
               @user_classes ||= []
               @user_classes << name
 
-              construct nil, false
+              construct nil, false, postpone_freezing
 
             when 73   # ?I
               ivar_index = @has_ivar.length
               @has_ivar.push true
 
-              obj = construct ivar_index, false
+              obj = construct ivar_index, false, true
 
               set_instance_variables obj if @has_ivar.pop
 
@@ -726,7 +727,12 @@ def construct(ivar_index = nil, call_proc = true)
               raise ArgumentError, "load error, unknown type #{type}"
             end
 
-      @proc.call(obj) if call_proc and @proc and @call
+      if @freeze && !postpone_freezing && !(Primitive.object_kind_of?(obj, Class) || Primitive.object_kind_of?(obj, Module))
+        obj = -obj if Primitive.class_of(obj) == String
+        Primitive.object_freeze_with_singleton_class(obj)
+      end
+
+      return @proc.call(obj) if call_proc and @proc and @call
 
       obj
     end
@@ -1388,8 +1394,8 @@ def consume_byte
   end
 
   class StringState < State
-    def initialize(stream, depth, prc)
-      super stream, depth, prc
+    def initialize(stream, depth, prc, freeze)
+      super stream, depth, prc, freeze
 
       if @stream
         @byte_array = stream.bytes
@@ -1426,7 +1432,7 @@ def self.dump(obj, an_io=nil, limit=nil)
     end
 
     depth = Primitive.rb_to_int limit
-    ms = State.new nil, depth, nil
+    ms = State.new nil, depth, nil, nil
 
     if an_io
       unless Primitive.object_respond_to? an_io, :write, false
@@ -1447,16 +1453,16 @@ def self.dump(obj, an_io=nil, limit=nil)
     end
   end
 
-  def self.load(obj, prc = nil)
+  def self.load(obj, prc = nil, freeze: false)
     if Primitive.object_respond_to? obj, :to_str, false
       data = obj.to_s
-      ms = StringState.new data, nil, prc
+      ms = StringState.new data, nil, prc, freeze
 
       major = ms.consume_byte
       minor = ms.consume_byte
     elsif Primitive.object_respond_to? obj, :read, false and
           Primitive.object_respond_to? obj, :getc, false
-      ms = IOState.new obj, nil, prc
+      ms = IOState.new obj, nil, prc, freeze
 
       major = ms.consume_byte
       minor = ms.consume_byte
