4 files changed, 125 insertions, 5 deletions

diff --git a/ext/psych/lib/psych.rb b/ext/psych/lib/psych.rb
index 6205061d0b..1efe8d5f62 100644
--- a/ext/psych/lib/psych.rb
+++ b/ext/psych/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/ext/psych/lib/psych/class_loader.rb b/ext/psych/lib/psych/class_loader.rb
index a5d1a7a4a9..088373cd66 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 fac0c42b9f..f473b95a3b 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 bf7c0bb8ca..05748dd81c 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