Merge pull request #7296 from rubygems/segiddins/implement-safemarshal-support-for-hash-compare-by-identity

Implement SafeMarshal support for hash compare by identity

(cherry picked from commit 57741e8)

Author: deivid-rodriguez

lib/rubygems/safe_marshal/elements.rb

@@ -133,6 +133,14 @@ def initialize(sign, data)
         end
         attr_reader :sign, :data
       end
+
+      class UserClass < Element
+        def initialize(name, wrapped_object)
+          @name = name
+          @wrapped_object = wrapped_object
+        end
+        attr_reader :name, :wrapped_object
+      end
     end
   end
 end

lib/rubygems/safe_marshal/reader.rb

@@ -299,7 +299,9 @@ def read_struct
       end
 
       def read_user_class
-        raise NotImplementedError, "Reading Marshal objects of type user_class is not implemented"
+        name = read_element
+        wrapped_object = read_element
+        Elements::UserClass.new(name, wrapped_object)
       end
     end
   end

lib/rubygems/safe_marshal/visitors/to_ruby.rb

@@ -247,6 +247,30 @@ def visit_Gem_SafeMarshal_Elements_Bignum(b)
         end
       end
 
+      def visit_Gem_SafeMarshal_Elements_UserClass(r)
+        if resolve_class(r.name) == ::Hash && r.wrapped_object.is_a?(Elements::Hash)
+
+          hash = register_object({}.compare_by_identity)
+
+          o = r.wrapped_object
+          o.pairs.each_with_index do |(k, v), i|
+            push_stack i
+            k = visit(k)
+            push_stack k
+            hash[k] = visit(v)
+          end
+
+          if o.is_a?(Elements::HashWithDefaultValue)
+            push_stack :default
+            hash.default = visit(o.default)
+          end
+
+          hash
+        else
+          raise UnsupportedError.new("Unsupported user class #{resolve_class(r.name)} in marshal stream", stack: formatted_stack)
+        end
+      end
+
       def resolve_class(n)
         @class_cache[n] ||= begin
           to_s = resolve_symbol_name(n)
@@ -375,6 +399,12 @@ def initialize(name:, stack:)
         end
       end
 
+      class UnsupportedError < Error
+        def initialize(message, stack:)
+          super "#{message} @ #{stack.join "."}"
+        end
+      end
+
       class FormatError < Error
       end
 

test/rubygems/test_gem_safe_marshal.rb

@@ -247,9 +247,41 @@ def test_hash_with_default_value
   end
 
   def test_hash_with_compare_by_identity
-    pend "`read_user_class` not yet implemented"
+    with_const(Gem::SafeMarshal, :PERMITTED_CLASSES, %w[Hash]) do
+      assert_safe_load_as Hash.new.compare_by_identity.tap {|h|
+                            h[+"a"] = 1
+                            h[+"a"] = 2 }, additional_methods: [:compare_by_identity?], equality: false
+      assert_safe_load_as Hash.new.compare_by_identity, additional_methods: [:compare_by_identity?]
+      assert_safe_load_as Hash.new(0).compare_by_identity.tap {|h|
+                            h[+"a"] = 1
+                            h[+"a"] = 2 }, additional_methods: [:compare_by_identity?, :default], equality: false
+    end
+  end
+
+  class StringSubclass < ::String
+  end
 
-    assert_safe_load_as Hash.new.compare_by_identity
+  def test_string_subclass
+    with_const(Gem::SafeMarshal, :PERMITTED_CLASSES, [StringSubclass.name]) do
+      with_const(Gem::SafeMarshal, :PERMITTED_IVARS, { StringSubclass.name => %w[E] }) do
+        e = assert_raise(Gem::SafeMarshal::Visitors::ToRuby::UnsupportedError) do
+          Gem::SafeMarshal.safe_load Marshal.dump StringSubclass.new("abc")
+        end
+        assert_equal "Unsupported user class #{StringSubclass.name} in marshal stream @ root.object", e.message
+      end
+    end
+  end
+
+  class ArraySubclass < ::Array
+  end
+
+  def test_array_subclass
+    with_const(Gem::SafeMarshal, :PERMITTED_CLASSES, [ArraySubclass.name]) do
+      e = assert_raise(Gem::SafeMarshal::Visitors::ToRuby::UnsupportedError) do
+        Gem::SafeMarshal.safe_load(Marshal.dump(ArraySubclass.new << "abc"))
+      end
+      assert_equal "Unsupported user class #{ArraySubclass.name} in marshal stream @ root", e.message
+    end
   end
 
   def test_frozen_object