* ext/psych/lib/psych.rb: Adding Psych.safe_load for loading a user

  defined, restricted subset of Ruby object types.
* ext/psych/lib/psych/class_loader.rb: A class loader for
  encapsulating the logic for which objects are allowed to be
  deserialized.
* ext/psych/lib/psych/deprecated.rb: Changes to use the class loader
* ext/psych/lib/psych/exception.rb: ditto
* ext/psych/lib/psych/json/stream.rb: ditto
* ext/psych/lib/psych/nodes/node.rb: ditto
* ext/psych/lib/psych/scalar_scanner.rb: ditto
* ext/psych/lib/psych/stream.rb: ditto
* ext/psych/lib/psych/streaming.rb: ditto
* ext/psych/lib/psych/visitors/json_tree.rb: ditto
* ext/psych/lib/psych/visitors/to_ruby.rb: ditto
* ext/psych/lib/psych/visitors/yaml_tree.rb: ditto
* ext/psych/psych_to_ruby.c: ditto
* test/psych/helper.rb: ditto
* test/psych/test_safe_load.rb: tests for restricted subset.
* test/psych/test_scalar_scanner.rb: ditto
* test/psych/visitors/test_to_ruby.rb: ditto
* test/psych/visitors/test_yaml_tree.rb: ditto

12 files changed, 253 insertions, 53 deletions

diff --git a/lib/psych.rb b/lib/psych.rb
index 66a0641..711b3c1 100644
--- a/lib/psych.rb
+++ b/lib/psych.rb
@@ -246,6 +246,55 @@ module Psych
   end
 
   ###
+  # Safely load the yaml string in +yaml+.  By default, only the following
+  # classes are allowed to be deserialized:
+  #
+  # * TrueClass
+  # * FalseClass
+  # * NilClass
+  # * Numeric
+  # * String
+  # * Array
+  # * Hash
+  #
+  # Recursive data structures are not allowed by default.  Arbitrary classes
+  # can be allowed by adding those classes to the +whitelist+.  They are
+  # additive.  For example, to allow Date deserialization:
+  #
+  #   Psych.safe_load(yaml, [Date])
+  #
+  # Now the Date class can be loaded in addition to the classes listed above.
+  #
+  # Aliases can be explicitly allowed by changing the +aliases+ parameter.
+  # For example:
+  #
+  #   x = []
+  #   x << x
+  #   yaml = Psych.dump x
+  #   Psych.safe_load yaml               # => raises an exception
+  #   Psych.safe_load yaml, [], [], true # => loads the aliases
+  #
+  # A Psych::DisallowedClass exception will be raised if the yaml contains a
+  # class that isn't in the whitelist.
+  #
+  # A Psych::BadAlias exception will be raised if the yaml contains aliases
+  # but the +aliases+ parameter is set to false.
+  def self.safe_load yaml, whitelist_classes = [], whitelist_symbols = [], aliases = false, filename = nil
+    result = parse(yaml, filename)
+    return unless result
+
+    class_loader = ClassLoader::Restricted.new(whitelist_classes.map(&:to_s),
+                                               whitelist_symbols.map(&:to_s))
+    scanner      = ScalarScanner.new class_loader
+    if aliases
+      visitor = Visitors::ToRuby.new scanner, class_loader
+    else
+      visitor = Visitors::NoAliasRuby.new scanner, class_loader
+    end
+    visitor.accept result
+  end
+
+  ###
   # Parse a YAML string in +yaml+.  Returns the Psych::Nodes::Document.
   # +filename+ is used in the exception message if a Psych::SyntaxError is
   # raised.
@@ -355,7 +404,7 @@ module Psych
       io      = nil
     end
 
-    visitor = Psych::Visitors::YAMLTree.new options
+    visitor = Psych::Visitors::YAMLTree.create options
     visitor << o
     visitor.tree.yaml io, options
   end
@@ -367,7 +416,7 @@ module Psych
   #
   #   Psych.dump_stream("foo\n  ", {}) # => "--- ! \"foo\\n  \"\n--- {}\n"
   def self.dump_stream *objects
-    visitor = Psych::Visitors::YAMLTree.new {}
+    visitor = Psych::Visitors::YAMLTree.create({})
     objects.each do |o|
       visitor << o
     end
@@ -377,7 +426,7 @@ module Psych
   ###
   # Dump Ruby +object+ to a JSON string.
   def self.to_json object
-    visitor = Psych::Visitors::JSONTree.new
+    visitor = Psych::Visitors::JSONTree.create
     visitor << object
     visitor.tree.yaml
   end
@@ -435,7 +484,7 @@ module Psych
   @load_tags = {}
   @dump_tags = {}
   def self.add_tag tag, klass
-    @load_tags[tag] = klass
+    @load_tags[tag] = klass.name
     @dump_tags[klass] = tag
   end
 
diff --git a/lib/psych/class_loader.rb b/lib/psych/class_loader.rb
new file mode 100644
index 0000000..46c6b93
--- /dev/null
+++ b/lib/psych/class_loader.rb
@@ -0,0 +1,101 @@
+require 'psych/omap'
+require 'psych/set'
+
+module Psych
+  class ClassLoader # :nodoc:
+    BIG_DECIMAL = 'BigDecimal'
+    COMPLEX     = 'Complex'
+    DATE        = 'Date'
+    DATE_TIME   = 'DateTime'
+    EXCEPTION   = 'Exception'
+    OBJECT      = 'Object'
+    PSYCH_OMAP  = 'Psych::Omap'
+    PSYCH_SET   = 'Psych::Set'
+    RANGE       = 'Range'
+    RATIONAL    = 'Rational'
+    REGEXP      = 'Regexp'
+    STRUCT      = 'Struct'
+    SYMBOL      = 'Symbol'
+
+    def initialize
+      @cache = CACHE.dup
+    end
+
+    def load klassname
+      return nil if !klassname || klassname.empty?
+
+      find klassname
+    end
+
+    def symbolize sym
+      symbol
+      sym.to_sym
+    end
+
+    constants.each do |const|
+      konst = const_get const
+      define_method(const.to_s.downcase) do
+        load konst
+      end
+    end
+
+    private
+
+    def find klassname
+      @cache[klassname] ||= resolve(klassname)
+    end
+
+    def resolve klassname
+      name    = klassname
+      retried = false
+
+      begin
+        path2class(name)
+      rescue ArgumentError, NameError => ex
+        unless retried
+          name    = "Struct::#{name}"
+          retried = ex
+          retry
+        end
+        raise retried
+      end
+    end
+
+    CACHE = Hash[constants.map { |const|
+      val = const_get const
+      begin
+        [val, ::Object.const_get(val)]
+      rescue
+        nil
+      end
+    }.compact]
+
+    class Restricted < ClassLoader
+      def initialize classes, symbols
+        @classes = classes
+        @symbols = symbols
+        super()
+      end
+
+      def symbolize sym
+        return super if @symbols.empty?
+
+        if @symbols.include? sym
+          super
+        else
+          raise DisallowedClass, 'Symbol'
+        end
+      end
+
+      private
+
+      def find klassname
+        if @classes.include? klassname
+          super
+        else
+          raise DisallowedClass, klassname
+        end
+      end
+    end
+  end
+end
diff --git a/lib/psych/deprecated.rb b/lib/psych/deprecated.rb
index 1e42859..8c310b3 100644
--- a/lib/psych/deprecated.rb
+++ b/lib/psych/deprecated.rb
@@ -35,7 +35,8 @@ module Psych
     warn "#{caller[0]}: detect_implicit is deprecated" if $VERBOSE
     return '' unless String === thing
     return 'null' if '' == thing
-    ScalarScanner.new.tokenize(thing).class.name.downcase
+    ss = ScalarScanner.new(ClassLoader.new)
+    ss.tokenize(thing).class.name.downcase
   end
 
   def self.add_ruby_type type_tag, &block
diff --git a/lib/psych/exception.rb b/lib/psych/exception.rb
index d96c527..ce9d2ca 100644
--- a/lib/psych/exception.rb
+++ b/lib/psych/exception.rb
@@ -4,4 +4,10 @@ module Psych
 
   class BadAlias < Exception
   end
+
+  class DisallowedClass < Exception
+    def initialize klass_name
+      super "Tried to load unspecified class: #{klass_name}"
+    end
+  end
 end
diff --git a/lib/psych/json/stream.rb b/lib/psych/json/stream.rb
index be1a0a8..fe2a6e9 100644
--- a/lib/psych/json/stream.rb
+++ b/lib/psych/json/stream.rb
@@ -6,6 +6,7 @@ module Psych
     class Stream < Psych::Visitors::JSONTree
       include Psych::JSON::RubyEvents
       include Psych::Streaming
+      extend Psych::Streaming::ClassMethods
 
       class Emitter < Psych::Stream::Emitter # :nodoc:
         include Psych::JSON::YAMLEvents
diff --git a/lib/psych/nodes/node.rb b/lib/psych/nodes/node.rb
index 0cefe44..83233a6 100644
--- a/lib/psych/nodes/node.rb
+++ b/lib/psych/nodes/node.rb
@@ -1,4 +1,6 @@
 require 'stringio'
+require 'psych/class_loader'
+require 'psych/scalar_scanner'
 
 module Psych
   module Nodes
@@ -32,7 +34,7 @@ module Psych
       #
       # See also Psych::Visitors::ToRuby
       def to_ruby
-        Visitors::ToRuby.new.accept self
+        Visitors::ToRuby.create.accept(self)
       end
       alias :transform :to_ruby
 
diff --git a/lib/psych/scalar_scanner.rb b/lib/psych/scalar_scanner.rb
index 8aa594e..5935e26 100644
--- a/lib/psych/scalar_scanner.rb
+++ b/lib/psych/scalar_scanner.rb
@@ -19,10 +19,13 @@ module Psych
                   |[-+]?(?:0|[1-9][0-9_]*) (?# base 10)
                   |[-+]?0x[0-9a-fA-F_]+    (?# base 16))$/x
 
+    attr_reader :class_loader
+
     # Create a new scanner
-    def initialize
+    def initialize class_loader
       @string_cache = {}
       @symbol_cache = {}
+      @class_loader = class_loader
     end
 
     # Tokenize +string+ returning the ruby object
@@ -63,7 +66,7 @@ module Psych
       when /^\d{4}-(?:1[012]|0\d|\d)-(?:[12]\d|3[01]|0\d|\d)$/
         require 'date'
         begin
-          Date.strptime(string, '%Y-%m-%d')
+          class_loader.date.strptime(string, '%Y-%m-%d')
         rescue ArgumentError
           string
         end
@@ -75,9 +78,9 @@ module Psych
         Float::NAN
       when /^:./
         if string =~ /^:(["'])(.*)\1/
-          @symbol_cache[string] = $2.sub(/^:/, '').to_sym
+          @symbol_cache[string] = class_loader.symbolize($2.sub(/^:/, ''))
         else
-          @symbol_cache[string] = string.sub(/^:/, '').to_sym
+          @symbol_cache[string] = class_loader.symbolize(string.sub(/^:/, ''))
         end
       when /^[-+]?[0-9][0-9_]*(:[0-5]?[0-9])+$/
         i = 0
@@ -117,6 +120,8 @@ module Psych
     ###
     # Parse and return a Time from +string+
     def parse_time string
+      klass = class_loader.load 'Time'
+
       date, time = *(string.split(/[ tT]/, 2))
       (yy, m, dd) = date.split('-').map { |x| x.to_i }
       md = time.match(/(\d+:\d+:\d+)(?:\.(\d*))?\s*(Z|[-+]\d+(:\d\d)?)?/)
@@ -124,10 +129,10 @@ module Psych
       (hh, mm, ss) = md[1].split(':').map { |x| x.to_i }
       us = (md[2] ? Rational("0.#{md[2]}") : 0) * 1000000
 
-      time = Time.utc(yy, m, dd, hh, mm, ss, us)
+      time = klass.utc(yy, m, dd, hh, mm, ss, us)
 
       return time if 'Z' == md[3]
-      return Time.at(time.to_i, us) unless md[3]
+      return klass.at(time.to_i, us) unless md[3]
 
       tz = md[3].match(/^([+\-]?\d{1,2})\:?(\d{1,2})?$/)[1..-1].compact.map { |digit| Integer(digit, 10) }
       offset = tz.first * 3600
@@ -138,7 +143,7 @@ module Psych
         offset += ((tz[1] || 0) * 60)
       end
 
-      Time.at((time - offset).to_i, us)
+      klass.at((time - offset).to_i, us)
     end
   end
 end
diff --git a/lib/psych/stream.rb b/lib/psych/stream.rb
index 567c1bb..88c4c4c 100644
--- a/lib/psych/stream.rb
+++ b/lib/psych/stream.rb
@@ -32,5 +32,6 @@ module Psych
     end
 
     include Psych::Streaming
+    extend Psych::Streaming::ClassMethods
   end
 end
diff --git a/lib/psych/streaming.rb b/lib/psych/streaming.rb
index c6fa109..9d94eb5 100644
--- a/lib/psych/streaming.rb
+++ b/lib/psych/streaming.rb
@@ -1,10 +1,15 @@
 module Psych
   module Streaming
-    ###
-    # Create a new streaming emitter.  Emitter will print to +io+.  See
-    # Psych::Stream for an example.
-    def initialize io
-      super({}, self.class.const_get(:Emitter).new(io))
+    module ClassMethods
+      ###
+      # Create a new streaming emitter.  Emitter will print to +io+.  See
+      # Psych::Stream for an example.
+      def new io
+        emitter      = const_get(:Emitter).new(io)
+        class_loader = ClassLoader.new
+        ss           = ScalarScanner.new class_loader
+        super(emitter, ss, {})
+      end
     end
 
     ###
diff --git a/lib/psych/visitors/json_tree.rb b/lib/psych/visitors/json_tree.rb
index 0350dd1..0127ac8 100644
--- a/lib/psych/visitors/json_tree.rb
+++ b/lib/psych/visitors/json_tree.rb
@@ -5,8 +5,11 @@ module Psych
     class JSONTree < YAMLTree
       include Psych::JSON::RubyEvents
 
-      def initialize options = {}, emitter = Psych::JSON::TreeBuilder.new
-        super
+      def self.create options = {}
+        emitter = Psych::JSON::TreeBuilder.new
+        class_loader = ClassLoader.new
+        ss           = ScalarScanner.new class_loader
+        new(emitter, ss, options)
       end
 
       def accept target
diff --git a/lib/psych/visitors/to_ruby.rb b/lib/psych/visitors/to_ruby.rb
index 75c7bc0..f770bb8 100644
--- a/lib/psych/visitors/to_ruby.rb
+++ b/lib/psych/visitors/to_ruby.rb
@@ -1,4 +1,5 @@
 require 'psych/scalar_scanner'
+require 'psych/class_loader'
 require 'psych/exception'
 
 unless defined?(Regexp::NOENCODING)
@@ -10,11 +11,20 @@ module Psych
     ###
     # This class walks a YAML AST, converting each node to ruby
     class ToRuby < Psych::Visitors::Visitor
-      def initialize ss = ScalarScanner.new
+      def self.create
+        class_loader = ClassLoader.new
+        scanner      = ScalarScanner.new class_loader
+        new(scanner, class_loader)
+      end
+
+      attr_reader :class_loader
+
+      def initialize ss, class_loader
         super()
         @st = {}
         @ss = ss
         @domain_types = Psych.domain_types
+        @class_loader = class_loader
       end
 
       def accept target
@@ -33,7 +43,7 @@ module Psych
       end
 
       def deserialize o
-        if klass = Psych.load_tags[o.tag]
+        if klass = resolve_class(Psych.load_tags[o.tag])
           instance = klass.allocate
 
           if instance.respond_to?(:init_with)
@@ -60,19 +70,23 @@ module Psych
           end
         when '!ruby/object:BigDecimal'
           require 'bigdecimal'
-          BigDecimal._load o.value
+          class_loader.big_decimal._load o.value
         when "!ruby/object:DateTime"
+          class_loader.date_time
           require 'date'
           @ss.parse_time(o.value).to_datetime
         when "!ruby/object:Complex"
+          class_loader.complex
           Complex(o.value)
         when "!ruby/object:Rational"
+          class_loader.rational
           Rational(o.value)
         when "!ruby/class", "!ruby/module"
           resolve_class o.value
         when "tag:yaml.org,2002:float", "!float"
           Float(@ss.tokenize(o.value))
         when "!ruby/regexp"
+          klass = class_loader.regexp
           o.value =~ /^\/(.*)\/([mixn]*)$/
           source  = $1
           options = 0
@@ -86,15 +100,16 @@ module Psych
             else lang = option
             end
           end
-          Regexp.new(*[source, options, lang].compact)
+          klass.new(*[source, options, lang].compact)
         when "!ruby/range"
+          klass = class_loader.range
           args = o.value.split(/([.]{2,3})/, 2).map { |s|
             accept Nodes::Scalar.new(s)
           }
           args.push(args.delete_at(1) == '...')
-          Range.new(*args)
+          klass.new(*args)
         when /^!ruby\/sym(bol)?:?(.*)?$/
-          o.value.to_sym
+          class_loader.symbolize o.value
         else
           @ss.tokenize o.value
         end
@@ -106,7 +121,7 @@ module Psych
       end
 
       def visit_Psych_Nodes_Sequence o
-        if klass = Psych.load_tags[o.tag]
+        if klass = resolve_class(Psych.load_tags[o.tag])
           instance = klass.allocate
 
           if instance.respond_to?(:init_with)
@@ -138,22 +153,24 @@ module Psych
       end
 
       def visit_Psych_Nodes_Mapping o
-        return revive(Psych.load_tags[o.tag], o) if Psych.load_tags[o.tag]
+        if Psych.load_tags[o.tag]
+          return revive(resolve_class(Psych.load_tags[o.tag]), o)
+        end
         return revive_hash({}, o) unless o.tag
 
         case o.tag
         when /^!ruby\/struct:?(.*)?$/
-          klass = resolve_class($1)
+          klass = resolve_class($1) if $1
 
           if klass
             s = register(o, klass.allocate)
 
             members = {}
-            struct_members = s.members.map { |x| x.to_sym }
+            struct_members = s.members.map { |x| class_loader.symbolize x }
             o.children.each_slice(2) do |k,v|
               member = accept(k)
               value  = accept(v)
-              if struct_members.include?(member.to_sym)
+              if struct_members.include?(class_loader.symbolize(member))
                 s.send("#{member}=", value)
               else
                 members[member.to_s.sub(/^@/, '')] = value
@@ -161,22 +178,27 @@ module Psych
             end
             init_with(s, members, o)
           else
+            klass = class_loader.struct
             members = o.children.map { |c| accept c }
             h = Hash[*members]
-            Struct.new(*h.map { |k,v| k.to_sym }).new(*h.map { |k,v| v })
+            klass.new(*h.map { |k,v|
+              class_loader.symbolize k
+            }).new(*h.map { |k,v| v })
           end
 
         when /^!ruby\/object:?(.*)?$/
           name = $1 || 'Object'
 
           if name == 'Complex'
+            class_loader.complex
             h = Hash[*o.children.map { |c| accept c }]
             register o, Complex(h['real'], h['image'])
           elsif name == 'Rational'
+            class_loader.rational
             h = Hash[*o.children.map { |c| accept c }]
             register o, Rational(h['numerator'], h['denominator'])
           else
-            obj = revive((resolve_class(name) || Object), o)
+            obj = revive((resolve_class(name) || class_loader.object), o)
             obj
           end
 
@@ -214,18 +236,19 @@ module Psych
           list
 
         when '!ruby/range'
+          klass = class_loader.range
           h = Hash[*o.children.map { |c| accept c }]
-          register o, Range.new(h['begin'], h['end'], h['excl'])
+          register o, klass.new(h['begin'], h['end'], h['excl'])
 
         when /^!ruby\/exception:?(.*)?$/
           h = Hash[*o.children.map { |c| accept c }]
 
-          e = build_exception((resolve_class($1) || Exception),
+          e = build_exception((resolve_class($1) || class_loader.exception),
                               h.delete('message'))
           init_with(e, h, o)
 
         when '!set', 'tag:yaml.org,2002:set'
-          set = Psych::Set.new
+          set = class_loader.psych_set.new
           @st[o.anchor] = set if o.anchor
           o.children.each_slice(2) do |k,v|
             set[accept(k)] = accept(v)
@@ -236,7 +259,7 @@ module Psych
           revive_hash resolve_class($1).new, o
 
         when '!omap', 'tag:yaml.org,2002:omap'
-          map = register(o, Psych::Omap.new)
+          map = register(o, class_loader.psych_omap.new)
           o.children.each_slice(2) do |l,r|
             map[accept(l)] = accept r
           end
@@ -336,21 +359,13 @@ module Psych
 
       # Convert +klassname+ to a Class
       def resolve_class klassname
-        return nil unless klassname and not klassname.empty?
-
-        name    = klassname
-        retried = false
-
-        begin
-          path2class(name)
-        rescue ArgumentError, NameError => ex
-          unless retried
-            name    = "Struct::#{name}"
-            retried = ex
-            retry
-          end
-          raise retried
-        end
+        class_loader.load klassname
+      end
+    end
+
+    class NoAliasRuby < ToRuby
+      def visit_Psych_Nodes_Alias o
+        raise BadAlias, "Unknown alias: #{o.anchor}"
       end
     end
   end
diff --git a/lib/psych/visitors/yaml_tree.rb b/lib/psych/visitors/yaml_tree.rb
index 96640e0..ddd745b 100644
--- a/lib/psych/visitors/yaml_tree.rb
+++ b/lib/psych/visitors/yaml_tree.rb
@@ -1,3 +1,7 @@
+require 'psych/tree_builder'
+require 'psych/scalar_scanner'
+require 'psych/class_loader'
+
 module Psych
   module Visitors
     ###
@@ -36,7 +40,14 @@ module Psych
       alias :finished? :finished
       alias :started? :started
 
-      def initialize options = {}, emitter = TreeBuilder.new, ss = ScalarScanner.new
+      def self.create options = {}, emitter = nil
+        emitter      ||= TreeBuilder.new
+        class_loader = ClassLoader.new
+        ss           = ScalarScanner.new class_loader
+        new(emitter, ss, options)
+      end
+
+      def initialize emitter, ss, options
         super()
         @started  = false
         @finished = false