From fffd958fb20fdbe2af311e6dfe40dc4960c897e8 Mon Sep 17 00:00:00 2001
From: Dimitri Bouniol <dimitribouniol@mochidev.com>
Date: Thu, 15 Feb 2024 05:48:08 -0800
Subject: [PATCH] Added the ability to specify a maximum depth for CBOR
 decoding

Closes #92
---
 Sources/CBORDecoder.swift                |  7 +++++
 Sources/CBOROptions.swift                |  9 +++++--
 Sources/Decoder/CodableCBORDecoder.swift | 14 ++++++++--
 Tests/CBORDecoderTests.swift             | 33 ++++++++++++++++++++++++
 4 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/Sources/CBORDecoder.swift b/Sources/CBORDecoder.swift
index c5659a4..4e26806 100644
--- a/Sources/CBORDecoder.swift
+++ b/Sources/CBORDecoder.swift
@@ -7,6 +7,7 @@ public enum CBORError : Error {
     case wrongTypeInsideSequence
     case tooLongSequence
     case incorrectUTF8String
+    case maximumDepthExceeded
 }
 
 extension CBOR {
@@ -18,6 +19,7 @@ extension CBOR {
 public class CBORDecoder {
     private var istream : CBORInputStream
     public var options: CBOROptions
+    private var currentDepth = 0
 
     public init(stream: CBORInputStream, options: CBOROptions = CBOROptions()) {
         self.istream = stream
@@ -108,6 +110,11 @@ public class CBORDecoder {
     }
 
     public func decodeItem() throws -> CBOR? {
+        guard currentDepth <= options.maximumDepth
+        else { throw CBORError.maximumDepthExceeded }
+        
+        currentDepth += 1
+        defer { currentDepth -= 1 }
         let b = try istream.popByte()
 
         switch b {
diff --git a/Sources/CBOROptions.swift b/Sources/CBOROptions.swift
index a60c1eb..f422857 100644
--- a/Sources/CBOROptions.swift
+++ b/Sources/CBOROptions.swift
@@ -2,15 +2,19 @@ public struct CBOROptions {
     let useStringKeys: Bool
     let dateStrategy: DateStrategy
     let forbidNonStringMapKeys: Bool
+    /// The maximum number of nested items, inclusive, to decode. A maximum set to 0 dissallows anything other than top-level primitives.
+    let maximumDepth: Int
 
     public init(
         useStringKeys: Bool = false,
         dateStrategy: DateStrategy = .taggedAsEpochTimestamp,
-        forbidNonStringMapKeys: Bool = false
+        forbidNonStringMapKeys: Bool = false,
+        maximumDepth: Int = .max
     ) {
         self.useStringKeys = useStringKeys
         self.dateStrategy = dateStrategy
         self.forbidNonStringMapKeys = forbidNonStringMapKeys
+        self.maximumDepth = maximumDepth
     }
 
     func toCodableEncoderOptions() -> CodableCBOREncoder._Options {
@@ -24,7 +28,8 @@ public struct CBOROptions {
     func toCodableDecoderOptions() -> CodableCBORDecoder._Options {
         return CodableCBORDecoder._Options(
             useStringKeys: self.useStringKeys,
-            dateStrategy: self.dateStrategy
+            dateStrategy: self.dateStrategy,
+            maximumDepth: self.maximumDepth
         )
     }
 }
diff --git a/Sources/Decoder/CodableCBORDecoder.swift b/Sources/Decoder/CodableCBORDecoder.swift
index 62a8ec9..4f77d43 100644
--- a/Sources/Decoder/CodableCBORDecoder.swift
+++ b/Sources/Decoder/CodableCBORDecoder.swift
@@ -7,14 +7,24 @@ final public class CodableCBORDecoder {
     struct _Options {
         let useStringKeys: Bool
         let dateStrategy: DateStrategy
+        let maximumDepth: Int
 
-        init(useStringKeys: Bool = false, dateStrategy: DateStrategy = .taggedAsEpochTimestamp) {
+        init(
+            useStringKeys: Bool = false,
+            dateStrategy: DateStrategy = .taggedAsEpochTimestamp,
+            maximumDepth: Int = .max
+        ) {
             self.useStringKeys = useStringKeys
             self.dateStrategy = dateStrategy
+            self.maximumDepth = maximumDepth
         }
 
         func toCBOROptions() -> CBOROptions {
-            return CBOROptions(useStringKeys: self.useStringKeys, dateStrategy: self.dateStrategy)
+            return CBOROptions(
+                useStringKeys: self.useStringKeys,
+                dateStrategy: self.dateStrategy,
+                maximumDepth: self.maximumDepth
+            )
         }
     }
 
diff --git a/Tests/CBORDecoderTests.swift b/Tests/CBORDecoderTests.swift
index 3ff2ff8..4328283 100644
--- a/Tests/CBORDecoderTests.swift
+++ b/Tests/CBORDecoderTests.swift
@@ -168,4 +168,37 @@ class CBORDecoderTests: XCTestCase {
 
         XCTAssertEqual(decoded, expected)
     }
+    
+    func testDecodeFailsForExtremelyDeepStructures() {
+        let justOverTags: [UInt8] = Array(repeating: 202, count: 1025) + [0]
+        XCTAssertThrowsError(try CBOR.decode(justOverTags, options: CBOROptions(maximumDepth: 1024))) { error in
+            XCTAssertEqual(error as? CBORError, CBORError.maximumDepthExceeded)
+        }
+        let endlessTags: [UInt8] = Array(repeating: 202, count: 10000) + [0]
+        XCTAssertThrowsError(try CBOR.decode(endlessTags, options: CBOROptions(maximumDepth: 1024))) { error in
+            XCTAssertEqual(error as? CBORError, CBORError.maximumDepthExceeded)
+        }
+    }
+    
+    func testDecodeFailsForSillyMaximumDepths() {
+        let singleItem: [UInt8] = [0]
+        XCTAssertThrowsError(try CBOR.decode(singleItem, options: CBOROptions(maximumDepth: -1))) { error in
+            XCTAssertEqual(error as? CBORError, CBORError.maximumDepthExceeded)
+        }
+    }
+    
+    func testDecodeSucceedsForAllowedDeepStructures() {
+        let singleItem: [UInt8] = [0]
+        XCTAssertNoThrow(try CBOR.decode(singleItem, options: CBOROptions(maximumDepth: 0)))
+        let endlessTags: [UInt8] = Array(repeating: 202, count: 1024) + [0]
+        XCTAssertNoThrow(try CBOR.decode(endlessTags, options: CBOROptions(maximumDepth: 1024)))
+    }
+    
+    func testRandomInputDoesNotHitStackLimits() {
+        for _ in 1...50 {
+            let length = Int.random(in: 1...1_000_000)
+            let randomData: [UInt8] = Array(repeating: UInt8.random(in: 0...255), count: length)
+            _ = try? CBOR.decode(randomData, options: CBOROptions(maximumDepth: 1024))
+        }
+    }
 }