Merge pull request #617 from Shopify/fix-safe-dump-symbol

Fix RestrictedYAMLTree allowing the Symbol class should allow all symbols

2 files changed, 7 insertions, 4 deletions

diff --git a/lib/psych/visitors/yaml_tree.rb b/lib/psych/visitors/yaml_tree.rb
index 3185879..5149178 100644
--- a/lib/psych/visitors/yaml_tree.rb
+++ b/lib/psych/visitors/yaml_tree.rb
@@ -568,7 +568,7 @@ module Psych
           raise BadAlias, "Tried to dump an aliased object"
         end
 
-        unless @permitted_classes[target.class]
+        unless Symbol === target || @permitted_classes[target.class]
           raise DisallowedClass.new('dump', target.class.name || target.class.inspect)
         end
 
@@ -576,7 +576,7 @@ module Psych
       end
 
       def visit_Symbol sym
-        unless @permitted_symbols[sym]
+        unless @permitted_classes[Symbol] || @permitted_symbols[sym]
           raise DisallowedClass.new('dump', "Symbol(#{sym.inspect})")
         end
 
diff --git a/test/psych/test_psych.rb b/test/psych/test_psych.rb
index 1abd69c..c977e79 100644
--- a/test/psych/test_psych.rb
+++ b/test/psych/test_psych.rb
@@ -419,12 +419,15 @@ eoyml
   end
 
   def test_safe_dump_symbols
+    assert_equal Psych.dump(:foo), Psych.safe_dump(:foo, permitted_classes: [Symbol])
+    assert_equal Psych.dump(:foo), Psych.safe_dump(:foo, permitted_symbols: [:foo])
+
     error = assert_raise Psych::DisallowedClass do
-      Psych.safe_dump(:foo, permitted_classes: [Symbol])
+      Psych.safe_dump(:foo)
     end
     assert_equal "Tried to dump unspecified class: Symbol(:foo)", error.message
 
-    assert_match(/\A--- :foo\n(?:\.\.\.\n)?\z/, Psych.safe_dump(:foo, permitted_classes: [Symbol], permitted_symbols: [:foo]))
+    assert_match(/\A--- :foo\n(?:\.\.\.\n)?\z/, Psych.safe_dump(:foo, permitted_symbols: [:foo]))
   end
 
   def test_safe_dump_aliases