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.

5 files changed, 182 insertions, 5 deletions

diff --git a/lib/psych.rb b/lib/psych.rb
index 7c6b8d0..72276d9 100644
--- a/lib/psych.rb
+++ b/lib/psych.rb
@@ -282,7 +282,8 @@ module Psych
   # * TrueClass
   # * FalseClass
   # * NilClass
-  # * Numeric
+  # * Integer
+  # * Float
   # * String
   # * Array
   # * Hash
@@ -513,6 +514,79 @@ module Psych
   end
 
   ###
+  # call-seq:
+  #   Psych.safe_dump(o)               -> string of yaml
+  #   Psych.safe_dump(o, options)      -> string of yaml
+  #   Psych.safe_dump(o, io)           -> io object passed in
+  #   Psych.safe_dump(o, io, options)  -> io object passed in
+  #
+  # Safely dump Ruby object +o+ to a YAML string. Optional +options+ may be passed in
+  # to control the output format.  If an IO object is passed in, the YAML will
+  # be dumped to that IO object. By default, only the following
+  # classes are allowed to be serialized:
+  #
+  # * TrueClass
+  # * FalseClass
+  # * NilClass
+  # * Integer
+  # * Float
+  # * String
+  # * Array
+  # * Hash
+  #
+  # Arbitrary classes can be allowed by adding those classes to the +permitted_classes+
+  # keyword argument.  They are additive.  For example, to allow Date serialization:
+  #
+  #   Psych.safe_dump(yaml, permitted_classes: [Date])
+  #
+  # Now the Date class can be dumped in addition to the classes listed above.
+  #
+  # A Psych::DisallowedClass exception will be raised if the object contains a
+  # class that isn't in the +permitted_classes+ list.
+  #
+  # Currently supported options are:
+  #
+  # [<tt>:indentation</tt>]   Number of space characters used to indent.
+  #                           Acceptable value should be in <tt>0..9</tt> range,
+  #                           otherwise option is ignored.
+  #
+  #                           Default: <tt>2</tt>.
+  # [<tt>:line_width</tt>]    Max character to wrap line at.
+  #
+  #                           Default: <tt>0</tt> (meaning "wrap at 81").
+  # [<tt>:canonical</tt>]     Write "canonical" YAML form (very verbose, yet
+  #                           strictly formal).
+  #
+  #                           Default: <tt>false</tt>.
+  # [<tt>:header</tt>]        Write <tt>%YAML [version]</tt> at the beginning of document.
+  #
+  #                           Default: <tt>false</tt>.
+  #
+  # Example:
+  #
+  #   # Dump an array, get back a YAML string
+  #   Psych.safe_dump(['a', 'b'])  # => "---\n- a\n- b\n"
+  #
+  #   # Dump an array to an IO object
+  #   Psych.safe_dump(['a', 'b'], StringIO.new)  # => #<StringIO:0x000001009d0890>
+  #
+  #   # Dump an array with indentation set
+  #   Psych.safe_dump(['a', ['b']], indentation: 3) # => "---\n- a\n-  - b\n"
+  #
+  #   # Dump an array to an IO with indentation set
+  #   Psych.safe_dump(['a', ['b']], StringIO.new, indentation: 3)
+  def self.safe_dump o, io = nil, options = {}
+    if Hash === io
+      options = io
+      io      = nil
+    end
+
+    visitor = Psych::Visitors::RestrictedYAMLTree.create options
+    visitor << o
+    visitor.tree.yaml io, options
+  end
+
+  ###
   # Dump a list of objects as separate documents to a document stream.
   #
   # Example:
diff --git a/lib/psych/class_loader.rb b/lib/psych/class_loader.rb
index a5d1a7a..088373c 100644
--- a/lib/psych/class_loader.rb
+++ b/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/lib/psych/exception.rb b/lib/psych/exception.rb
index fac0c42..f473b95 100644
--- a/lib/psych/exception.rb
+++ b/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/lib/psych/visitors/yaml_tree.rb b/lib/psych/visitors/yaml_tree.rb
index bf7c0bb..05748dd 100644
--- a/lib/psych/visitors/yaml_tree.rb
+++ b/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
diff --git a/test/psych/test_psych.rb b/test/psych/test_psych.rb
index 256ed91..c9d39c5 100644
--- a/test/psych/test_psych.rb
+++ b/test/psych/test_psych.rb
@@ -381,4 +381,61 @@ hoge:
     result = Psych.safe_load(yaml, symbolize_names: true)
     assert_equal result, { foo: { bar: "baz", 1 => 2 }, hoge: [{ fuga: "piyo" }] }
   end
+
+  def test_safe_dump_defaults
+    yaml = <<-eoyml
+---
+array:
+- 1
+float: 13.12
+booleans:
+- true
+- false
+eoyml
+
+    payload = YAML.safe_dump({
+      "array" => [1],
+      "float" => 13.12,
+      "booleans" => [true, false],
+    })
+    assert_equal yaml, payload
+  end
+
+  def test_safe_dump_unpermitted_class
+    error = assert_raises Psych::DisallowedClass do
+      YAML.safe_dump(Object.new)
+    end
+    assert_equal "Tried to dump unspecified class: Object", error.message
+
+    hash_subclass = Class.new(Hash)
+    error = assert_raises Psych::DisallowedClass do
+      YAML.safe_dump(hash_subclass.new)
+    end
+    assert_equal "Tried to dump unspecified class: #{hash_subclass.inspect}", error.message
+  end
+
+  def test_safe_dump_extra_permitted_classes
+    assert_equal "--- !ruby/object {}\n", YAML.safe_dump(Object.new, permitted_classes: [Object])
+  end
+
+  def test_safe_dump_symbols
+    error = assert_raises Psych::DisallowedClass do
+      YAML.safe_dump(:foo, permitted_classes: [Symbol])
+    end
+    assert_equal "Tried to dump unspecified class: Symbol(:foo)", error.message
+
+    assert_equal "--- :foo\n", YAML.safe_dump(:foo, permitted_classes: [Symbol], permitted_symbols: [:foo])
+  end
+
+  def test_safe_dump_aliases
+    x = []
+    x << x
+    error = assert_raises Psych::BadAlias do
+      YAML.safe_dump(x)
+    end
+    assert_equal "Tried to dump an aliased object", error.message
+
+    assert_equal "--- &1\n" + "- *1\n", YAML.safe_dump(x, aliases: true)
+  end
+
 end