Merge pull request #384 from mattpolzin/bugfix/382/url_template_stack…

…_overflow2 Fix `URLTemplate` stack overflow
mattpolzin · Oct 4, 2024 · 35dd374 · 35dd374
2 parents 03840ad + 24d7950
commit 35dd374
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 53 deletions.
diff --git a/Sources/OpenAPIKitCore/URLTemplate/URLTemplate+Parsing.swift b/Sources/OpenAPIKitCore/URLTemplate/URLTemplate+Parsing.swift
@@ -7,64 +7,55 @@
 
 extension URLTemplate {
     internal static func scan(
-        _ string: String,
-        partialToken: PartialToken?,
-        from remainder: Substring,
-        addingTo tokens: [Component]
+        _ string: String
     ) throws -> [Component] {
-        guard let next = remainder.first else {
-            guard partialToken == nil || partialToken?.type == .constant else {
-                throw ParsingError.unterminatedVariable(name: String(partialToken?.string ?? ""))
-            }
-            return tokens + tokenArray(from: partialToken)
-        }
-        let nextFirstIndex = remainder.index(remainder.startIndex, offsetBy: 1, limitedBy: remainder.endIndex)
+        var tokens = [Component]()
+        var remainder = string[...]
+        var partialToken: PartialToken? = nil
+
+        while let next = remainder.first {
+            let nextFirstIndex = remainder.index(remainder.startIndex, offsetBy: 1, limitedBy: remainder.endIndex)
 
-        switch (partialToken?.type, next) {
-        case (nil, "{"),
-             (.constant, "{"):
-            guard let newFirstIndex = nextFirstIndex else {
-                throw ParsingError.unterminatedVariable(name: "")
-            }
-            let newTokens = tokens + tokenArray(from: partialToken)
-            return try scan(
-                string,
-                partialToken: .init(type: .variable, string: remainder[newFirstIndex..<newFirstIndex]),
-                from: remainder.dropFirst(),
-                addingTo: newTokens
-            )
+            switch (partialToken?.type, next) {
+            case (nil, "{"),
+                (.constant, "{"):
+                guard let newFirstIndex = nextFirstIndex else {
+                    throw ParsingError.unterminatedVariable(name: "")
+                }
+                tokens += tokenArray(from: partialToken)
+                partialToken = .init(type: .variable, string: remainder[newFirstIndex..<newFirstIndex])
+                remainder = remainder.dropFirst()
 
-        case (.variable, "}"):
-            let newTokens = tokens + tokenArray(from: partialToken)
-            return try scan(string, partialToken: nil, from: remainder.dropFirst(), addingTo: newTokens)
+            case (.variable, "}"):
+                tokens += tokenArray(from: partialToken)
+                partialToken = nil
+                remainder = remainder.dropFirst()
 
-        case (nil, "}"),
-             (.constant, "}"):
-            throw ParsingError.variableEndedWithoutStarting(name: partialToken.map { String($0.string) } ?? "")
+            case (nil, "}"),
+                (.constant, "}"):
+                throw ParsingError.variableEndedWithoutStarting(name: partialToken.map { String($0.string) } ?? "")
 
-        case (.variable, "{"):
-            throw ParsingError.variableStartedWithinVariable(name: partialToken.map { String($0.string) } ?? "")
+            case (.variable, "{"):
+                throw ParsingError.variableStartedWithinVariable(name: partialToken.map { String($0.string) } ?? "")
 
-        case (nil, _):
-            return try scan(
-                string,
-                partialToken: .init(type: .constant, string: remainder[remainder.startIndex...remainder.startIndex]),
-                from: remainder.dropFirst(),
-                addingTo: tokens
-            )
+            case (nil, _):
+                partialToken = .init(type: .constant, string: remainder[remainder.startIndex...remainder.startIndex])
+                remainder = remainder.dropFirst()
 
-        case (.constant, _),
-             (.variable, _):
-            guard nextFirstIndex != nil, let reifiedPartialToken = partialToken else {
-                return tokens + tokenArray(from: partialToken)
+            case (.constant, _),
+                (.variable, _):
+                guard nextFirstIndex != nil, let reifiedPartialToken = partialToken else {
+                    tokens += tokenArray(from: partialToken)
+                    continue
+                }
+                    partialToken = reifiedPartialToken.advancingStringByOne(within: string)
+                    remainder = remainder.dropFirst()
             }
-            return try scan(
-                string,
-                partialToken: reifiedPartialToken.advancingStringByOne(within: string),
-                from: remainder.dropFirst(),
-                addingTo: tokens
-            )
         }
+        guard partialToken == nil || partialToken?.type == .constant else {
+            throw ParsingError.unterminatedVariable(name: String(partialToken?.string ?? ""))
+        }
+        return tokens + tokenArray(from: partialToken)
     }
 
     internal static func tokenArray(from partial: PartialToken?) -> [Component] {

diff --git a/Sources/OpenAPIKitCore/URLTemplate/URLTemplate.swift b/Sources/OpenAPIKitCore/URLTemplate/URLTemplate.swift
@@ -111,10 +111,7 @@ public struct URLTemplate: Hashable, RawRepresentable {
     public init(templateString: String) throws {
         rawValue = templateString
         components = try URLTemplate.scan(
-            templateString,
-            partialToken: nil,
-            from: templateString[...],
-            addingTo: []
+            templateString
         )
     }
 

diff --git a/Tests/OpenAPIKitCoreTests/URLTemplate/URLTemplateTests.swift b/Tests/OpenAPIKitCoreTests/URLTemplate/URLTemplateTests.swift
@@ -375,3 +375,37 @@ extension URLTemplateTests {
 fileprivate struct TemplatedURLWrapper: Codable {
     let url: URLTemplate?
 }
+
+// MARK: - Stack Overflow Regression Test
+#if swift(>=5.5)
+import Dispatch
+
+extension URLTemplateTests {
+    struct StackFoo: Decodable {
+        var val: URLTemplate
+    }
+
+    static func stackWork() throws {
+        let data = Data("""
+    {
+        \"val\": \"https://\(Array(repeating: "foo.", count: 1000).joined())com/\"
+    }
+
+    """.utf8)
+        let document = try JSONDecoder().decode(
+            StackFoo.self,
+            from: data
+        )
+        print(document)
+    }
+
+    func test_avoid_stack_overflow() async throws {
+        try await withThrowingTaskGroup(of: Void.self) { group in
+            group.addTask {
+                try URLTemplateTests.stackWork()
+            }
+            try await group.waitForAll()
+        }
+    }
+}
+#endif