Prevent Stack Overflow With Nested Structures (#5)

Author: fabianfett

Sources/PureSwiftJSONParsing/JSONParser.swift

@@ -18,13 +18,21 @@ public struct JSONParser {
 @usableFromInline struct JSONParserImpl {
 
   @usableFromInline var reader: DocumentReader
+  @usableFromInline var depth : Int = 0
 
   @inlinable init<Bytes: Collection>(bytes: Bytes) where Bytes.Element == UInt8 {
     self.reader = DocumentReader(bytes: bytes)
   }
 
   @usableFromInline mutating func parse() throws -> JSONValue {
     let value = try parseValue()
+    #if DEBUG
+    defer {
+      guard self.depth == 0 else {
+        preconditionFailure()
+      }
+    }
+    #endif
 
     // handle extra character if top level was number
     if case .number(_) = value {
@@ -236,13 +244,19 @@ public struct JSONParser {
   }
 
   mutating func parseArray() throws -> [JSONValue] {
-    // parse first value or immidiate end
-    precondition(reader.value == UInt8(ascii: "["))
+    assert(reader.value == UInt8(ascii: "["))
+    guard depth < 512 else {
+      throw JSONError.tooManyNestedArraysOrDictionaries(characterIndex: reader.index)
+    }
+    depth += 1
+    defer { depth -= 1 }
     var state = ArrayState.expectValueOrEnd
 
     var array = [JSONValue]()
     array.reserveCapacity(10)
 
+    // parse first value or immidiate end
+    
     do {
       let value = try parseValue()
       array.append(value)
@@ -336,6 +350,13 @@ public struct JSONParser {
   }
 
   mutating func parseObject() throws -> [String: JSONValue] {
+    assert(reader.value == UInt8(ascii: "{"))
+    guard depth < 512 else {
+      throw JSONError.tooManyNestedArraysOrDictionaries(characterIndex: reader.index)
+    }
+    depth += 1
+    defer { depth -= 1 }
+    
     var state = ObjectState.expectKeyOrEnd
 
     // parse first key or end immidiatly

Sources/PureSwiftJSONParsing/JSONValue.swift

@@ -2,6 +2,7 @@
 public enum JSONError: Swift.Error {
   case unexpectedCharacter(ascii: UInt8)
   case unexpectedEndOfFile
+  case tooManyNestedArraysOrDictionaries(characterIndex: Int)
 }
 
 public enum JSONValue {

Tests/JSONParsingTests/ArrayParserTests.swift

@@ -56,4 +56,21 @@ class ArrayParserTests: XCTestCase {
       XCTFail("Unexpected error: \(error)")
     }
   }
+  
+  func testHighlyNestedArray() throws {
+    // test 512 should succeed
+    let passingString = String(repeating: "[", count: 512) +  String(repeating: "]", count: 512)
+    _  = try JSONParser().parse(bytes: [UInt8](passingString.utf8))
+    
+    let failingString = String(repeating: "[", count: 513)
+    do {
+      _  = try JSONParser().parse(bytes: [UInt8](failingString.utf8))
+    }
+    catch JSONError.tooManyNestedArraysOrDictionaries(characterIndex: 512) {
+      //expected case
+    }
+    catch {
+      XCTFail("Unexpected error: \(error)")
+    }
+  }
 }

Tests/JSONParsingTests/JSONParserTests.swift

@@ -65,6 +65,23 @@ class JSONParserTests: XCTestCase {
     catch {
       XCTFail("Unexpected error: \(error)")
     }
+  }
+  
+  func testHighlyNestedObjectAndDictionary() throws {
+    // test 512 should succeed
+    let passingString = String(repeating: #"{"a":["#, count: 256) + "null" + String(repeating: "]}", count: 256)
+    _  = try JSONParser().parse(bytes: [UInt8](passingString.utf8))
 
+    let failingString = String(repeating: #"{"a":["#, count: 257)
+    do {
+      _  = try JSONParser().parse(bytes: [UInt8](failingString.utf8))
+    }
+    catch JSONError.tooManyNestedArraysOrDictionaries(characterIndex: 1536) {
+      //expected case
+    }
+    catch {
+      XCTFail("Unexpected error: \(error)")
+    }
   }
+
 }

Tests/JSONParsingTests/ObjectParserTests.swift

@@ -27,4 +27,21 @@ class ObjectParserTests: XCTestCase {
     let result = try parser.parseObject()
     XCTAssertEqual(result, ["hello": .string("world"), "haha": .bool(true)])
   }
+  
+  func testHighlyNestedObject() throws {
+    // test 512 should succeed
+    let passingString = String(repeating: #"{"a":"#, count: 512) + "null" + String(repeating: "}", count: 512)
+    _  = try JSONParser().parse(bytes: [UInt8](passingString.utf8))
+    
+    let failingString = String(repeating: #"{"a":"#, count: 513)
+    do {
+      _  = try JSONParser().parse(bytes: [UInt8](failingString.utf8))
+    }
+    catch JSONError.tooManyNestedArraysOrDictionaries(characterIndex: 2560) {
+      //expected case
+    }
+    catch {
+      XCTFail("Unexpected error: \(error)")
+    }
+  }
 }