Merge pull request #647 from headius/jruby_default_limits

Allow setting defaults for SnakeYAML limits
ruby · Sep 18, 2023 · 48cdab3 · 48cdab3
2 parents 587c50f + 0c1754e
commit 48cdab3
Show file tree

Hide file tree

Showing 2 changed files with 98 additions and 1 deletion.
diff --git a/ext/java/org/jruby/ext/psych/PsychParser.java b/ext/java/org/jruby/ext/psych/PsychParser.java
@@ -41,6 +41,7 @@
 import org.jruby.RubyIO;
 import org.jruby.RubyKernel;
 import org.jruby.RubyModule;
+import org.jruby.RubyNumeric;
 import org.jruby.RubyObject;
 import org.jruby.RubyString;
 import org.jruby.anno.JRubyMethod;
@@ -113,6 +114,13 @@ public static void initPsychParser(Ruby runtime, RubyModule psych) {
         psychParser.defineConstant("UTF16BE", runtime.newFixnum(YAML_UTF16BE_ENCODING.ordinal()));
 
         psychParser.defineAnnotatedMethods(PsychParser.class);
+
+        // defaults for SnakeYAML load settings
+        LoadSettings defaults = LoadSettings.builder().build();
+        psychParser.setInternalVariable("max_aliases_for_collections", runtime.newFixnum(defaults.getMaxAliasesForCollections()));
+        psychParser.setInternalVariable("allow_duplicate_keys", runtime.newBoolean(defaults.getAllowDuplicateKeys()));
+        psychParser.setInternalVariable("allow_recursive_keys", runtime.newBoolean(defaults.getAllowRecursiveKeys()));
+        psychParser.setInternalVariable("code_point_limit", runtime.newFixnum(defaults.getCodePointLimit()));
     }
 
     public PsychParser(Ruby runtime, RubyClass klass) {
@@ -131,7 +139,15 @@ public PsychParser(Ruby runtime, RubyClass klass) {
         this.start_mapping = sites[Call.start_mapping.ordinal()];
         this.end_mapping = sites[Call.end_mapping.ordinal()];
         this.end_stream = sites[Call.end_stream.ordinal()];
-        this.loadSettingsBuilder = LoadSettings.builder().setSchema(new CoreSchema());
+
+        // prepare settings builder and apply global defaults
+        LoadSettingsBuilder lsb = LoadSettings.builder();
+        lsb.setSchema(new CoreSchema());
+        lsb.setMaxAliasesForCollections(((IRubyObject) klass.getInternalVariable("max_aliases_for_collections")).convertToInteger().getIntValue());
+        lsb.setAllowDuplicateKeys(((IRubyObject) klass.getInternalVariable("allow_duplicate_keys")).isTrue());
+        lsb.setAllowRecursiveKeys(((IRubyObject) klass.getInternalVariable("allow_recursive_keys")).isTrue());
+        lsb.setCodePointLimit(((IRubyObject) klass.getInternalVariable("code_point_limit")).convertToInteger().getIntValue());
+        this.loadSettingsBuilder = lsb;
     }
 
     private IRubyObject stringOrNilForAnchor(ThreadContext context, Optional<Anchor> value) {
@@ -563,6 +579,68 @@ public IRubyObject code_point_limit(ThreadContext context) {
         return context.runtime.newFixnum(buildSettings().getCodePointLimit());
     }
 
+    // class-level accessors for default values
+
+    @JRubyMethod(name = "max_aliases_for_collections=", meta = true)
+    public static IRubyObject max_aliases_for_collections_set(ThreadContext context, IRubyObject self, IRubyObject max) {
+        int maxAliasesForCollections = RubyNumeric.num2int(max);
+
+        if (maxAliasesForCollections <= 0) {
+            throw context.runtime.newRangeError("max_aliases_for_collections must be positive");
+        }
+
+        self.getInternalVariables().setInternalVariable("max_aliases_for_collections", max);
+
+        return max;
+    }
+
+    @JRubyMethod(name = "max_aliases_for_collections")
+    public static IRubyObject max_aliases_for_collections(ThreadContext context, IRubyObject self) {
+        return (IRubyObject) self.getInternalVariables().getInternalVariable("max_aliases_for_collections");
+    }
+
+    @JRubyMethod(name = "allow_duplicate_keys=", meta = true)
+    public static IRubyObject allow_duplicate_keys_set(IRubyObject self, IRubyObject allow) {
+        self.getInternalVariables().setInternalVariable("allow_duplicate_keys", allow);
+
+        return allow;
+    }
+
+    @JRubyMethod(name = "allow_duplicate_keys", meta = true)
+    public static IRubyObject allow_duplicate_keys(ThreadContext context, IRubyObject self) {
+        return (IRubyObject) self.getInternalVariables().getInternalVariable("allow_duplicate_keys");
+    }
+
+    @JRubyMethod(name = "allow_recursive_keys=", meta = true)
+    public static IRubyObject allow_recursive_keys_set(IRubyObject self, IRubyObject allow) {
+        self.getInternalVariables().setInternalVariable("allow_recursive_keys", allow);
+
+        return allow;
+    }
+
+    @JRubyMethod(name = "allow_recursive_keys", meta = true)
+    public static IRubyObject allow_recursive_keys(ThreadContext context, IRubyObject self) {
+        return (IRubyObject) self.getInternalVariables().getInternalVariable("allow_recursive_keys");
+    }
+
+    @JRubyMethod(name = "code_point_limit=", meta = true)
+    public static IRubyObject code_point_limit_set(ThreadContext context, IRubyObject self, IRubyObject limit) {
+        int codePointLimit = RubyNumeric.num2int(limit);
+
+        if (codePointLimit <= 0) {
+            throw context.runtime.newRangeError("code_point_limit must be positive");
+        }
+
+        self.getInternalVariables().setInternalVariable("code_point_limit", limit);
+
+        return limit;
+    }
+
+    @JRubyMethod(name = "code_point_limit", meta = true)
+    public static IRubyObject code_point_limit(ThreadContext context, IRubyObject self) {
+        return (IRubyObject) self.getInternalVariables().getInternalVariable("code_point_limit");
+    }
+
     private LoadSettings buildSettings() {
         return loadSettingsBuilder.build();
     }

diff --git a/test/psych/test_parser.rb b/test/psych/test_parser.rb
@@ -384,6 +384,25 @@ def test_event_location
                      [:end_stream, [2, 0, 2, 0]]], events
     end
 
+    if Psych::Parser.method_defined?(:code_point_limit)
+      def test_code_point_limit
+        yaml = "foo: bar\n" * 500_000
+        assert_raise(org.snakeyaml.engine.v2.exceptions.YamlEngineException) do
+          Psych.load(yaml)
+        end
+
+        assert_nothing_raised do
+          begin
+            old_code_point_limit, Psych::Parser.code_point_limit = Psych::Parser::code_point_limit, 5_000_000
+
+            Psych.load(yaml)
+          ensure
+            Psych::Parser.code_point_limit = old_code_point_limit
+          end
+        end
+      end
+    end
+
     def assert_called call, with = nil, parser = @parser
       if with
         call = parser.handler.calls.find { |x|