[ruby/psych] Implement YAML.safe_dump to make safe_load more usable.

In case where Psych is used as a two way serializers,
e.g. to serialize some cache or config, it is preferable
to have the same restrictions on both load and dump.

Otherwise you might dump and persist some objects payloads
that you later won't be able to read.

https://github.com/ruby/psych/commit/441958396f

3 files changed, 50 insertions, 4 deletions

diff --git a/ext/psych/lib/psych/class_loader.rb b/ext/psych/lib/psych/class_loader.rb
index a5d1a7a4a92..088373cd669 100644
--- a/ext/psych/lib/psych/class_loader.rb
+++ b/ext/psych/lib/psych/class_loader.rb
@@ -86,7 +86,7 @@ module Psych
         if @symbols.include? sym
           super
         else
-          raise DisallowedClass, 'Symbol'
+          raise DisallowedClass.new('load', 'Symbol')
         end
       end
 
@@ -96,7 +96,7 @@ module Psych
         if @classes.include? klassname
           super
         else
-          raise DisallowedClass, klassname
+          raise DisallowedClass.new('load', klassname)
         end
       end
     end
diff --git a/ext/psych/lib/psych/exception.rb b/ext/psych/lib/psych/exception.rb
index fac0c42b9fb..f473b95a3bb 100644
--- a/ext/psych/lib/psych/exception.rb
+++ b/ext/psych/lib/psych/exception.rb
@@ -7,8 +7,8 @@ module Psych
   end
 
   class DisallowedClass < Exception
-    def initialize klass_name
-      super "Tried to load unspecified class: #{klass_name}"
+    def initialize action, klass_name
+      super "Tried to #{action} unspecified class: #{klass_name}"
     end
   end
 end
diff --git a/ext/psych/lib/psych/visitors/yaml_tree.rb b/ext/psych/lib/psych/visitors/yaml_tree.rb
index bf7c0bb8ca9..05748dd81cb 100644
--- a/ext/psych/lib/psych/visitors/yaml_tree.rb
+++ b/ext/psych/lib/psych/visitors/yaml_tree.rb
@@ -535,5 +535,51 @@ module Psych
         end
       end
     end
+
+    class RestrictedYAMLTree < YAMLTree
+      DEFAULT_PERMITTED_CLASSES = {
+        TrueClass => true,
+        FalseClass => true,
+        NilClass => true,
+        Integer => true,
+        Float => true,
+        String => true,
+        Array => true,
+        Hash => true,
+      }.compare_by_identity.freeze
+
+      def initialize emitter, ss, options
+        super
+        @permitted_classes = DEFAULT_PERMITTED_CLASSES.dup
+        Array(options[:permitted_classes]).each do |klass|
+          @permitted_classes[klass] = true
+        end
+        @permitted_symbols = {}.compare_by_identity
+        Array(options[:permitted_symbols]).each do |symbol|
+          @permitted_symbols[symbol] = true
+        end
+        @aliases = options.fetch(:aliases, false)
+      end
+
+      def accept target
+        if !@aliases && @st.key?(target)
+          raise BadAlias, "Tried to dump an aliased object"
+        end
+
+        unless @permitted_classes[target.class]
+          raise DisallowedClass.new('dump', target.class.name || target.class.inspect)
+        end
+
+        super
+      end
+
+      def visit_Symbol sym
+        unless @permitted_symbols[sym]
+          raise DisallowedClass.new('dump', "Symbol(#{sym.inspect})")
+        end
+
+        super
+      end
+    end
   end
 end