recursive ref tests

onthegomap · Sep 29, 2024 · 6233221 · 6233221
1 parent 948d245
commit 6233221
Show file tree

Hide file tree

Showing 2 changed files with 83 additions and 7 deletions.
diff --git a/planetiler-core/src/main/java/com/onthegomap/planetiler/util/YAML.java b/planetiler-core/src/main/java/com/onthegomap/planetiler/util/YAML.java
@@ -1,16 +1,20 @@
 package com.onthegomap.planetiler.util;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.onthegomap.planetiler.reader.FileFormatException;
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.UncheckedIOException;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.Collections;
+import java.util.IdentityHashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import org.snakeyaml.engine.v2.api.Load;
 import org.snakeyaml.engine.v2.api.LoadSettings;
 
@@ -43,32 +47,43 @@ public static <T> T load(InputStream stream, Class<T> clazz) {
     }
   }
 
+  private static void handleMergeOperator(Object parsed) {
+    handleMergeOperator(parsed, Collections.newSetFromMap(new IdentityHashMap<>()));
+  }
+
   /**
    * SnakeYaml doesn't handle the <a href="https://yaml.org/type/merge.html">merge operator</a> so manually post-process
    * the parsed yaml object to merge referenced objects into the parent one.
    */
-  private static void handleMergeOperator(Object parsed) {
+  private static void handleMergeOperator(Object parsed, Set<Object> parentNodes) {
+    if (!parentNodes.add(parsed)) {
+      throw new FileFormatException("Illegal recursive reference in yaml file");
+    }
     if (parsed instanceof Map<?, ?> map) {
       Object toMerge = map.remove("<<");
       if (toMerge != null) {
         var orig = new LinkedHashMap<>(map);
         // to preserve the map key order we insert the merged operator objects first, then the original ones
         map.clear();
-        mergeInto(map, toMerge, false);
-        mergeInto(map, orig, true);
+        mergeInto(map, toMerge, false, parentNodes);
+        mergeInto(map, orig, true, parentNodes);
       }
       for (var value : map.values()) {
-        handleMergeOperator(value);
+        handleMergeOperator(value, parentNodes);
       }
     } else if (parsed instanceof List<?> list) {
       for (var item : list) {
-        handleMergeOperator(item);
+        handleMergeOperator(item, parentNodes);
       }
     }
+    parentNodes.remove(parsed);
   }
 
   @SuppressWarnings("rawtypes")
-  private static void mergeInto(Map dest, Object source, boolean replace) {
+  private static void mergeInto(Map dest, Object source, boolean replace, Set<Object> parentNodes) {
+    if (!parentNodes.add(source)) {
+      throw new FileFormatException("Illegal recursive reference in yaml file");
+    }
     if (source instanceof Map<?, ?> map) {
       if (replace) {
         dest.putAll(map);
@@ -77,9 +92,10 @@ private static void mergeInto(Map dest, Object source, boolean replace) {
       }
     } else if (source instanceof List<?> nesteds) {
       for (var nested : nesteds) {
-        mergeInto(dest, nested, replace);
+        mergeInto(dest, nested, replace, parentNodes);
       }
     }
+    parentNodes.remove(source);
   }
 
   public static <T> T load(String config, Class<T> clazz) {

diff --git a/planetiler-core/src/test/java/com/onthegomap/planetiler/util/YamlTest.java b/planetiler-core/src/test/java/com/onthegomap/planetiler/util/YamlTest.java
@@ -1,7 +1,9 @@
 package com.onthegomap.planetiler.util;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
+import com.onthegomap.planetiler.reader.FileFormatException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -235,6 +237,64 @@ void testMergeOperatorFromDraft3() {
       """);
   }
 
+  @Test
+  void testAnchorAndAliasMap() {
+    assertSameYaml("""
+      source: &label
+        a: 1
+      dest: *label
+      """, """
+      source:
+        a: 1
+      dest:
+        a: 1
+      """);
+  }
+
+  @Test
+  void testAnchorAndAliasList() {
+    assertSameYaml("""
+      source: &label
+        - 1
+      dest: *label
+      """, """
+      source: [1]
+      dest: [1]
+      """);
+  }
+
+  @Test
+  void testAllowRefInMergeDoc() {
+    assertSameYaml("""
+      source: &label
+        a: &label1
+          c: 1
+        b: *label1
+        d:
+          <<: *label1
+      dest: *label
+      """, """
+      source: {a: {c: 1}, b: {c: 1}, d: {c: 1}}
+      dest: {a: {c: 1}, b: {c: 1}, d: {c: 1}}
+      """);
+  }
+
+  @Test
+  void testFailsOnRecursiveRefs() {
+    assertThrows(FileFormatException.class, () -> YAML.load("""
+      source: &label
+        - *label
+      """, Object.class));
+    assertThrows(FileFormatException.class, () -> YAML.load("""
+      source: &label
+        <<: *label
+      """, Object.class));
+    assertThrows(FileFormatException.class, () -> YAML.load("""
+      source: &label
+        a: *label
+      """, Object.class));
+  }
+
   private static void assertSameYaml(String a, String b) {
     assertEquals(
       YAML.load(b, Object.class),