feat: allow to set a maximum amount of aliases while parsing to prote…

…ct against billion laughs attack and similar (#620)

src/commonMain/kotlin/com/charleskorn/kaml/Yaml.kt

@@ -77,8 +77,11 @@ public class Yaml(
 
     internal fun parseToYamlNode(source: Source): YamlNode {
         val parser = YamlParser(source, configuration.codePointLimit)
-        val reader =
-            YamlNodeReader(parser, configuration.extensionDefinitionPrefix, configuration.allowAnchorsAndAliases)
+        val reader = YamlNodeReader(
+            parser,
+            configuration.extensionDefinitionPrefix,
+            configuration.anchorsAndAliases.maxAliasCount,
+        )
         val node = reader.read()
         parser.ensureEndOfStreamReached()
         return node

src/commonMain/kotlin/com/charleskorn/kaml/YamlConfiguration.kt

@@ -37,7 +37,7 @@ import kotlinx.serialization.ExperimentalSerializationApi
  * * [multiLineStringStyle]: the style in which a multi line String value is written. Can be overruled for a specific field with the [YamlMultiLineStringStyle] annotation.
  * * [ambiguousQuoteStyle]: how strings should be escaped when [singleLineStringStyle] is [SingleLineStringStyle.PlainExceptAmbiguous] and the value is ambiguous
  * * [sequenceBlockIndent]: number of spaces to use as indentation for sequences, if [sequenceStyle] set to [SequenceStyle.Block]
- * * [allowAnchorsAndAliases]: set to true to allow anchors and aliases when decoding YAML (defaults to `false`)
+ * * [anchorsAndAliases]: set to [AnchorsAndAliases.Permitted] to allow anchors and aliases when decoding YAML (defaults to [AnchorsAndAliases.Forbidden])
  * * [yamlNamingStrategy]: The system that converts the field names in to the names used in the Yaml.
  * * [codePointLimit]: the maximum amount of code points allowed in the input YAML document (defaults to 3 MB)
  * * [decodeEnumCaseInsensitive]: set to true to allow case-insensitive decoding of enums (defaults to `false`)
@@ -55,7 +55,7 @@ public data class YamlConfiguration(
     internal val multiLineStringStyle: MultiLineStringStyle = singleLineStringStyle.multiLineStringStyle,
     internal val ambiguousQuoteStyle: AmbiguousQuoteStyle = AmbiguousQuoteStyle.DoubleQuoted,
     internal val sequenceBlockIndent: Int = 0,
-    internal val allowAnchorsAndAliases: Boolean = false,
+    internal val anchorsAndAliases: AnchorsAndAliases = AnchorsAndAliases.Forbidden,
     internal val yamlNamingStrategy: YamlNamingStrategy? = null,
     internal val codePointLimit: Int? = null,
     @ExperimentalSerializationApi
@@ -124,3 +124,16 @@ public enum class AmbiguousQuoteStyle {
     DoubleQuoted,
     SingleQuoted,
 }
+
+public sealed class AnchorsAndAliases {
+    internal abstract val maxAliasCount: UInt?
+
+    public data object Forbidden : AnchorsAndAliases() {
+        override val maxAliasCount: UInt = 0u
+    }
+
+    /**
+     * [maxAliasCount]: the maximum amount of aliases allowed in the input YAML document if allowed at all, `null` allows any amount (defaults to `100`)
+     */
+    public data class Permitted(override val maxAliasCount: UInt? = 100u) : AnchorsAndAliases()
+}

src/commonMain/kotlin/com/charleskorn/kaml/YamlNodeReader.kt

@@ -29,37 +29,41 @@ import it.krzeminski.snakeyaml.engine.kmp.events.SequenceStartEvent
 internal class YamlNodeReader(
     private val parser: YamlParser,
     private val extensionDefinitionPrefix: String? = null,
-    private val allowAnchorsAndAliases: Boolean = false,
+    private val maxAliasCount: UInt? = 0u,
 ) {
-    private val aliases = mutableMapOf<Anchor, YamlNode>()
+    private val aliases = mutableMapOf<Anchor, WeightedNode>()
+    private var aliasCount = 0u
 
-    fun read(): YamlNode = readNode(YamlPath.root)
+    fun read(): YamlNode = readNode(YamlPath.root).node
 
-    private fun readNode(path: YamlPath): YamlNode = readNodeAndAnchor(path).first
+    private fun readNode(path: YamlPath): WeightedNode = readNodeAndAnchor(path).first
 
-    private fun readNodeAndAnchor(path: YamlPath): Pair<YamlNode, Anchor?> {
+    private fun readNodeAndAnchor(path: YamlPath): Pair<WeightedNode, Anchor?> {
         val event = parser.consumeEvent(path)
-        val node = readFromEvent(event, path)
+        val (node, weight) = readFromEvent(event, path)
 
         if (event is NodeEvent) {
-            event.anchor?.let {
-                if (!allowAnchorsAndAliases) {
-                    throw ForbiddenAnchorOrAliasException("Parsing anchors and aliases is disabled.", path)
-                }
+            if (event !is AliasEvent) {
+                event.anchor?.let {
+                    if (maxAliasCount == 0u) {
+                        throw ForbiddenAnchorOrAliasException("Parsing anchors and aliases is disabled.", path)
+                    }
 
-                aliases.put(it, node.withPath(YamlPath.forAliasDefinition(it.value, event.location)))
+                    val anchor = node.withPath(YamlPath.forAliasDefinition(it.value, event.location))
+                    aliases[it] = WeightedNode(anchor, weight)
+                }
             }
 
-            return node to event.anchor
+            return WeightedNode(node, weight) to event.anchor
         }
 
-        return node to null
+        return WeightedNode(node, weight = 0u) to null
     }
 
-    private fun readFromEvent(event: Event, path: YamlPath): YamlNode = when (event) {
-        is ScalarEvent -> readScalarOrNull(event, path).maybeToTaggedNode(event.tag)
-        is SequenceStartEvent -> readSequence(path).maybeToTaggedNode(event.tag)
-        is MappingStartEvent -> readMapping(path).maybeToTaggedNode(event.tag)
+    private fun readFromEvent(event: Event, path: YamlPath): WeightedNode = when (event) {
+        is ScalarEvent -> WeightedNode(readScalarOrNull(event, path).maybeToTaggedNode(event.tag), weight = 0u)
+        is SequenceStartEvent -> readSequence(path).let { it.copy(node = it.node.maybeToTaggedNode(event.tag)) }
+        is MappingStartEvent -> readMapping(path).let { it.copy(node = it.node.maybeToTaggedNode(event.tag)) }
         is AliasEvent -> readAlias(event, path)
         else -> throw MalformedYamlException("Unexpected ${event.eventId}", path.withError(event.location))
     }
@@ -72,33 +76,39 @@ internal class YamlNodeReader(
         }
     }
 
-    private fun readSequence(path: YamlPath): YamlList {
+    private fun readSequence(path: YamlPath): WeightedNode {
         val items = mutableListOf<YamlNode>()
+        var sequenceWeight = 0u
 
         while (true) {
             val event = parser.peekEvent(path)
 
             when (event.eventId) {
                 Event.ID.SequenceEnd -> {
                     parser.consumeEventOfType(Event.ID.SequenceEnd, path)
-                    return YamlList(items, path)
+                    return WeightedNode(YamlList(items, path), sequenceWeight)
                 }
 
-                else -> items += readNode(path.withListEntry(items.size, event.location))
+                else -> {
+                    val (node, weight) = readNode(path.withListEntry(items.size, event.location))
+                    sequenceWeight += weight
+                    items += node
+                }
             }
         }
     }
 
-    private fun readMapping(path: YamlPath): YamlMap {
+    private fun readMapping(path: YamlPath): WeightedNode {
         val items = mutableMapOf<YamlScalar, YamlNode>()
+        var mapWeight = 0u
 
         while (true) {
             val event = parser.peekEvent(path)
 
             when (event.eventId) {
                 Event.ID.MappingEnd -> {
                     parser.consumeEventOfType(Event.ID.MappingEnd, path)
-                    return YamlMap(doMerges(items), path)
+                    return WeightedNode(YamlMap(doMerges(items), path), mapWeight)
                 }
 
                 else -> {
@@ -108,14 +118,15 @@ internal class YamlNodeReader(
 
                     val valueLocation = parser.peekEvent(keyNode.path).location
                     val valuePath = if (isMerge(keyNode)) path.withMerge(valueLocation) else keyNode.path.withMapElementValue(valueLocation)
-                    val (value, anchor) = readNodeAndAnchor(valuePath)
+                    val (weightedNode, anchor) = readNodeAndAnchor(valuePath)
+                    mapWeight += weightedNode.weight
 
                     if (path == YamlPath.root && extensionDefinitionPrefix != null && key.startsWith(extensionDefinitionPrefix)) {
                         if (anchor == null) {
                             throw NoAnchorForExtensionException(key, extensionDefinitionPrefix, path.withError(event.location))
                         }
                     } else {
-                        items += (keyNode to value)
+                        items += (keyNode to weightedNode.node)
                     }
                 }
             }
@@ -189,22 +200,43 @@ internal class YamlNodeReader(
         return merged
     }
 
-    private fun readAlias(event: AliasEvent, path: YamlPath): YamlNode {
-        if (!allowAnchorsAndAliases) {
+    private fun readAlias(event: AliasEvent, path: YamlPath): WeightedNode {
+        if (maxAliasCount == 0u) {
             throw ForbiddenAnchorOrAliasException("Parsing anchors and aliases is disabled.", path)
         }
 
-        val anchor = event.anchor!!
+        val anchor = event.alias
 
-        val resolvedNode = aliases.getOrElse(anchor) {
+        val (resolvedNode, resolvedNodeWeight) = aliases.getOrElse(anchor) {
             throw UnknownAnchorException(anchor.value, path.withError(event.location))
         }
 
-        return resolvedNode.withPath(path.withAliasReference(anchor.value, event.location).withAliasDefinition(anchor.value, resolvedNode.location))
+        val resultWeight = resolvedNodeWeight + 1u
+        aliasCount += resultWeight
+
+        if ((maxAliasCount != null) && (aliasCount > maxAliasCount)) {
+            throw ForbiddenAnchorOrAliasException(
+                "Maximum number of aliases has been reached.",
+                path,
+            )
+        }
+
+        return WeightedNode(
+            node = resolvedNode.withPath(
+                path.withAliasReference(anchor.value, event.location)
+                    .withAliasDefinition(anchor.value, resolvedNode.location),
+            ),
+            weight = resultWeight,
+        )
     }
 
     private fun <T> Iterable<T>.second(): T = this.drop(1).first()
 
     private val Event.location: Location
         get() = Location(startMark!!.line + 1, startMark!!.column + 1)
 }
+
+private data class WeightedNode(
+    val node: YamlNode,
+    val weight: UInt,
+)

src/commonTest/kotlin/com/charleskorn/kaml/YamlReadingTest.kt

@@ -773,7 +773,7 @@ class YamlReadingTest : FlatFunSpec({
                 """.trimIndent()
 
                 context("parsing anchors and aliases is disabled") {
-                    val configuration = YamlConfiguration(extensionDefinitionPrefix = ".", allowAnchorsAndAliases = false)
+                    val configuration = YamlConfiguration(extensionDefinitionPrefix = ".", anchorsAndAliases = AnchorsAndAliases.Forbidden)
                     val yaml = Yaml(configuration = configuration)
 
                     test("throws an appropriate exception") {
@@ -788,7 +788,7 @@ class YamlReadingTest : FlatFunSpec({
                 }
 
                 context("parsing anchors and aliases is enabled") {
-                    val configuration = YamlConfiguration(extensionDefinitionPrefix = ".", allowAnchorsAndAliases = true)
+                    val configuration = YamlConfiguration(extensionDefinitionPrefix = ".", anchorsAndAliases = AnchorsAndAliases.Permitted())
                     val yaml = Yaml(configuration = configuration)
                     val result = yaml.decodeFromString(SimpleStructure.serializer(), input)
 
@@ -798,6 +798,109 @@ class YamlReadingTest : FlatFunSpec({
                 }
             }
 
+            context("restricting max alias count") {
+                val input = """
+                    .some-extension: &name Jamie
+
+                    members: [*name, *name]
+                """.trimIndent()
+
+                context("parsing anchors and aliases is disabled") {
+                    val configuration = YamlConfiguration(extensionDefinitionPrefix = ".", anchorsAndAliases = AnchorsAndAliases.Forbidden)
+                    val yaml = Yaml(configuration = configuration)
+
+                    test("throws an appropriate exception") {
+                        val exception = shouldThrow<ForbiddenAnchorOrAliasException> { yaml.decodeFromString(Team.serializer(), input) }
+
+                        exception.asClue {
+                            it.message shouldBe "Parsing anchors and aliases is disabled."
+                            it.line shouldBe 1
+                            it.column shouldBe 18
+                        }
+                    }
+                }
+
+                context("parsing anchors and aliases is enabled, max aliases count 0") {
+                    val configuration = YamlConfiguration(extensionDefinitionPrefix = ".", anchorsAndAliases = AnchorsAndAliases.Permitted(maxAliasCount = 0u))
+                    val yaml = Yaml(configuration = configuration)
+
+                    test("throws an appropriate exception") {
+                        val exception = shouldThrow<ForbiddenAnchorOrAliasException> { yaml.decodeFromString(Team.serializer(), input) }
+
+                        exception.asClue {
+                            it.message shouldBe "Parsing anchors and aliases is disabled."
+                            it.line shouldBe 1
+                            it.column shouldBe 18
+                        }
+                    }
+                }
+
+                context("parsing anchors and aliases is enabled, max aliases count 1") {
+                    val configuration = YamlConfiguration(extensionDefinitionPrefix = ".", anchorsAndAliases = AnchorsAndAliases.Permitted(maxAliasCount = 1u))
+                    val yaml = Yaml(configuration = configuration)
+
+                    test("throws an appropriate exception") {
+                        val exception = shouldThrow<ForbiddenAnchorOrAliasException> { yaml.decodeFromString(Team.serializer(), input) }
+
+                        exception.asClue {
+                            it.message shouldBe "Maximum number of aliases has been reached."
+                            it.line shouldBe 3
+                            it.column shouldBe 18
+                        }
+                    }
+                }
+
+                context("parsing anchors and aliases is enabled, max aliases count 2") {
+                    val configuration = YamlConfiguration(extensionDefinitionPrefix = ".", anchorsAndAliases = AnchorsAndAliases.Permitted(maxAliasCount = 2u))
+                    val yaml = Yaml(configuration = configuration)
+                    val result = yaml.decodeFromString(Team.serializer(), input)
+
+                    test("deserializes it to a Kotlin object, replacing the reference to the extension with the extension") {
+                        result shouldBe Team(listOf("Jamie", "Jamie"))
+                    }
+                }
+
+                context("parsing anchors and aliases is enabled, max aliases count null") {
+                    val configuration = YamlConfiguration(extensionDefinitionPrefix = ".", anchorsAndAliases = AnchorsAndAliases.Permitted(maxAliasCount = null))
+                    val yaml = Yaml(configuration = configuration)
+                    val result = yaml.decodeFromString(Team.serializer(), input)
+
+                    test("deserializes it to a Kotlin object, replacing the reference to the extension with the extension") {
+                        result shouldBe Team(listOf("Jamie", "Jamie"))
+                    }
+                }
+
+                context("parsing anchors and aliases is enabled, billion laughs is prevented by default") {
+                    val configuration = YamlConfiguration(extensionDefinitionPrefix = ".", anchorsAndAliases = AnchorsAndAliases.Permitted())
+                    val yaml = Yaml(configuration = configuration)
+
+                    test("throws an appropriate exception") {
+                        val exception = shouldThrow<ForbiddenAnchorOrAliasException> {
+                            yaml.decodeFromString(
+                                Team.serializer(),
+                                """
+                                    a: &a ["lol","lol","lol","lol","lol","lol","lol","lol","lol"]
+                                    b: &b [*a,*a,*a,*a,*a,*a,*a,*a,*a]
+                                    c: &c [*b,*b,*b,*b,*b,*b,*b,*b,*b]
+                                    d: &d [*c,*c,*c,*c,*c,*c,*c,*c,*c]
+                                    e: &e [*d,*d,*d,*d,*d,*d,*d,*d,*d]
+                                    f: &f [*e,*e,*e,*e,*e,*e,*e,*e,*e]
+                                    g: &g [*f,*f,*f,*f,*f,*f,*f,*f,*f]
+                                    h: &h [*g,*g,*g,*g,*g,*g,*g,*g,*g]
+                                    i: &i [*h,*h,*h,*h,*h,*h,*h,*h,*h]
+                                """.trimIndent(),
+                            )
+                        }
+
+                        exception.asClue {
+                            it.message shouldBe "Maximum number of aliases has been reached."
+                            it.line shouldBe 4
+                            it.column shouldBe 8
+                        }
+                    }
+                }
+            }
+
             context("given some input with an additional unknown field") {
                 val input = """
                     name: Blah Blahson