Allow alias expansion to be limited via configuration (#50)

* Allow alias expansion to be limited via configuration - make internal configuration copies - move location to separate class - enable disable of marks for snakeyaml-engine Signed-off-by: Michael Edgar <michael@xlate.io> * Only count alias expansion when configured maximum is less than max long Signed-off-by: Michael Edgar <michael@xlate.io> Signed-off-by: Michael Edgar <michael@xlate.io>
xlate · Sep 5, 2022 · 7b690a2 · 7b690a2
1 parent 6788d2e
commit 7b690a2
Show file tree

Hide file tree

Showing 14 changed files with 192 additions and 74 deletions.
diff --git a/src/main/java/io/xlate/yamljson/SettingsBuilder.java b/src/main/java/io/xlate/yamljson/SettingsBuilder.java
@@ -1,12 +1,14 @@
 package io.xlate.yamljson;
 
 import java.util.Map;
-import java.util.function.Consumer;
+import java.util.Objects;
+import java.util.function.Function;
 import java.util.function.Supplier;
 
 import org.snakeyaml.engine.v2.api.DumpSettings;
 import org.snakeyaml.engine.v2.api.DumpSettingsBuilder;
 import org.snakeyaml.engine.v2.api.LoadSettings;
+import org.snakeyaml.engine.v2.api.LoadSettingsBuilder;
 import org.yaml.snakeyaml.DumperOptions;
 import org.yaml.snakeyaml.LoaderOptions;
 
@@ -30,35 +32,39 @@ default <T> T loadProvider(Supplier<T> providerSupplier, String providerModule)
         }
     }
 
-    default LoaderOptions buildLoaderOptions(Map<String, ?> properties) {
+    default LoaderOptions buildLoaderOptions(Map<String, Object> properties) {
+        LoaderOptions options = new LoaderOptions();
+        replace(properties, Yaml.Settings.LOAD_MAX_ALIAS_EXPANSION_SIZE, Long::valueOf, Long.MAX_VALUE);
         // No load properties supported currently
-        return new LoaderOptions();
+        return options;
     }
 
-    default DumperOptions buildDumperOptions(Map<String, ?> properties) {
+    default DumperOptions buildDumperOptions(Map<String, Object> properties) {
         DumperOptions settings = new DumperOptions();
-        setBoolean(properties, Yaml.Settings.DUMP_EXPLICIT_START, settings::setExplicitStart);
-        setBoolean(properties, Yaml.Settings.DUMP_EXPLICIT_END, settings::setExplicitEnd);
+        settings.setExplicitStart(getProperty(properties, Yaml.Settings.DUMP_EXPLICIT_START, Boolean::valueOf, false));
+        settings.setExplicitEnd(getProperty(properties, Yaml.Settings.DUMP_EXPLICIT_END, Boolean::valueOf, false));
         return settings;
     }
 
-    default LoadSettings buildLoadSettings(Map<String, ?> properties) {
-        // No load properties supported currently
-        return LoadSettings.builder().build();
+    default LoadSettings buildLoadSettings(Map<String, Object> properties) {
+        LoadSettingsBuilder settings = LoadSettings.builder();
+        settings.setUseMarks(getProperty(properties, Yaml.Settings.LOAD_USE_MARKS, Boolean::valueOf, true));
+        replace(properties, Yaml.Settings.LOAD_MAX_ALIAS_EXPANSION_SIZE, Long::valueOf, Long.MAX_VALUE);
+        return settings.build();
     }
 
-    default DumpSettings buildDumpSettings(Map<String, ?> properties) {
+    default DumpSettings buildDumpSettings(Map<String, Object> properties) {
         DumpSettingsBuilder settings = DumpSettings.builder();
-        setBoolean(properties, Yaml.Settings.DUMP_EXPLICIT_START, settings::setExplicitStart);
-        setBoolean(properties, Yaml.Settings.DUMP_EXPLICIT_END, settings::setExplicitEnd);
+        settings.setExplicitStart(getProperty(properties, Yaml.Settings.DUMP_EXPLICIT_START, Boolean::valueOf, false));
+        settings.setExplicitEnd(getProperty(properties, Yaml.Settings.DUMP_EXPLICIT_END, Boolean::valueOf, false));
         return settings.build();
     }
 
-    default void setBoolean(Map<String, ?> properties, String key, Consumer<Boolean> setter) {
-        setter.accept(getBoolean(properties, key));
+    default <T> T getProperty(Map<String, ?> properties, String key, Function<String, T> parser, T defaultValue) {
+        return parser.apply(String.valueOf(Objects.requireNonNullElse(properties.get(key), defaultValue)));
     }
 
-    default Boolean getBoolean(Map<String, ?> properties, String key) {
-        return Boolean.valueOf(String.valueOf(properties.get(key)));
+    default <T> void replace(Map<String, Object> properties, String key, Function<String, T> parser, T defaultValue) {
+        properties.put(key, getProperty(properties, key, parser, defaultValue));
     }
 }
diff --git a/src/main/java/io/xlate/yamljson/SnakeYamlEngineParser.java b/src/main/java/io/xlate/yamljson/SnakeYamlEngineParser.java
@@ -17,17 +17,18 @@
 
 import java.io.Reader;
 import java.util.Iterator;
+import java.util.Map;
 
 final class SnakeYamlEngineParser extends YamlParser<org.snakeyaml.engine.v2.events.Event, org.snakeyaml.engine.v2.exceptions.Mark> {
 
-    SnakeYamlEngineParser(Iterator<org.snakeyaml.engine.v2.events.Event> yamlEvents, Reader yamlReader) {
-        super(yamlEvents, yamlReader);
+    SnakeYamlEngineParser(Iterator<org.snakeyaml.engine.v2.events.Event> yamlEvents, Reader yamlReader, Map<String, ?> properties) {
+        super(yamlEvents, yamlReader, properties);
     }
 
     @Override
-    protected org.snakeyaml.engine.v2.exceptions.Mark getMark() {
-        if (currentYamlEvent != null) {
-            return currentYamlEvent.getStartMark().orElse(null);
+    protected org.snakeyaml.engine.v2.exceptions.Mark getMark(org.snakeyaml.engine.v2.events.Event event) {
+        if (event != null) {
+            return event.getStartMark().orElse(null);
         }
 
         return null;

diff --git a/src/main/java/io/xlate/yamljson/SnakeYamlParser.java b/src/main/java/io/xlate/yamljson/SnakeYamlParser.java
@@ -17,17 +17,18 @@
 
 import java.io.Reader;
 import java.util.Iterator;
+import java.util.Map;
 
 final class SnakeYamlParser extends YamlParser<org.yaml.snakeyaml.events.Event, org.yaml.snakeyaml.error.Mark> {
 
-    SnakeYamlParser(Iterator<org.yaml.snakeyaml.events.Event> yamlEvents, Reader yamlReader) {
-        super(yamlEvents, yamlReader);
+    SnakeYamlParser(Iterator<org.yaml.snakeyaml.events.Event> yamlEvents, Reader yamlReader, Map<String, ?> properties) {
+        super(yamlEvents, yamlReader, properties);
     }
 
     @Override
-    protected org.yaml.snakeyaml.error.Mark getMark() {
-        if (currentYamlEvent != null) {
-            return currentYamlEvent.getStartMark();
+    protected org.yaml.snakeyaml.error.Mark getMark(org.yaml.snakeyaml.events.Event event) {
+        if (event != null) {
+            return event.getStartMark();
         }
 
         return null;

diff --git a/src/main/java/io/xlate/yamljson/Yaml.java b/src/main/java/io/xlate/yamljson/Yaml.java
@@ -119,11 +119,30 @@ private Settings() {
         public static final String YAML_VERSION = "io.xlate.yamljson.YAML_VERSION";
 
         /**
+         * The maximum number of scalars to which an alias (or chain of aliases via arrays/objects)
+         * may expand.
+         *
+         * @since 0.1.0
+         */
+        public static final String LOAD_MAX_ALIAS_EXPANSION_SIZE = "LOAD_MAX_ALIAS_EXPANSION_SIZE";
+
+        /**
+         * Requires snakeyaml-engine, not supported with snakeyaml.
+         *
+         * @see org.snakeyaml.engine.v2.api.LoadSettingsBuilder#setUseMarks(boolean)
+         *
+         * @since 0.1.0
+         */
+        public static final String LOAD_USE_MARKS = "LOAD_USE_MARKS";
+
+        /**
+         * @see org.yaml.snakeyaml.DumperOptions#setExplicitStart(boolean)
          * @see org.snakeyaml.engine.v2.api.DumpSettingsBuilder#setExplicitStart(boolean)
          */
         public static final String DUMP_EXPLICIT_START = "DUMP_EXPLICIT_START";
 
         /**
+         * @see org.yaml.snakeyaml.DumperOptions#setExplicitEnd(boolean)
          * @see org.snakeyaml.engine.v2.api.DumpSettingsBuilder#setExplicitEnd(boolean)
          */
         public static final String DUMP_EXPLICIT_END = "DUMP_EXPLICIT_END";

diff --git a/src/main/java/io/xlate/yamljson/YamlGeneratorFactory.java b/src/main/java/io/xlate/yamljson/YamlGeneratorFactory.java
@@ -21,6 +21,7 @@
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
 
@@ -29,20 +30,20 @@
 
 class YamlGeneratorFactory implements JsonGeneratorFactory, SettingsBuilder {
 
-    private final Map<String, ?> properties;
+    private final Map<String, Object> properties;
     private final boolean useSnakeYamlEngine;
     private final Object snakeYamlSettings;
 
     YamlGeneratorFactory(Map<String, ?> properties) {
-        this.properties = properties;
+        this.properties = new HashMap<>(properties);
 
         Object version = properties.get(Yaml.Settings.YAML_VERSION);
         useSnakeYamlEngine = Yaml.Versions.V1_2.equals(version);
 
         if (useSnakeYamlEngine) {
-            snakeYamlSettings = loadProvider(() -> buildDumpSettings(properties), MOD_SNAKEYAML_ENGINE);
+            snakeYamlSettings = loadProvider(() -> buildDumpSettings(this.properties), MOD_SNAKEYAML_ENGINE);
         } else {
-            snakeYamlSettings = loadProvider(() -> buildDumperOptions(properties), MOD_SNAKEYAML);
+            snakeYamlSettings = loadProvider(() -> buildDumperOptions(this.properties), MOD_SNAKEYAML);
         }
     }
 

diff --git a/src/main/java/io/xlate/yamljson/YamlLocation.java b/src/main/java/io/xlate/yamljson/YamlLocation.java
@@ -0,0 +1,32 @@
+package io.xlate.yamljson;
+
+import jakarta.json.stream.JsonLocation;
+
+class YamlLocation implements JsonLocation {
+
+    final long lineNumber;
+    final long columnNumber;
+    final long streamOffset;
+
+    public YamlLocation(long lineNumber, long columnNumber, long streamOffset) {
+        this.lineNumber = lineNumber;
+        this.columnNumber = columnNumber;
+        this.streamOffset = streamOffset;
+    }
+
+    @Override
+    public long getLineNumber() {
+        return lineNumber;
+    }
+
+    @Override
+    public long getColumnNumber() {
+        return columnNumber;
+    }
+
+    @Override
+    public long getStreamOffset() {
+        return streamOffset;
+    }
+
+}
diff --git a/src/main/java/io/xlate/yamljson/YamlParser.java b/src/main/java/io/xlate/yamljson/YamlParser.java
@@ -37,7 +37,7 @@
 import jakarta.json.stream.JsonParser;
 import jakarta.json.stream.JsonParsingException;
 
-abstract class YamlParser<E, M> implements JsonParser, JsonLocation {
+abstract class YamlParser<E, M> implements JsonParser {
 
     enum NumberType {
         INTEGER,
@@ -79,9 +79,12 @@ enum NumberType {
     static final Set<String> VALUES_NAN = Set.of(YamlNumbers.CANONICAL_NAN, ".NaN", ".NAN");
 
     static final BigDecimal UNSET_NUMBER = new BigDecimal(0);
+    static final YamlLocation UNKNOWN_LOCATION = new YamlLocation(-1, -1, -1);
 
     final Reader yamlSource;
     final Iterator<E> yamlEvents;
+    final Map<String, ?> properties;
+    final long maxAliasExpansionSize;
 
     final Deque<E> yamlEventQueue = new ArrayDeque<>();
     final Deque<Boolean> aliasExpansionQueue = new ArrayDeque<>();
@@ -138,9 +141,11 @@ static class AnchoredDataEvent<E> extends AnchoredEvent<E> {
 
     final Map<String, List<AnchoredEvent<E>>> anchoredEvents = new HashMap<>();
 
-    YamlParser(Iterator<E> yamlEvents, Reader yamlReader) {
+    YamlParser(Iterator<E> yamlEvents, Reader yamlReader, Map<String, ?> properties) {
         this.yamlEvents = yamlEvents;
         this.yamlSource = yamlReader;
+        this.properties = properties;
+        this.maxAliasExpansionSize = (Long) properties.get(Yaml.Settings.LOAD_MAX_ALIAS_EXPANSION_SIZE);
     }
 
     void advanceEvent() {
@@ -413,10 +418,40 @@ void enqueueAlias(E yamlEvent, Boolean needKeyName) {
             // Enqueue the alias event, specifying that the JSON event is KEY_NAME
             enqueue(yamlEvent, Event.KEY_NAME, NumberType.NONE, "", UNSET_NUMBER);
         } else {
+            if (maxAliasExpansionSize < Long.MAX_VALUE) {
+                long expansionSize = countExpansion(alias, maxAliasExpansionSize);
+
+                if (expansionSize >= maxAliasExpansionSize) {
+                    String message = String.format("Alias '%s' expands to too many scalars: %d", alias, expansionSize);
+                    throw new JsonParsingException(message, getLocation(yamlEvent));
+                }
+            }
+
             enqueue(yamlEvent, Event.VALUE_NULL, NumberType.NONE, "", UNSET_NUMBER);
         }
     }
 
+    long countExpansion(String alias, long limit) {
+        List<AnchoredEvent<E>> anchored = anchoredEvents.get(alias);
+        long count = 0;
+
+        for (AnchoredEvent<E> event : anchored) {
+            String nestedAlias = getAlias(event.yamlEvent);
+
+            if (nestedAlias != null) {
+                count += countExpansion(nestedAlias, limit);
+            } else if ("Scalar".equals(getEventId(event.yamlEvent))) {
+                count++;
+            }
+
+            if (count >= limit) {
+                break;
+            }
+        }
+
+        return count;
+    }
+
     Boolean isKeyExpected() {
         return depth > -1 ? this.valueIsKey[depth] : null;
     }
@@ -533,14 +568,16 @@ void fillQueues() {
                     break;
                 }
             }
+        } catch (JsonException je) {
+            throw je;
         } catch (Exception re) {
             Throwable cause = re.getCause();
 
             if (cause instanceof IOException) {
                 throw new JsonException("IOException encountered reading YAML", cause);
             }
 
-            throw new JsonParsingException("Exception reading YAML", re, this);
+            throw new JsonParsingException("Exception reading YAML", re, getLocation());
         }
     }
 
@@ -655,32 +692,25 @@ public boolean isNaN() {
         return this.currentNumberType == NumberType.NAN;
     }
 
-    // JsonLocation
-
     @Override
     public JsonLocation getLocation() {
-        return this;
+        return getLocation(currentYamlEvent);
     }
 
-    @Override
-    public long getLineNumber() {
-        M mark = getMark();
-        return mark != null ? getMarkLine(mark) + 1 : -1;
-    }
+    JsonLocation getLocation(E yamlEvent) {
+        M mark = getMark(yamlEvent);
+        YamlLocation location = null;
 
-    @Override
-    public long getColumnNumber() {
-        M mark = getMark();
-        return mark != null ? getMarkColumn(mark) + 1 : -1;
-    }
+        if (mark != null) {
+            location = new YamlLocation(getMarkLine(mark) + 1L, getMarkColumn(mark) + 1L, getMarkIndex(mark));
+        } else {
+            location = UNKNOWN_LOCATION;
+        }
 
-    @Override
-    public long getStreamOffset() {
-        M mark = getMark();
-        return mark != null ? getMarkIndex(mark) : -1;
+        return location;
     }
 
-    protected abstract M getMark();
+    protected abstract M getMark(E event);
     protected abstract int getMarkLine(M mark);
     protected abstract int getMarkColumn(M mark);
     protected abstract int getMarkIndex(M mark);