diff --git a/Sources/Parsing/EventStreamComposer.ExpansionBudget.swift b/Sources/Parsing/EventStreamComposer.ExpansionBudget.swift new file mode 100644 index 0000000..1e8a2a8 --- /dev/null +++ b/Sources/Parsing/EventStreamComposer.ExpansionBudget.swift @@ -0,0 +1,26 @@ +extension PureYAML.Parsing.EventStreamComposer { + /// Charges `cost` new expanded nodes against the document's remaining + /// ``EventStreamComposer/expandedNodeBudget``, throwing `expansionTooLarge` the moment it + /// runs out. + /// + /// `enforceNestingLimit` (`Parser.Events.swift`) bounds PHYSICAL collection depth in the raw + /// event stream; it does nothing against an anchor-expansion ("billion laughs") bomb, where a + /// handful of shallow, small anchors are each aliased many times, each alias nested inside the + /// next, so the LOGICAL tree an alias resolves to is exponential in the number of alias levels + /// while the physical document stays kilobytes -- a 25-level, fan-out-9 document already + /// expands past 9^25 logical nodes. + /// + /// Called once per scalar, alias resolution, or completed container -- never re-derived from + /// a walk -- so a bomb is rejected within a few alias levels, long before its true size could + /// ever be materialized: each charge is the MARGINAL cost this one node adds, with any nested + /// alias expansion it contains already folded into a prior charge (cached in + /// ``EventStreamComposer/anchorSizes``) and looked up in O(1), never re-counted here. + /// ``EventStreamComposer/expandedNodeBudget``'s initial 5,000,000 is far past any real + /// document's node count. + mutating func chargeExpansion(_ cost: Int, mark: PureYAML.Parsing.Mark) throws { + expandedNodeBudget -= cost + guard expandedNodeBudget >= 0 else { + throw PureYAML.Parsing.ParseError.expansionTooLarge(line: mark.line, column: mark.column) + } + } +} diff --git a/Sources/Parsing/EventStreamComposer.swift b/Sources/Parsing/EventStreamComposer.swift index f1498e1..b33b026 100644 --- a/Sources/Parsing/EventStreamComposer.swift +++ b/Sources/Parsing/EventStreamComposer.swift @@ -6,6 +6,10 @@ extension PureYAML.Parsing { // default value path never sets this, so it stays byte-identical to before. var tracksLocations = false var anchors: [String: PureYAML.Model.Value] = [:] + /// Each anchor's cached EXPANDED node count; see `chargeExpansion` for why. + var anchorSizes: [String: Int] = [:] + // The expanded-node budget remaining for this document; see `chargeExpansion`. + var expandedNodeBudget = 5_000_000 var containers: [EventStreamContainer] = [] var documentValues: [PureYAML.Model.Value] = [] var documentLocated: [PureYAML.Model.LocatedValue] = [] @@ -103,6 +107,8 @@ private extension PureYAML.Parsing.EventStreamComposer { throw PureYAML.Parsing.ParseError.unsupportedMultiDocumentStream(line: mark.line) } anchors = [:] + anchorSizes = [:] + expandedNodeBudget = 5_000_000 currentDocumentValue = nil } @@ -126,8 +132,10 @@ private extension PureYAML.Parsing.EventStreamComposer { marks: (start: PureYAML.Parsing.Mark, end: PureYAML.Parsing.Mark), ) throws { let model = try scalarParser.composeScalarValue(value, tag: tag, style: style, mark: marks.start) + try chargeExpansion(1, mark: marks.start) if let anchor { anchors[anchor] = model + anchorSizes[anchor] = 1 } try complete(.init( value: model, @@ -139,7 +147,7 @@ private extension PureYAML.Parsing.EventStreamComposer { ), mark: marks.start, located: tracksLocations ? .scalar(model, range(start: marks.start, end: marks.end)) : nil, - )) + ), cost: 1) } mutating func consumeAlias( @@ -153,6 +161,12 @@ private extension PureYAML.Parsing.EventStreamComposer { column: mark.column, ) } + // The target's own cached expanded size (itself plus every descendant, its own nested + // aliases already resolved): an alias to a value built from other aliases charges the + // FULL compounded cost here, in O(1), which is exactly what makes a nested bomb's true + // size visible without ever materializing it. + let cost = anchorSizes[anchor] ?? 1 + try chargeExpansion(cost, mark: mark) try complete(.init( value: value, mappingKey: .init( @@ -163,15 +177,20 @@ private extension PureYAML.Parsing.EventStreamComposer { ), mark: mark, located: tracksLocations ? .synthesized(from: value, at: range(at: mark)) : nil, - )) + ), cost: cost) } mutating func completeSequence(mark: PureYAML.Parsing.Mark) throws { guard let container = containers.popLast(), container.kind == .sequence else { throw unexpectedEvent(expected: "sequence start", mark: mark) } + // Only the container's own wrapper is a NEW charge here; every child already charged its + // own cost (including any alias expansion) as it completed, accumulated into + // `logicalSize`. The cache for a later alias to THIS container needs the full total. + try chargeExpansion(1, mark: container.mark) + let totalCost = container.logicalSize + 1 let value = PureYAML.Model.Value.sequence(container.values.map(\.value)) - store(value, anchor: container.anchor) + store(value, anchor: container.anchor, cost: totalCost) try complete(.init( value: value, mappingKey: .init( @@ -185,7 +204,7 @@ private extension PureYAML.Parsing.EventStreamComposer { located: tracksLocations ? .sequence(container.locatedValues, range(start: container.mark, end: mark)) : nil, - )) + ), cost: totalCost) } mutating func completeMapping(mark: PureYAML.Parsing.Mark) throws { @@ -195,12 +214,14 @@ private extension PureYAML.Parsing.EventStreamComposer { guard container.pendingKey == nil else { throw unexpectedEvent(expected: "mapping value", mark: mark) } + try chargeExpansion(1, mark: container.mark) + let totalCost = container.logicalSize + 1 let pairs = mergedPairs( from: container.mergeSources, localPairs: container.localPairs, ) let value = PureYAML.Model.Value.mapping(.init(pairs)) - store(value, anchor: container.anchor) + store(value, anchor: container.anchor, cost: totalCost) let located: PureYAML.Model.LocatedValue? = tracksLocations ? .mapping( mergedLocatedPairs(from: container.mergeLocatedSources, localPairs: container.localLocatedPairs), @@ -217,10 +238,10 @@ private extension PureYAML.Parsing.EventStreamComposer { ), mark: container.mark, located: located, - )) + ), cost: totalCost) } - mutating func complete(_ node: PureYAML.Parsing.EventStreamComposedNode) throws { + mutating func complete(_ node: PureYAML.Parsing.EventStreamComposedNode, cost: Int) throws { guard !containers.isEmpty else { currentDocumentValue = node.value if tracksLocations { @@ -232,11 +253,13 @@ private extension PureYAML.Parsing.EventStreamComposer { switch containers[containers.count - 1].kind { case .sequence: containers[containers.count - 1].values.append(node) + containers[containers.count - 1].logicalSize += cost if tracksLocations { containers[containers.count - 1].locatedValues.append(locatedNode(node)) } case .mapping: try completeMappingNode(node) + containers[containers.count - 1].logicalSize += cost } } @@ -277,11 +300,13 @@ private extension PureYAML.Parsing.EventStreamComposer { mutating func store( _ value: PureYAML.Model.Value, anchor: String?, + cost: Int, ) { guard let anchor else { return } anchors[anchor] = value + anchorSizes[anchor] = cost } func mergeSources( diff --git a/Sources/Parsing/EventStreamContainer.swift b/Sources/Parsing/EventStreamContainer.swift index 74bc644..2538e13 100644 --- a/Sources/Parsing/EventStreamContainer.swift +++ b/Sources/Parsing/EventStreamContainer.swift @@ -7,6 +7,11 @@ extension PureYAML.Parsing { var localPairs: [PureYAML.Model.Pair] = [] var pendingKey: PureYAML.Parsing.EventComposer.ComposedMappingKey? var mark: PureYAML.Parsing.Mark + /// The running sum of every child's own expanded-node cost (see + /// `EventStreamComposer.expandedNodeBudget`), so this container's total logical size -- + /// itself plus every descendant, alias expansion already accounted for -- is available in + /// O(1) the moment it completes, without re-walking anything. + var logicalSize = 0 // Location tracking (parallel to the above, only when the composer tracks locations). var locatedValues: [PureYAML.Model.LocatedValue] = [] var mergeLocatedSources: [[PureYAML.Model.LocatedPair]] = [] diff --git a/Sources/Parsing/ParseError.swift b/Sources/Parsing/ParseError.swift index 4d95610..a76aa93 100644 --- a/Sources/Parsing/ParseError.swift +++ b/Sources/Parsing/ParseError.swift @@ -23,6 +23,13 @@ public extension PureYAML.Parsing { case unexpectedEvent(expected: String, actual: String, line: Int, column: Int) case unexpectedToken(expected: String, actual: String, line: Int, column: Int) case nestingTooDeep(line: Int, column: Int) + /// The document's total EXPANDED node count (after resolving every `*alias` to its + /// target) exceeds the budget. `nestingTooDeep` bounds physical collection depth in the + /// event stream; it does nothing against an anchor-expansion ("billion laughs") bomb, + /// where a handful of shallow, small anchors are each aliased many times, each alias + /// nested inside the next, so the LOGICAL tree an alias resolves to is exponential in + /// the number of alias levels while the physical document stays tiny. + case expansionTooLarge(line: Int, column: Int) case malformed(reason: String, line: Int, column: Int) public var description: String { @@ -71,6 +78,8 @@ public extension PureYAML.Parsing { "expected \(expected) at line \(line), column \(column), found \(actual)" case let .nestingTooDeep(line, column): "nesting is too deep at line \(line), column \(column)" + case let .expansionTooLarge(line, column): + "the document's expanded size (after resolving aliases) is too large at line \(line), column \(column)" case let .malformed(reason, line, column): "\(reason) at line \(line), column \(column)" } @@ -122,6 +131,8 @@ public extension PureYAML.Parsing { "unexpectedToken" case .nestingTooDeep: "nestingTooDeep" + case .expansionTooLarge: + "expansionTooLarge" case .malformed: "malformed" } @@ -152,6 +163,7 @@ public extension PureYAML.Parsing { let .unexpectedEvent(_, _, line, _), let .unexpectedToken(_, _, line, _), let .nestingTooDeep(line, _), + let .expansionTooLarge(line, _), let .invalidControlCharacter(line, _), let .malformed(_, line, _): line @@ -182,6 +194,7 @@ public extension PureYAML.Parsing { let .unexpectedEvent(_, _, _, column), let .unexpectedToken(_, _, _, column), let .nestingTooDeep(_, column), + let .expansionTooLarge(_, column), let .invalidControlCharacter(_, column), let .malformed(_, _, column): column diff --git a/Tests/ParsingAliasTests.swift b/Tests/ParsingAliasTests.swift index 7ef19a7..6f75a99 100644 --- a/Tests/ParsingAliasTests.swift +++ b/Tests/ParsingAliasTests.swift @@ -19,4 +19,46 @@ struct ParsingAliasTests { } #expect(second["name"] == .string("Example")) } + + /// The regression this guards: `enforceNestingLimit` (Parser.Events.swift) bounds PHYSICAL + /// collection depth in the raw event stream, which does nothing against an anchor-expansion + /// ("billion laughs") bomb -- a handful of shallow, small anchors each aliased many times, + /// each alias nested inside the next, so the LOGICAL tree an alias resolves to is exponential + /// in the number of alias levels while the physical document stays kilobytes. A 25-level, + /// fan-out-9 document like this one expands past 9^25 logical nodes; before this fix, + /// `PureYAML.parse` accepted it in under 2ms (the nesting cap never fires, since raw event + /// depth per line stays at 2) and any full-tree walk downstream (JSON representability + /// checking, JSON writing, validation) never terminated. + @Test("Rejects a nested anchor-expansion bomb instead of accepting an astronomically large tree") + func test_rejectsAnchorExpansionBomb() { + var yaml = "a0: &a0 [x, x, x]\n" + for level in 1 ... 25 { + let previous = "*a\(level - 1)" + let refs = Array(repeating: previous, count: 9).joined(separator: ", ") + yaml += "a\(level): &a\(level) [\(refs)]\n" + } + yaml += "bomb: *a25\n" + + #expect { + try PureYAML.parse(yaml) + } throws: { error in + guard case .expansionTooLarge = error as? PureYAML.Parsing.ParseError else { return false } + return true + } + } + + /// The fix must not reject ordinary, non-exponential alias reuse: the same small anchor + /// referenced many times at ONE level (no nesting) stays comfortably inside the budget. + @Test("A large but non-exponential alias fan-out still parses") + func test_wideNonExponentialFanOutStillParses() throws { + let refs = Array(repeating: "*item", count: 5000).joined(separator: ", ") + let yaml = "item: &item {name: Example}\nlist: [\(refs)]\n" + + let root = try requireMapping(PureYAML.parse(yaml)) + guard case let .sequence(list)? = root?["list"] else { + recordIssue("expected a sequence") + return + } + #expect(list.count == 5000) + } }