[4.1] {JSON,Plist}{En,De}coder defer container pops

stdlib/public/SDK/Foundation/JSONEncoder.swift

@@ -722,6 +722,7 @@ extension _JSONEncoder {
         switch self.options.dateEncodingStrategy {
         case .deferredToDate:
             // Must be called with a surrounding with(pushedKey:) call.
+            // Dates encode as single-value objects; this can't both throw and push a container, so no need to catch the error.
             try date.encode(to: self)
             return self.storage.popContainer()
 
@@ -743,7 +744,16 @@ extension _JSONEncoder {
 
         case .custom(let closure):
             let depth = self.storage.count
-            try closure(date, self)
+            do {
+                try closure(date, self)
+            } catch {
+                // If the value pushed a container before throwing, pop it back off to restore state.
+                if self.storage.count > depth {
+                    let _ = self.storage.popContainer()
+                }
+
+                throw error
+            }
 
             guard self.storage.count > depth else {
                 // The closure didn't encode anything. Return the default keyed container.
@@ -759,15 +769,36 @@ extension _JSONEncoder {
         switch self.options.dataEncodingStrategy {
         case .deferredToData:
             // Must be called with a surrounding with(pushedKey:) call.
-            try data.encode(to: self)
+            let depth = self.storage.count
+            do {
+                try data.encode(to: self)
+            } catch {
+                // If the value pushed a container before throwing, pop it back off to restore state.
+                // This shouldn't be possible for Data (which encodes as an array of bytes), but it can't hurt to catch a failure.
+                if self.storage.count > depth {
+                    let _ = self.storage.popContainer()
+                }
+
+                throw error
+            }
+
             return self.storage.popContainer()
 
         case .base64:
             return NSString(string: data.base64EncodedString())
 
         case .custom(let closure):
             let depth = self.storage.count
-            try closure(data, self)
+            do {
+                try closure(data, self)
+            } catch {
+                // If the value pushed a container before throwing, pop it back off to restore state.
+                if self.storage.count > depth {
+                    let _ = self.storage.popContainer()
+                }
+
+                throw error
+            }
 
             guard self.storage.count > depth else {
                 // The closure didn't encode anything. Return the default keyed container.
@@ -801,7 +832,16 @@ extension _JSONEncoder {
 
         // The value should request a container from the _JSONEncoder.
         let depth = self.storage.count
-        try value.encode(to: self)
+        do {
+            try value.encode(to: self)
+        } catch {
+            // If the value pushed a container before throwing, pop it back off to restore state.
+            if self.storage.count > depth {
+                let _ = self.storage.popContainer()
+            }
+
+            throw error
+        }
 
         // The top container should be a new container.
         guard self.storage.count > depth else {
@@ -2237,9 +2277,8 @@ extension _JSONDecoder {
         switch self.options.dateDecodingStrategy {
         case .deferredToDate:
             self.storage.push(container: value)
-            let date = try Date(from: self)
-            self.storage.popContainer()
-            return date
+            defer { self.storage.popContainer() }
+            return try Date(from: self)
 
         case .secondsSince1970:
             let double = try self.unbox(value, as: Double.self)!
@@ -2271,9 +2310,8 @@ extension _JSONDecoder {
 
         case .custom(let closure):
             self.storage.push(container: value)
-            let date = try closure(self)
-            self.storage.popContainer()
-            return date
+            defer { self.storage.popContainer() }
+            return try closure(self)
         }
     }
 
@@ -2283,9 +2321,8 @@ extension _JSONDecoder {
         switch self.options.dataDecodingStrategy {
         case .deferredToData:
             self.storage.push(container: value)
-            let data = try Data(from: self)
-            self.storage.popContainer()
-            return data
+            defer { self.storage.popContainer() }
+            return try Data(from: self)
 
         case .base64:
             guard let string = value as? String else {
@@ -2300,9 +2337,8 @@ extension _JSONDecoder {
 
         case .custom(let closure):
             self.storage.push(container: value)
-            let data = try closure(self)
-            self.storage.popContainer()
-            return data
+            defer { self.storage.popContainer() }
+            return try closure(self)
         }
     }
 
@@ -2319,13 +2355,10 @@ extension _JSONDecoder {
     }
 
     fileprivate func unbox<T : Decodable>(_ value: Any, as type: T.Type) throws -> T? {
-        let decoded: T
         if type == Date.self || type == NSDate.self {
-            guard let date = try self.unbox(value, as: Date.self) else { return nil }
-            decoded = date as! T
+            return try self.unbox(value, as: Date.self) as? T
         } else if type == Data.self || type == NSData.self {
-            guard let data = try self.unbox(value, as: Data.self) else { return nil }
-            decoded = data as! T
+            return try self.unbox(value, as: Data.self) as? T
         } else if type == URL.self || type == NSURL.self {
             guard let urlString = try self.unbox(value, as: String.self) else {
                 return nil
@@ -2336,17 +2369,14 @@ extension _JSONDecoder {
                                                                         debugDescription: "Invalid URL string."))
             }
 
-            decoded = (url as! T)
+            return url as! T
         } else if type == Decimal.self || type == NSDecimalNumber.self {
-            guard let decimal = try self.unbox(value, as: Decimal.self) else { return nil }
-            decoded = decimal as! T
+            return try self.unbox(value, as: Decimal.self) as? T
         } else {
             self.storage.push(container: value)
-            decoded = try type.init(from: self)
-            self.storage.popContainer()
+            defer { self.storage.popContainer() }
+            return try type.init(from: self)
         }
-
-        return decoded
     }
 }
 

stdlib/public/SDK/Foundation/PlistEncoder.swift

@@ -490,7 +490,16 @@ extension _PlistEncoder {
 
         // The value should request a container from the _PlistEncoder.
         let depth = self.storage.count
-        try value.encode(to: self)
+        do {
+            try value.encode(to: self)
+        } catch let error {
+            // If the value pushed a container before throwing, pop it back off to restore state.
+            if self.storage.count > depth {
+                let _ = self.storage.popContainer()
+            }
+
+            throw error
+        }
 
         // The top container should be a new container.
         guard self.storage.count > depth else {
@@ -1755,20 +1764,15 @@ extension _PlistDecoder {
     }
 
     fileprivate func unbox<T : Decodable>(_ value: Any, as type: T.Type) throws -> T? {
-        let decoded: T
         if type == Date.self || type == NSDate.self {
-            guard let date = try self.unbox(value, as: Date.self) else { return nil }
-            decoded = date as! T
+            return try self.unbox(value, as: Date.self) as? T
         } else if type == Data.self || type == NSData.self {
-            guard let data = try self.unbox(value, as: Data.self) else { return nil }
-            decoded = data as! T
+            return try self.unbox(value, as: Data.self) as? T
         } else {
             self.storage.push(container: value)
-            decoded = try type.init(from: self)
-            self.storage.popContainer()
+            defer { self.storage.popContainer() }
+            return try type.init(from: self)
         }
-
-        return decoded
     }
 }
 

test/stdlib/TestJSONEncoder.swift

@@ -799,6 +799,119 @@ class TestJSONEncoder : TestJSONEncoderSuper {
       expectEqual(type(of: decoded), Employee.self, "Expected decoded value to be of type Employee; got \(type(of: decoded)) instead.")
   }
 
+  // MARK: - Encoder State
+  // SR-6078
+  func testEncoderStateThrowOnEncode() {
+    struct ReferencingEncoderWrapper<T : Encodable> : Encodable {
+      let value: T
+      init(_ value: T) { self.value = value }
+
+      func encode(to encoder: Encoder) throws {
+        // This approximates a subclass calling into its superclass, where the superclass encodes a value that might throw.
+        // The key here is that getting the superEncoder creates a referencing encoder.
+        var container = encoder.unkeyedContainer()
+        let superEncoder = container.superEncoder()
+
+        // Pushing a nested container on leaves the referencing encoder with multiple containers.
+        var nestedContainer = superEncoder.unkeyedContainer()
+        try nestedContainer.encode(value)
+      }
+    }
+
+    // The structure that would be encoded here looks like
+    //
+    //   [[[Float.infinity]]]
+    //
+    // The wrapper asks for an unkeyed container ([^]), gets a super encoder, and creates a nested container into that ([[^]]).
+    // We then encode an array into that ([[[^]]]), which happens to be a value that causes us to throw an error.
+    //
+    // The issue at hand reproduces when you have a referencing encoder (superEncoder() creates one) that has a container on the stack (unkeyedContainer() adds one) that encodes a value going through box_() (Array does that) that encodes something which throws (Float.infinity does that).
+    // When reproducing, this will cause a test failure via fatalError().
+    _ = try? JSONEncoder().encode(ReferencingEncoderWrapper([Float.infinity]))
+  }
+
+  func testEncoderStateThrowOnEncodeCustomDate() {
+    // This test is identical to testEncoderStateThrowOnEncode, except throwing via a custom Date closure.
+    struct ReferencingEncoderWrapper<T : Encodable> : Encodable {
+      let value: T
+      init(_ value: T) { self.value = value }
+      func encode(to encoder: Encoder) throws {
+        var container = encoder.unkeyedContainer()
+        let superEncoder = container.superEncoder()
+        var nestedContainer = superEncoder.unkeyedContainer()
+        try nestedContainer.encode(value)
+      }
+    }
+
+    // The closure needs to push a container before throwing an error to trigger.
+    let encoder = JSONEncoder()
+    encoder.dateEncodingStrategy = .custom({ _, encoder in
+      let _ = encoder.unkeyedContainer()
+      enum CustomError : Error { case foo }
+      throw CustomError.foo
+    })
+
+    _ = try? encoder.encode(ReferencingEncoderWrapper(Date()))
+  }
+
+  func testEncoderStateThrowOnEncodeCustomData() {
+    // This test is identical to testEncoderStateThrowOnEncode, except throwing via a custom Data closure.
+    struct ReferencingEncoderWrapper<T : Encodable> : Encodable {
+      let value: T
+      init(_ value: T) { self.value = value }
+      func encode(to encoder: Encoder) throws {
+        var container = encoder.unkeyedContainer()
+        let superEncoder = container.superEncoder()
+        var nestedContainer = superEncoder.unkeyedContainer()
+        try nestedContainer.encode(value)
+      }
+    }
+
+    // The closure needs to push a container before throwing an error to trigger.
+    let encoder = JSONEncoder()
+    encoder.dataEncodingStrategy = .custom({ _, encoder in
+      let _ = encoder.unkeyedContainer()
+      enum CustomError : Error { case foo }
+      throw CustomError.foo
+    })
+
+    _ = try? encoder.encode(ReferencingEncoderWrapper(Data()))
+  }
+
+  // MARK: - Decoder State
+  // SR-6048
+  func testDecoderStateThrowOnDecode() {
+    // The container stack here starts as [[1,2,3]]. Attempting to decode as [String] matches the outer layer (Array), and begins decoding the array.
+    // Once Array decoding begins, 1 is pushed onto the container stack ([[1,2,3], 1]), and 1 is attempted to be decoded as String. This throws a .typeMismatch, but the container is not popped off the stack.
+    // When attempting to decode [Int], the container stack is still ([[1,2,3], 1]), and 1 fails to decode as [Int].
+    let json = "[1,2,3]".data(using: .utf8)!
+    let _ = try! JSONDecoder().decode(EitherDecodable<[String], [Int]>.self, from: json)
+  }
+
+  func testDecoderStateThrowOnDecodeCustomDate() {
+    // This test is identical to testDecoderStateThrowOnDecode, except we're going to fail because our closure throws an error, not because we hit a type mismatch.
+    let decoder = JSONDecoder()
+    decoder.dateDecodingStrategy = .custom({ decoder in
+      enum CustomError : Error { case foo }
+      throw CustomError.foo
+    })
+
+    let json = "{\"value\": 1}".data(using: .utf8)!
+    let _ = try! decoder.decode(EitherDecodable<TopLevelWrapper<Date>, TopLevelWrapper<Int>>.self, from: json)
+  }
+
+  func testDecoderStateThrowOnDecodeCustomData() {
+    // This test is identical to testDecoderStateThrowOnDecode, except we're going to fail because our closure throws an error, not because we hit a type mismatch.
+    let decoder = JSONDecoder()
+    decoder.dataDecodingStrategy = .custom({ decoder in
+      enum CustomError : Error { case foo }
+      throw CustomError.foo
+    })
+
+    let json = "{\"value\": 1}".data(using: .utf8)!
+    let _ = try! decoder.decode(EitherDecodable<TopLevelWrapper<Data>, TopLevelWrapper<Int>>.self, from: json)
+  }
+
   // MARK: - Helper Functions
   private var _jsonEmptyDictionary: Data {
     return "{}".data(using: .utf8)!
@@ -1391,6 +1504,20 @@ fileprivate struct DoubleNaNPlaceholder : Codable, Equatable {
   }
 }
 
+fileprivate enum EitherDecodable<T : Decodable, U : Decodable> : Decodable {
+  case t(T)
+  case u(U)
+
+  init(from decoder: Decoder) throws {
+    let container = try decoder.singleValueContainer()
+    do {
+      self = .t(try container.decode(T.self))
+    } catch {
+      self = .u(try container.decode(U.self))
+    }
+  }
+}
+
 // MARK: - Run Tests
 
 #if !FOUNDATION_XCTEST
@@ -1439,5 +1566,11 @@ JSONEncoderTests.test("testInterceptDecimal") { TestJSONEncoder().testInterceptD
 JSONEncoderTests.test("testInterceptURL") { TestJSONEncoder().testInterceptURL() }
 JSONEncoderTests.test("testTypeCoercion") { TestJSONEncoder().testTypeCoercion() }
 JSONEncoderTests.test("testDecodingConcreteTypeParameter") { TestJSONEncoder().testDecodingConcreteTypeParameter() }
+JSONEncoderTests.test("testEncoderStateThrowOnEncode") { TestJSONEncoder().testEncoderStateThrowOnEncode() }
+JSONEncoderTests.test("testEncoderStateThrowOnEncodeCustomDate") { TestJSONEncoder().testEncoderStateThrowOnEncodeCustomDate() }
+JSONEncoderTests.test("testEncoderStateThrowOnEncodeCustomData") { TestJSONEncoder().testEncoderStateThrowOnEncodeCustomData() }
+JSONEncoderTests.test("testDecoderStateThrowOnDecode") { TestJSONEncoder().testDecoderStateThrowOnDecode() }
+JSONEncoderTests.test("testDecoderStateThrowOnDecodeCustomDate") { TestJSONEncoder().testDecoderStateThrowOnDecodeCustomDate() }
+JSONEncoderTests.test("testDecoderStateThrowOnDecodeCustomData") { TestJSONEncoder().testDecoderStateThrowOnDecodeCustomData() }
 runAllTests()
 #endif