Merge pull request #495 from Shopify/safe-dumpHEADmaster

Implement YAML.safe_dump to make safe_load more usable.

5 files changed, 182 insertions, 5 deletions

diff --git a/lib/psych.rb b/lib/psych.rb
index 6205061..1efe8d5 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